Impulse generator

Generator for impulses of different shapes.

::: {#cell-3 .cell 0=‘e’ 1=‘x’ 2=‘p’ 3=‘o’ 4=‘r’ 5=‘t’}

import numpy as np
from typing import Tuple, Optional
from einops import einsum

:::

import matplotlib.pyplot as plt

::: {#cell-5 .cell 0=‘e’ 1=‘x’ 2=‘p’ 3=‘o’ 4=‘r’ 5=‘t’}

class Generator:
    pass

:::

::: {#cell-6 .cell 0=‘e’ 1=‘x’ 2=‘p’ 3=‘o’ 4=‘r’ 5=‘t’}

class Gaussian(Generator):
    def __init__(
        self,
        num_points: int = 100,
    ):
        self.x = np.linspace(0, 1, num_points)
        self.dx = self.x[1] - self.x[0]

    def __call__(
        self,
        mean: float = 0.5,  # in percentage w.r.t. the number of points
        std: float = 0.05,  # in percentage w.r.t. the number of points
    ):
        std_corr = np.max([std, self.dx])  # To avoid too narrow gaussians
        return np.exp(-((self.x - mean) ** 2 / (2 * std_corr**2)))

:::

::: {#cell-7 .cell 0=‘t’ 1=‘e’ 2=‘s’ 3=‘t’}

y = Gaussian(num_points=501)(mean=0.5, std=0.0000005)
assert len(y) == 501
# plot the gaussian
plt.plot(y)

:::

::: {#cell-8 .cell 0=‘e’ 1=‘x’ 2=‘p’ 3=‘o’ 4=‘r’ 5=‘t’}

class Gaussian2d(Generator):
    """This class generates a 2D gaussian distribution."""

    def __init__(
        self,
        num_points_x: int = 100,
        aspect_ratio: float = 1.0,  # aspect ratio of the lengths of the two axes, ly/lx
    ):
        self.aspect_ratio = aspect_ratio
        num_points_y = int(np.floor((num_points_x - 1) * aspect_ratio)) + 1
        x = np.linspace(0, 1, num_points_x)
        y = np.linspace(0, aspect_ratio, num_points_y)
        self.dx = x[1] - x[0]
        self.dy = y[1] - y[0]
        self.grid_x, self.grid_y = np.meshgrid(x, y, indexing="ij")

    def __call__(
        self,
        mean: tuple[float, float] = (0.5, 0.5),  # respect to the plate aspect ratio
        std: float = 0.05,  # in percentage w.r.t. the number of points
    ):
        std_corr = np.max([std, self.dx, self.dy])  # To avoid too narrow gaussians
        return np.exp(
            -(
                (self.grid_x - mean[0]) ** 2
                + (self.grid_y - self.aspect_ratio * mean[1]) ** 2
            )
            / (2 * std**2)
        )

:::

::: {#cell-9 .cell 0=‘t’ 1=‘e’ 2=‘s’ 3=‘t’}

nx = 101
aspect_ratio = 0.7
z = Gaussian2d(num_points_x=nx, aspect_ratio=aspect_ratio)(mean=(0.9, 1.0), std=0.05)
ny = int(np.floor((nx - 1) * aspect_ratio)) + 1
print(z.shape)
print(nx, ny)
assert z.shape == (nx, ny)
# plot the gaussian
plt.imshow(z, vmin=0, vmax=1)

:::

::: {#cell-10 .cell 0=‘e’ 1=‘x’ 2=‘p’ 3=‘o’ 4=‘r’ 5=‘t’}

class NoiseBurst(Generator):
    def __init__(
        self,
        num_points: int = 100,
    ):
        self.num_points = num_points
        self.x = np.linspace(0, 1, num_points)
        self.gaussian = Gaussian(num_points)
        self.dx = self.x[1] - self.x[0]

    def __call__(
        self,
        rng: np.random.Generator,
        noise_range=[0, 1],
        burst_mean: float = 0.5,
        burst_std: float = 0.1,
    ):
        std_corr = np.max([burst_std, self.dx])  # To avoid too narrow gaussian envelope
        y = rng.uniform(*noise_range, self.num_points) * self.gaussian(
            burst_mean, burst_std
        )
        return y

:::

::: {#cell-11 .cell 0=‘t’ 1=‘e’ 2=‘s’ 3=‘t’}

rng = np.random.default_rng(42)
y = NoiseBurst(num_points=500)(rng, noise_range=[0, 1], burst_mean=0.5, burst_std=0.05)
assert len(y) == 500

:::

::: {#cell-12 .cell 0=‘e’ 1=‘x’ 2=‘p’ 3=‘o’ 4=‘r’ 5=‘t’}

class Noise(Generator):
    def __init__(
        self,
        num_points: int = 100,
    ):
        self.num_points = num_points
        self.x = np.linspace(0, 1, num_points)
        self.dx = self.x[1] - self.x[0]

    def __call__(
        self,
        rng: np.random.Generator,
        noise_range=[0, 1],
    ):
        y = rng.uniform(*noise_range, self.num_points)
        return y

:::

::: {#cell-13 .cell 0=‘e’ 1=‘x’ 2=‘p’ 3=‘o’ 4=‘r’ 5=‘t’}

class Noise2d(Generator):
    def __init__(
        self,
        num_points_x: int = 100,
        aspect_ratio: float = 1.0,  # aspect ratio of the lengths of the two axes, ly/lx
    ):
        self.num_points_x = num_points_x
        num_points_y = int(np.floor((num_points_x - 1) * aspect_ratio)) + 1
        self.num_points_y = num_points_y
        self.dx = 1 / (num_points_x - 1)

    def __call__(
        self,
        rng: np.random.Generator,
        noise_range=[0, 1],
    ):
        z = rng.uniform(*noise_range, (self.num_points_x, self.num_points_y))
        return z

:::

::: {#cell-14 .cell 0=‘t’ 1=‘e’ 2=‘s’ 3=‘t’}

rng = np.random.default_rng(42)
z = Noise2d(num_points_x=100, aspect_ratio=0.7)(rng, noise_range=[0, 1])
assert z.shape == (100, 70)
plt.imshow(z, vmin=0, vmax=1)

:::

::: {#cell-15 .cell 0=‘e’ 1=‘x’ 2=‘p’ 3=‘o’ 4=‘r’ 5=‘t’}

class Triangle(Generator):
    def __init__(
        self,
        num_points: int = 100,
        num_modes: int = 25,
        length: float = 1.0,
    ):
        self.length = length
        self.num_points = num_points
        self.x = np.linspace(0, 1, num_points)
        self.dx = self.x[1] - self.x[0]

        # Wavenumbers
        n = np.arange(1, num_modes + 1)  # Mode numbers
        self.k = n * np.pi / length

    def __call__(
        self,
        height: float = 1.0,  # height of the excitation
        position: float = 0.5,  # normalized position of the excitation [0, 1]
    ):
        # Modal coefficients (based on the pluck excitation)
        A_n = (
            height
            * (
                self.length
                / (self.length - position)
                * np.sin(self.k * position)
                / (self.k * position)
            )
            / self.k
        )

        # Sum of the modes
        return np.sum(
            A_n[:, None] * np.sin(self.k[:, None] * self.x),
            axis=0,
        )

:::

::: {#cell-16 .cell 0=‘t’ 1=‘e’ 2=‘s’ 3=‘t’}

y = Triangle(num_points=101, num_modes=25)(height=1.0, position=0.1)
# plot the triangle
plt.plot(y)

:::

::: {#cell-17 .cell 0=‘e’ 1=‘x’ 2=‘p’ 3=‘o’ 4=‘r’ 5=‘t’}

class SineMode(Generator):
    def __init__(
        self,
        num_points: int = 100,
    ):
        self.x = np.linspace(0, 1, num_points)
        self.dx = self.x[1] - self.x[0]

    def __call__(
        self,
        k: int = 1,
    ):
        assert k > 0, "k must be positive"
        return np.sin(np.pi * k * self.x)

:::

::: {#cell-18 .cell 0=‘e’ 1=‘x’ 2=‘p’ 3=‘o’ 4=‘r’ 5=‘t’}

def make_pluck_hammer(
    y: np.ndarray,
    ic_type: str = "pluck",  # "pluck" or "hammer"
) -> Tuple[np.ndarray, np.ndarray]:
    if ic_type == "pluck":
        return y, np.zeros_like(y)
    elif ic_type == "hammer":
        return np.zeros_like(y), y
    else:
        raise ValueError(f"ic_type should be either 'pluck' or 'hammer', got {ic_type}")


def generate_initial_condition(
    rng: np.random.Generator = np.random.default_rng(42),
    generator: Generator = Gaussian(),
    ic_type: str = "pluck",  # "pluck" or "hammer"
    ic_max_amplitude: float = 1.0,  # Amplitude of the initial condition, when ic_amplitude_random is True, this is the upper bound
    ic_min_amplitude: float = 0.0,  # only used when ic_amplitude_random is True
    ic_amplitude_random: bool = False,  # If True, the amplitude is chosen randomly between ic_min_amplitude and ic_max_amplitude
    ic_sine_k: int = 1,  # only used when ic_type is "sine"
) -> Tuple[np.ndarray, np.ndarray]:  # a tuple of position and velocity
    # Check that ic_max_amplitude is larger than ic_min_amplitude
    if ic_amplitude_random:
        assert ic_max_amplitude > ic_min_amplitude, (
            f"ic_max_amplitude should be larger than ic_min_amplitude, got "
            f"ic_max_amplitude={ic_max_amplitude} and ic_min_amplitude={ic_min_amplitude}"
        )
    # TODO: Mean and std are hard-coded, could be made more flexible
    mean = rng.uniform(0.3, 0.7)
    min_std = 2 * generator.dx
    std = rng.uniform(min_std, 0.1)
    if isinstance(generator, Gaussian):
        y = generator(mean, std)
    elif isinstance(generator, NoiseBurst):
        y = generator(rng, noise_range=[-1, 1], burst_mean=mean, burst_std=std)
    elif isinstance(generator, Noise):
        y = generator(rng, noise_range=[-1, 1])
    elif isinstance(generator, SineMode):
        y = generator(k=ic_sine_k)
    elif isinstance(generator, Triangle):
        y = generator(height=ic_max_amplitude, position=rng.uniform(0.1, 0.9))
    elif isinstance(generator, Gaussian2d):
        mean = (rng.uniform(0.3, 0.7), rng.uniform(0.3, 0.7))
        y = generator(mean, std)
    elif isinstance(generator, Noise2d):
        y = generator(rng, noise_range=[-1, 1])
    else:
        raise TypeError(
            f"generator should be either Gaussian, Noise, Sine or NoiseBurst, got {type(generator)}"
        )
    # Normalize the amplitude to the desired value
    if ic_amplitude_random:
        amplitude = rng.uniform(ic_min_amplitude, ic_max_amplitude)
    else:
        amplitude = ic_max_amplitude

    y = y * amplitude / np.max(np.abs(y))
    return make_pluck_hammer(y, ic_type)

:::

::: {#cell-19 .cell 0=‘t’ 1=‘e’ 2=‘s’ 3=‘t’}

rng = np.random.default_rng()
ic_type = "pluck"
ic_max_amplitude = 0.7
ic_min_amplitude = 0.1
ic_amplitude_random = False
num_points = 500
u, v = generate_initial_condition(
    rng,
    SineMode(num_points=num_points),
    ic_type=ic_type,
    ic_max_amplitude=ic_max_amplitude,
    ic_min_amplitude=ic_min_amplitude,
    ic_amplitude_random=ic_amplitude_random,
    ic_sine_k=10,
)
if ic_type == "pluck":
    assert np.all(v == 0)
    if ic_amplitude_random:
        assert np.max(np.abs(u)) <= ic_max_amplitude
        assert np.max(np.abs(u)) >= ic_min_amplitude
    else:
        assert np.max(np.abs(u)) == ic_max_amplitude
elif ic_type == "hammer":
    assert np.all(u == 0)
    if ic_amplitude_random:
        assert np.max(np.abs(v)) <= ic_max_amplitude
        assert np.max(np.abs(v)) >= ic_min_amplitude
    else:
        assert np.max(np.abs(v)) == ic_max_amplitude

# Plot the initial conditions
fig, ax = plt.subplots()
ax.plot(u)
ax.plot(v)

:::

::: {#cell-20 .cell 0=‘t’ 1=‘e’ 2=‘s’ 3=‘t’}

# Test the 2d functions

rng = np.random.default_rng()
ic_type = "hammer"
ic_max_amplitude = 0.5
ic_min_amplitude = 0.1
ic_amplitude_random = False
num_points = 500
u, v = generate_initial_condition(
    rng,
    Noise2d(num_points_x=num_points, aspect_ratio=0.7),
    ic_type=ic_type,
    ic_max_amplitude=ic_max_amplitude,
    ic_min_amplitude=ic_min_amplitude,
    ic_amplitude_random=ic_amplitude_random,
)
print(np.max(np.abs(v)))
if ic_type == "pluck":
    assert np.all(v == 0)
    if ic_amplitude_random:
        assert np.max(np.abs(u)) <= ic_max_amplitude
        assert np.max(np.abs(u)) >= ic_min_amplitude
    else:
        assert np.allclose(np.max(np.abs(u)), ic_max_amplitude)
elif ic_type == "hammer":
    assert np.all(u == 0)
    if ic_amplitude_random:
        assert np.max(np.abs(v)) <= ic_max_amplitude
        assert np.max(np.abs(v)) >= ic_min_amplitude
    else:
        assert np.allclose(np.max(np.abs(v)), ic_max_amplitude)

# Plot the initial conditions
fig, ax = plt.subplots(1, 2)
ax[0].imshow(u)
ax[1].imshow(v)

:::