Repository: 0x23/WaveSimulator2D
Branch: main
Commit: 990124d7c641
Files: 17
Total size: 67.7 KB

Directory structure:
gitextract_7z9h8v6g/

├── README.md
├── requirements.txt
└── wave_sim2d/
    ├── __init__.py
    ├── develop_tests.py
    ├── examples/
    │   ├── example0.py
    │   ├── example1.py
    │   ├── example2.py
    │   ├── example3.py
    │   └── example4.py
    ├── main.py
    ├── scene_objects/
    │   ├── source.py
    │   ├── static_dampening.py
    │   ├── static_image_scene.py
    │   ├── static_refractive_index.py
    │   └── strain_refractive_index.py
    ├── wave_simulation.py
    └── wave_visualizer.py

================================================
FILE CONTENTS
================================================

================================================
FILE: README.md
================================================
# 2D Wave Simulation on the GPU

This repository contains a lightweight 2D wave simulator running on the GPU using CuPy library (probably requires a NVIDIA GPU). It can be used for 2D light and sound simulations.
A simple visualizer shows the field and its intensity on the screen and writes a movie file for each to disks. The goal is to provide a fast, easy to use but still felxible wave simulator.

<div style="display: flex;">
    <img src="images/simulation_1.jpg" alt="Example Image 1" width="49%">
    <img src="images/simulation_2.jpg" alt="Example Image 2" width="49%">
</div>

### Update 06.04.2025

* Scene objects can now draw to a visualization layer (most of them do not yet, feel free to contribute) !
* Example 4 now shows a two-mirror optical cavity and how standing waves emerge.
* Added new Line Sources
* Added Refractive index Polygon object (StaticRefractiveIndexPolygon)
* Added Refractive index Box object (StaticRefractiveIndexBox)
* Fixed some issues with the examples

<div style="display: flex;">
    <img src="images/optical_cavity.jpg" alt="Example 4 - Optical Cavity with Standing Waves" width="50%">
</div>

### Update 01.04.2024

* Refactored the code to support a more flexible scene description. A simulation scene now consists of a list of objects that add their contribution to the fields.
They can be combined to build complex and time dependent simulations. The refactoring also made the core simulation code even simpler.
* Added a few new custom colormaps that work well for wave simulations.
* Added new examples, which should make it easier to understand the usage of the program and how you can setup your own simulations: [examples](source/examples).

<div style="display: flex;">
    <img src="images/simulation_3.jpg" alt="Example Image 3" width="45%">
    <img src="images/simulation_4.jpg" alt="Example Image 4" width="45%">
</div>

The old image based scene description is still available as a scene object. You can continue to use the convenience of an image editing software and create simulations
without much programming.

###  Image Scene Decsription Usage ###

When using the 'StaticImageScene' class the simulation scenes can given as an 8Bit RGB image with the following channel semantics:
* Red:   The Refractive index times 100 (for refractive index 1.5 you would use value 150)
* Green: Each pixel with a green value above 0 is a sinusoidal wave source. The green value defines its frequency.
* Blue:  Absorbtion field. Larger values correspond to higher dampening of the waves, use graduated transitions to avoid reflections

WARNING: Do not use anti-aliasing for the green channel ! The shades produced are interpreted as different source frequencies, which yields weird results.

<div style="display: flex;">
    <img src="images/source_antialiasing.png" alt="Example Image 5" width="50%">
</div>

### Recommended Installation ###

1. Install Python and PyCharm IDE
2. Clone the Project to you hard disk
3. Open the folder as a Project using PyCharm
4. If prompted to install requirements, accept (or install requirements using pip -r requirements.txt)
5. Right click on one of the examples in wave_sim2d/examples and select run

NOTE: If you have issues installing the `cupy` library
1. Make sure you have the `nvidia-cuda-toolkit` installed. 
You can check it by running `nvcc --version`.
1. In the *requirements.txt* file, replace `cupy` by `cupy-cuda[version-number]x`. 
   Where the version number displayed when running `nvcc --version` (example: `cupy-cuda11x`).





================================================
FILE: requirements.txt
================================================
numpy
opencv-python
matplotlib
cupy

================================================
FILE: wave_sim2d/__init__.py
================================================


================================================
FILE: wave_sim2d/develop_tests.py
================================================
import wave_visualizer
import wave_visualizer as vis
import wave_simulation as sim
import numpy as np
import cv2
import math
import json
from scene_objects.static_dampening import StaticDampening
from scene_objects.static_refractive_index import StaticRefractiveIndex
from scene_objects.static_image_scene import StaticImageScene
from scene_objects.source import PointSource, ModulatorSmoothSquare, ModulatorDiscreteSignal


def build_example_scene1(scene_image):
    """
    This example uses the old image scene description. See 'StaticImageScene' for more information.
    """
    scene_objects = [StaticImageScene(scene_image)]
    return scene_objects


def build_example_scene2(width, height):
    """
    In this example, a new scene is created from scratch and a few emitters are places manually.
    One of the emitters uses an amplitude modulation object to change brightness over time
    """
    objects = []

    # Add a static dampening field without any dampending in the interior (value 1.0 means no dampening)
    # However a dampening layer at the border is added to avoid reflections (see parameter 'border thickness')
    objects.append(StaticDampening(np.ones((height, width)), 48))

    # add a constant refractive index field
    objects.append(StaticRefractiveIndex(np.full((height, width), 1.5)))

    # add a simple point source
    objects.append(PointSource(200, 250, 0.19, 5))

    # add a point source with an amplitude modulator
    amplitude_modulator = ModulatorDiscreteSignal(np.random.randint(2, size=64), 0.0006)
    objects.append(PointSource(200, 350, 0.19, 5, amp_modulator=amplitude_modulator))

    return objects


def simulate(scene_image_fn, num_iterations,
             simulation_steps_per_frame, write_videos,
             field_colormap, intensity_colormap,
             background_image_fn=None):
    # reset random number generator
    np.random.seed(0)

    # load scene image
    scene_image = cv2.cvtColor(cv2.imread(scene_image_fn), cv2.COLOR_BGR2RGB)

    background_image = None
    if background_image_fn is not None:
        background_image = cv2.imread(background_image_fn)
        background_image = cv2.resize(background_image, (scene_image.shape[1], scene_image.shape[0]))

    # create simulator and visualizer objects
    simulator = sim.WaveSimulator2D(scene_image.shape[1], scene_image.shape[0])
    visualizer = vis.WaveVisualizer(field_colormap=field_colormap, intensity_colormap=intensity_colormap)

    # build simulation scene
    simulator.scene_objects = build_example_scene2(scene_image.shape[1], scene_image.shape[0])

    # create video writers
    if write_videos:
        video_writer1 = cv2.VideoWriter('simulation_field.avi', cv2.VideoWriter_fourcc(*'FFV1'),
                                       60, (scene_image.shape[1], scene_image.shape[0]))
        video_writer2 = cv2.VideoWriter('simulation_intensity.avi', cv2.VideoWriter_fourcc(*'FFV1'),
                                       60, (scene_image.shape[1], scene_image.shape[0]))

    # run simulation
    for i in range(num_iterations):
        simulator.update_scene()
        simulator.update_field()
        visualizer.update(simulator)

        if i % simulation_steps_per_frame == 0:
            frame_int = visualizer.render_intensity(1.0)
            frame_field = visualizer.render_field(1.0)

            if background_image is not None:
                frame_int = cv2.add(background_image, frame_int)
                frame_field = cv2.add(background_image, frame_field)

           # frame_int = cv2.pyrDown(frame_int)
           # frame_field = cv2.pyrDown(frame_field)
            cv2.imshow("Wave Simulation", frame_field) #cv2.resize(frame_int, dsize=(1024, 1024)))
            cv2.waitKey(1)

            if write_videos:
                video_writer1.write(frame_field)
                video_writer2.write(frame_int)

        if i % 128 == 0:
            print(f'{int((i+1)/num_iterations*100)}%')


if __name__ == "__main__":
    print('This file contains tests for development and you may not bve able to run it without errors')
    print('Please take a look at the previded examples')

    # increase simulation_steps_per_frame to better utilize GPU
    # good colormaps for field: RdBu[invert=True], colormap_wave1, colormap_wave2, colormap_wave4, icefire
    simulate('../exxample_data/scene_lens_doubleslit.png',
             20000,
             simulation_steps_per_frame=16,
             write_videos=True,
             field_colormap=vis.get_colormap_lut('colormap_wave4', invert=False, black_level=-0.05),
#             field_colormap=vis.get_colormap_lut('RdBu', invert=True, make_symmetric=True),
             intensity_colormap=vis.get_colormap_lut('afmhot', invert=False, black_level=0.0),
             background_image_fn=None)



================================================
FILE: wave_sim2d/examples/example0.py
================================================
import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '../'))  # noqa

import cv2
import wave_sim2d.wave_visualizer as vis
import wave_sim2d.wave_simulation as sim
from wave_sim2d.scene_objects.source import *
from wave_sim2d.scene_objects.static_refractive_index import *

def build_scene():
    """
    This example creates the simplest possible simulation using a single emitter.
    """
    width = 512
    height = 512
    objects = [PointSource(200, 256, 0.1, 5)]
    # objects.append(StaticRefractiveIndexPolygon([[400, 255], [300, 200], [300, 300]], 1.5))
    # objects = [LineSource((200, 265), (250, 105), 0.2, 0.5)]

    return objects, width, height


def main():
    # create colormaps
    field_colormap = vis.get_colormap_lut('colormap_wave1', invert=False, black_level=-0.05)
    intensity_colormap = vis.get_colormap_lut('afmhot', invert=False, black_level=0.0)

    # build simulation scene
    scene_objects, w, h = build_scene()

    # create simulator and visualizer objects
    simulator = sim.WaveSimulator2D(w, h, scene_objects)
    visualizer = vis.WaveVisualizer(field_colormap=field_colormap, intensity_colormap=intensity_colormap)

    # run simulation
    for i in range(1000):
        simulator.update_scene()
        simulator.update_field()
        visualizer.update(simulator)

        # show field
        frame_field = visualizer.render_field(1.0)
        cv2.imshow("Wave Simulation Field", frame_field)

        # show intensity
        # frame_int = visualizer.render_intensity(1.0)
        # cv2.imshow("Wave Simulation Intensity", frame_int)

        cv2.waitKey(1)


if __name__ == "__main__":
    main()



================================================
FILE: wave_sim2d/examples/example1.py
================================================
import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '../'))  # noqa

import numpy as np
import cv2

import wave_sim2d.wave_visualizer as vis
import wave_sim2d.wave_simulation as sim
from wave_sim2d.scene_objects.static_image_scene import StaticImageScene


def build_scene(scene_image_path):
    """
    This example uses the 'old' image scene description. See 'StaticImageScene' for more information.
    """
    # load scene image
    scene_image = cv2.cvtColor(cv2.imread(scene_image_path), cv2.COLOR_BGR2RGB)

    # create the scene object list with an 'StaticImageScene' entry as the only scene object
    # more scene objects can be added to the list to build more complex scenes
    scene_objects = [StaticImageScene(scene_image, source_fequency_scale=2.0)]

    return scene_objects, scene_image.shape[1], scene_image.shape[0]


def main():
    # Set scene image path. The image encodes refractive index, dampening and emitters in its color channels
    # see 'static_image_scene.StaticImageScene' class for a more detailed description.
    # please take a look at the image to understand what is happening in the simulation
    scene_image_path = '../../example_data/scene_lens_doubleslit.png'

    # create colormaps
    field_colormap = vis.get_colormap_lut('colormap_wave1', invert=False, black_level=-0.05)
    intensity_colormap = vis.get_colormap_lut('afmhot', invert=False, black_level=0.0)

    # reset random number generator
    np.random.seed(0)

    # build simulation scene
    scene_objects, w, h = build_scene(scene_image_path)

    # create simulator and visualizer objects
    simulator = sim.WaveSimulator2D(w, h, scene_objects)
    visualizer = vis.WaveVisualizer(field_colormap=field_colormap, intensity_colormap=intensity_colormap)

    # run simulation
    for i in range(2000):
        simulator.update_scene()
        simulator.update_field()
        visualizer.update(simulator)

        # visualize very N frames
        if (i % 4) == 0:
            # show field
            frame_field = visualizer.render_field(1.0)
            cv2.imshow("Wave Simulation Field", frame_field)

            # show intensity
            # frame_int = visualizer.render_intensity(1.0)
            # cv2.imshow("Wave Simulation Intensity", frame_int)

        cv2.waitKey(1)


if __name__ == "__main__":
    main()



================================================
FILE: wave_sim2d/examples/example2.py
================================================
import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '../'))  # noqa

import numpy as np
import cv2
import wave_sim2d.wave_visualizer as vis
import wave_sim2d.wave_simulation as sim
from wave_sim2d.scene_objects.static_dampening import StaticDampening
from wave_sim2d.scene_objects.static_refractive_index import StaticRefractiveIndex
from wave_sim2d.scene_objects.source import PointSource, ModulatorSmoothSquare


def build_scene():
    """
    In this example, a new scene is created from scratch and a few emitters are places manually.
    One of the emitters uses an amplitude modulation object to change brightness over time
    """
    width = 600
    height = 600
    objects = []

    # Add a static dampening field without any dampending in the interior (value 1.0 means no dampening)
    # However a dampening layer at the border is added to avoid reflections (see parameter 'border thickness')
    objects.append(StaticDampening(np.ones((height, width)), 32))

    # add a constant refractive index field
    objects.append(StaticRefractiveIndex(np.full((height, width), 1.5)))

    # add a simple point source
    objects.append(PointSource(200, 220, 0.2, 8))

    # add a point source with an amplitude modulator
    amplitude_modulator = ModulatorSmoothSquare(0.025, 0.0, smoothness=0.5)
    objects.append(PointSource(200, 380, 0.2, 8, amp_modulator=amplitude_modulator))

    return objects, width, height


def main():
    # create colormaps
    field_colormap = vis.get_colormap_lut('colormap_wave4', invert=False, black_level=-0.05)
    intensity_colormap = vis.get_colormap_lut('afmhot', invert=False, black_level=0.0)

    # reset random number generator
    np.random.seed(0)

    # build simulation scene
    scene_objects, w, h = build_scene()

    # create simulator and visualizer objects
    simulator = sim.WaveSimulator2D(w, h, scene_objects)
    visualizer = vis.WaveVisualizer(field_colormap=field_colormap, intensity_colormap=intensity_colormap)

    # run simulation
    for i in range(2000):
        simulator.update_scene()
        simulator.update_field()
        visualizer.update(simulator)

        # visualize very N frames
        if (i % 2) == 0:
            # show field
            frame_field = visualizer.render_field(1.0)
            cv2.imshow("Wave Simulation Field", frame_field)

            # show intensity
            # frame_int = visualizer.render_intensity(1.0)
            # cv2.imshow("Wave Simulation Intensity", frame_int)

        cv2.waitKey(1)


if __name__ == "__main__":
    main()



================================================
FILE: wave_sim2d/examples/example3.py
================================================
import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '../'))  # noqa

import numpy as np
import cupy as cp
import math
import cv2

import wave_sim2d.wave_visualizer as vis
import wave_sim2d.wave_simulation as sim
from wave_sim2d.scene_objects.static_dampening import StaticDampening
from wave_sim2d.scene_objects.static_refractive_index import StaticRefractiveIndex


def gaussian_kernel(size, sigma):
    """
    creates gaussian kernel with side length `l` and a sigma of `sig`
    """
    ax = np.linspace(-(size - 1) / 2., (size - 1) / 2., size)
    gauss = np.exp(-0.5 * np.square(ax) / np.square(sigma))
    kernel = np.outer(gauss, gauss)
    return kernel / np.sum(kernel)


class MovingCharge(sim.SceneObject):
    """
    Implements a point source scene object. The amplitude can be optionally modulated using a modulator object.
    :param x: center position x.
    :param y: center position y.
    :param frequency: motion frequency
    :param amplitude: motion amplitude
    """
    def __init__(self, x, y, frequency, amplitude):
        self.x = x
        self.y = y
        self.frequency = frequency
        self.amplitude = amplitude
        self.size = 11

        # create a smooth source shape
        self.source_array = cp.array(gaussian_kernel(self.size, self.size/3))

    def render(self, field, wave_speed_field, dampening_field):
        # no changes to the refractive index or dampening field required for this class
        pass

    def render_visualization(self, image: np.ndarray):
        pass

    def update_field(self, field, t):
        fade_in = math.sin(min(t*0.1, math.pi/2))

        # write the moving charge to the field
        x = self.x + math.sin(self.frequency * t*0.05)*200
        y = self.y + math.sin(self.frequency * t)*self.amplitude

        # copy source shape to current position into field
        wh = self.source_array.shape[1]//2
        hh = self.source_array.shape[0]//2
        field[y-hh:y+hh+1, x-wh:x+wh+1] += self.source_array * fade_in * 0.25


def build_scene():
    """
    In this example, a custom scene object is implemented and used to simulate a moving field disturbance.
    """
    width = 600
    height = 600
    objects = []

    # Add a static dampening field without any dampending in the interior (value 1.0 means no dampening)
    # However a dampening layer at the border is added to avoid reflections (see parameter 'border thickness')
    objects.append(StaticDampening(np.ones((height, width)), 64))

    # add a constant refractive index field
    objects.append(StaticRefractiveIndex(np.full((height, width), 1.5)))

    # add a simple point source
    objects.append(MovingCharge(300, 300, 0.1, 10))

    return objects, width, height


def main():
    # create colormaps
    field_colormap = vis.get_colormap_lut('colormap_wave1', invert=False, black_level=-0.05)
    intensity_colormap = vis.get_colormap_lut('afmhot', invert=False, black_level=0.0)

    # reset random number generator
    np.random.seed(0)

    # build simulation scene
    scene_objects, w, h = build_scene()

    # create simulator and visualizer objects
    simulator = sim.WaveSimulator2D(w, h, scene_objects)
    visualizer = vis.WaveVisualizer(field_colormap=field_colormap, intensity_colormap=intensity_colormap)

    # run simulation
    for i in range(8000):
        simulator.update_scene()
        simulator.update_field()
        visualizer.update(simulator)

        # visualize very N frames
        if (i % 2) == 0:
            # show field
            frame_field = visualizer.render_field(1.0)
            cv2.imshow("Wave Simulation Field", frame_field)

            # show intensity
            # frame_int = visualizer.render_intensity(1.0)
            # cv2.imshow("Wave Simulation Intensity", frame_int)

        cv2.waitKey(1)


if __name__ == "__main__":
    main()



================================================
FILE: wave_sim2d/examples/example4.py
================================================
import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '../'))  # noqa

import cv2
import numpy as np
import cupy as cp
import wave_sim2d.wave_visualizer as vis
import wave_sim2d.wave_simulation as sim
from wave_sim2d.scene_objects.source import *
from wave_sim2d.scene_objects.static_refractive_index import *
from wave_sim2d.scene_objects.static_dampening import *


def build_scene():
    """
    This example creates fabry pirot cavity and shows the standing waves
    """
    width = 768
    height = 512
    objects = []

    # Add a static dampening field without any dampening in the interior (value 1.0 means no dampening)
    # However a dampening layer at the border is added to avoid reflections (see parameter 'border thickness')
    objects.append(StaticDampening(np.ones((height, width)), 48))

    # add nonlinear refractive index field
    objects.append(StaticRefractiveIndexBox((50, height//2), (50, int(height*0.8)), 0.0, 100.0))
    objects.append(StaticRefractiveIndexBox((width-180, height//2), (40, int(height*0.8)), 0.0, 10.0))

    # add a point source with an amplitude modulator
    # objects.append(LineSource((77, height//2-140), (77, height//2+140), 0.0215, amplitude=0.5))
    objects.append(LineSource((77, height//2-140), (77, height//2+140), 0.1003, amplitude=0.3))

    return objects, width, height


def show_field(field, brightness_scale):
    gray = (cp.clip(field*brightness_scale, -1.0, 1.0) * 127 + 127).astype(np.uint8)
    img = gray.get()
    cv2.imshow("Strain Simulation Field", cv2.cvtColor(img, cv2.COLOR_RGB2BGR))


def main():
    write_videos = False
    write_video_frame_every = 2

    # create colormaps
    field_colormap = vis.get_colormap_lut('colormap_wave1', invert=False, black_level=-0.05)
    intensity_colormap = vis.get_colormap_lut('afmhot', invert=False, black_level=0.0)

    # build simulation scene
    scene_objects, w, h = build_scene()

    # create simulator and visualizer objects
    simulator = sim.WaveSimulator2D(w, h, scene_objects)
    visualizer = vis.WaveVisualizer(field_colormap=field_colormap, intensity_colormap=intensity_colormap)

    # optional create video writers
    if write_videos:
        video_writer1 = cv2.VideoWriter('simulation_field.avi', cv2.VideoWriter_fourcc(*'FFV1'), 60, (w, h))
        video_writer2 = cv2.VideoWriter('simulation_intensity.avi', cv2.VideoWriter_fourcc(*'FFV1'), 60, (w, h))

    # run simulation
    for i in range(100000):
        simulator.update_scene()
        simulator.update_field()

        visualizer.update(simulator)
        # show field
        frame_field = visualizer.render_field(1.0)
        cv2.imshow("Wave Simulation Field", frame_field)

        # show intensity
        frame_int = visualizer.render_intensity(1.0)
        # cv2.imshow("Wave Simulation Intensity", frame_int)

        if write_videos and (i % write_video_frame_every) == 0:
            video_writer1.write(frame_field)
            video_writer2.write(frame_int)

        cv2.waitKey(1)


if __name__ == "__main__":
    main()



================================================
FILE: wave_sim2d/main.py
================================================
if __name__ == "__main__":
  print('please run one of the examples from the source/example folder...')


================================================
FILE: wave_sim2d/scene_objects/source.py
================================================
from wave_sim2d.wave_simulation import SceneObject
import cupy as cp
import numpy as np
import math


class PointSource(SceneObject):
    """
    Implements a point source scene object. The amplitude can be optionally modulated using a modulator object.
    :param x: source position x.
    :param y: source position y.
    :param frequency: emitting frequency.
    :param amplitude: emitting amplitude, not used when an amplitude modulator is given
    :param phase: emitter phase
    :param amp_modulator: optional amplitude modulator. This can be used to change the amplitude of the source
                          over time.
    """
    def __init__(self, x, y, frequency, amplitude=1.0, phase=0, amp_modulator=None):
        self.x = x
        self.y = y
        self.frequency = frequency
        self.amplitude = amplitude
        self.phase = phase
        self.amplitude_modulator = amp_modulator

    def set_amplitude_modulator(self, func):
        self.amplitude_modulator = func

    def render(self, field: cp.ndarray, wave_speed_field: cp.ndarray, dampening_field: cp.ndarray):
        pass

    def update_field(self, field, t):
        if self.amplitude_modulator is not None:
            amplitude = self.amplitude_modulator(t) * self.amplitude
        else:
            amplitude = self.amplitude

        v = cp.sin(self.phase + self.frequency * t) * amplitude
        field[self.y, self.x] = v

    def render_visualization(self, image: np.ndarray):
        """ renders a visualization of the scene object to the image """
        pass


class LineSource(SceneObject):
    """
    Implements a line source scene object. The amplitude can be optionally modulated using a modulator object.
    The source emits along a line defined by a start and end point.
    :param start: starting (x, y) coordinates of the line as a tuple.
    :param end: ending (x, y) coordinates of the line as a tuple.
    :param frequency: emitting frequency.
    :param amplitude: emitting amplitude, not used when an amplitude modulator is given
    :param phase: emitter phase
    :param amp_modulator: optional amplitude modulator. This can be used to change the amplitude of the source
                          over time.
    """
    def __init__(self, start, end, frequency, amplitude=1.0, phase=0, amp_modulator=None):
        self.start = start
        self.end = end
        self.frequency = frequency
        self.amplitude = amplitude
        self.phase = phase
        self.amplitude_modulator = amp_modulator

    def set_amplitude_modulator(self, func):
        self.amplitude_modulator = func

    def render(self, field: cp.ndarray, wave_speed_field: cp.ndarray, dampening_field: cp.ndarray):
        pass

    def update_field(self, field, t):
        if self.amplitude_modulator is not None:
            amplitude = self.amplitude_modulator(t) * self.amplitude
        else:
            amplitude = self.amplitude

        v = cp.sin(self.phase + self.frequency * t) * amplitude

        # Determine the points along the line using NumPy
        x1, y1 = self.start
        x2, y2 = self.end

        distance = np.sqrt((x2 - x1)**2 + (y2 - y1)**2)
        num_points = int(distance) + 1

        if num_points > 0:
            x_coords = cp.linspace(x1, x2, num_points).round().astype(int)
            y_coords = cp.linspace(y1, y2, num_points).round().astype(int)

            # Create boolean masks for valid indices
            valid_x = (x_coords >= 0) & (x_coords < field.shape[1])
            valid_y = (y_coords >= 0) & (y_coords < field.shape[0])
            valid_indices = valid_x & valid_y

            # Use these valid indices to update the field directly
            valid_y_coords = y_coords[valid_indices]
            valid_x_coords = x_coords[valid_indices]
            field[valid_y_coords, valid_x_coords] = v

    def render_visualization(self, image: np.ndarray):
        """ renders a visualization of the scene object to the image """
        pass

# --- Modulators -------------------------------------------------------------------------------------------------------


class ModulatorSmoothSquare:
    """
    A modulator that creates a smoothed square wave
    """
    def __init__(self, frequency, phase, smoothness=0.5):
        self.frequency = frequency
        self.phase = phase
        self.smoothness = min(max(smoothness, 1e-4), 1.0)

    def __call__(self, t):
        s = math.pow(self.smoothness, 4.0)
        a = (0.5 / math.atan(1.0/s)) * math.atan(math.sin(t * self.frequency + self.phase) / s)+0.5
        return a


class ModulatorDiscreteSignal:
    """
    A modulator that creates a smoothed binary signal
    """
    def __init__(self, signal_array, time_factor, transition_slope=8.0):
        self.signal_array = signal_array
        self.time_factor = time_factor
        self.transition_slope = transition_slope

    def __call__(self, t):
        def smooth_step(t):
            return t * t * (3 - 2 * t)

        # Wrap around the position if it's outside the array range
        sl = len(self.signal_array)
        t = math.fmod(t*self.time_factor, sl)

        # Find the indices of the neighboring values
        index_low = int(t)
        index_high = (index_low + 1) % sl

        # Calculate the interpolation factor
        tf = (t - index_low)
        tf = max(0.0, min(1.0, (tf-0.5)*self.transition_slope+0.5))

        # Use smooth step to interpolate between neighboring values
        l = smooth_step(tf)
        interpolated_value = (1 - l) * self.signal_array[index_low] + l * self.signal_array[index_high]

        return interpolated_value


================================================
FILE: wave_sim2d/scene_objects/static_dampening.py
================================================
from wave_sim2d.wave_simulation import SceneObject
import cupy as cp
import numpy as np


class StaticDampening(SceneObject):
    """
    Implements a static dampening field that overwrites the entire domain.
    Therefore, us this as base layer in your scene.
    """

    def __init__(self, dampening_field, border_thickness):
        """
        Creates a static dampening field object
        @param dampening_field: A NxM array with dampening factors (1.0 equals no dampening) of the same size as the simulation domain.
        @param pml_thickness: Thickness of the Perfectly Matched Layer (PML) at the boundaries to prevent reflections.
        """
        w = dampening_field.shape[1]
        h = dampening_field.shape[0]
        self.d = cp.ones((h, w), dtype=cp.float32)
        self.d = cp.clip(cp.array(dampening_field), 0.0, 1.0)

        # apply border dampening
        for i in range(border_thickness):
            v = (i / border_thickness) ** 0.5
            self.d[i, i:w - i] = v
            self.d[-(1 + i), i:w - i] = v
            self.d[i:h - i, i] = v
            self.d[i:h - i, -(1 + i)] = v

    def render(self, field: cp.ndarray, wave_speed_field: cp.ndarray, dampening_field: cp.ndarray):
        assert (dampening_field.shape == self.d.shape)

        # overwrite existing dampening field
        dampening_field[:] = self.d

    def update_field(self, field: cp.ndarray, t):
        pass

    def render_visualization(self, image: np.ndarray):
        """ renders a visualization of the scene object to the image """
        pass


================================================
FILE: wave_sim2d/scene_objects/static_image_scene.py
================================================
from wave_sim2d.wave_simulation import SceneObject

import numpy as np
import cupy as cp
from wave_sim2d.scene_objects.static_dampening import StaticDampening
from wave_sim2d.scene_objects.static_refractive_index import StaticRefractiveIndex


class StaticImageScene(SceneObject):
    """
    Implements static scene, where the RGB channels of the input image encode the refractive index, the dampening and sources.
    This class allows to use an image editor to create scenes.
    """
    def __init__(self, scene_image, source_amplitude=1.0, source_fequency_scale=1.0):
        """
        load source from an image description
        The simulation scenes are given as an 8Bit RGB image with the following channel semantics:
            * Red:   The Refractive index times 100 (for refractive index 1.5 you would use value 150)
            * Green: Each pixel with a green value above 0 is a sinusoidal wave source. The green value
                     defines its frequency. WARNING: Do not use antialiasing for the green channel !
            * Blue:  Absorbtion field. Larger values correspond to higher dampening of the waves,
                     use graduated transitions to avoid reflections
        """
        # Set the opacity of source pixels to incoming waves. If the opacity is 0.0
        # the field will be completely overwritten by the source term
        # a nonzero value (e.g 0.5) allows for antialiasing of sources to work
        self.source_opacity = 0.9

        # set refractive index field
        self.refractive_index = StaticRefractiveIndex(scene_image[:, :, 0] / 100)

        # set absorber field
        self.dampening = StaticDampening(1.0 - scene_image[:, :, 2] / 255, border_thickness=48)

        # set sources, each entry describes a source with the following parameters:
        # (x, y, phase, amplitude, frequency)
        sources_pos = np.flip(np.argwhere(scene_image[:, :, 1] > 0), axis=1)
        phase_amplitude_freq = np.tile(np.array([0, source_amplitude, 0.3]), (sources_pos.shape[0], 1))
        self.sources = np.concatenate((sources_pos, phase_amplitude_freq), axis=1)

        # set source frequency to channel value
        self.sources[:, 4] = scene_image[sources_pos[:, 1], sources_pos[:, 0], 1] / 255 * 0.5 * source_fequency_scale
        self.sources = cp.array(self.sources).astype(cp.float32)

    def render(self, field: cp.ndarray, wave_speed_field: cp.ndarray, dampening_field: cp.ndarray):
        """
        render the stat
        """
        self.dampening.render(field, wave_speed_field, dampening_field)
        self.refractive_index.render(field, wave_speed_field, dampening_field)

    def update_field(self, field: cp.ndarray, t):
        # Update the sources in the simulation field based on their properties.
        v = cp.sin(self.sources[:, 2]+self.sources[:, 4]*t)*self.sources[:, 3]
        coords = self.sources[:, 0:2].astype(cp.int32)

        o = self.source_opacity
        field[coords[:, 1], coords[:, 0]] = field[coords[:, 1], coords[:, 0]]*o + v*(1.0-o)

    def render_visualization(self, image: np.ndarray):
        """ renders a visualization of the scene object to the image """
        pass


================================================
FILE: wave_sim2d/scene_objects/static_refractive_index.py
================================================
from wave_sim2d.wave_simulation import SceneObject
import cupy as cp
import numpy as np
import cv2


class StaticRefractiveIndex(SceneObject):
    """
    Implements a static refractive index field that overwrites the entire domain with a constant IOR value.
    Use this as base layer in your scene.
    """

    def __init__(self, refractive_index_field):
        """
        Creates a static refractive index field object
        :param refractive_index_field: The refractive index field, same size as the source.
                                       Note that values below 0.9 are clipped to prevent the simulation
                                       from becoming instable
        """
        shape = refractive_index_field.shape
        self.c = cp.ones((shape[0], shape[1]), dtype=cp.float32)
        self.c = 1.0/cp.clip(cp.array(refractive_index_field), 0.9, 10.0)

    def render(self, field: cp.ndarray, wave_speed_field: cp.ndarray, dampening_field: cp.ndarray):
        assert (wave_speed_field.shape == self.c.shape)
        wave_speed_field[:] = self.c

    def update_field(self, field: cp.ndarray, t):
        pass

    def render_visualization(self, image: np.ndarray):
        """ renders a visualization of the scene object to the image """
        pass


class StaticRefractiveIndexPolygon(SceneObject):
    """
    Draws a static polygon with a given refractive index into the wave_speed_field using an
    anti-aliased mask and indexing. Caches the pixel coordinates and mask values.
    """

    def __init__(self, vertices, refractive_index):
        """
        Initializes the StaticRefractiveIndexPolygon.

        Args:
            vertices (list or np.ndarray): A list or array of (x, y) coordinates defining the polygon.
            refractive_index (float): The refractive index of the polygon. Values are clamped to [0.9, 10.0].
        """
        self.vertices = np.array(vertices, dtype=np.float32)
        self.refractive_index = min(max(refractive_index, 0.9), 10.0)
        self._cached_coords = None
        self._cached_mask_values = None
        self._cached_field_shape = (0, 0)

    def _create_polygon_data(self, field_shape):
        """
        Creates and caches the pixel coordinates and anti-aliased mask values for the polygon.

        Args:
            field_shape (tuple): The shape (rows, cols) of the simulation field.

        Returns:
            tuple: A tuple containing:
                - coords (tuple of cp.ndarray): (y_coordinates, x_coordinates) of the polygon pixels within the field.
                - mask_values (cp.ndarray): Corresponding anti-aliased mask values (0.0 to 1.0).
        """
        if self._cached_coords is not None and self._cached_field_shape == field_shape:
            return self._cached_coords, self._cached_mask_values

        rows, cols = field_shape

        # Find the bounding box of the polygon
        min_x = np.min(self.vertices[:, 0])
        max_x = np.max(self.vertices[:, 0])
        min_y = np.min(self.vertices[:, 1])
        max_y = np.max(self.vertices[:, 1])

        mask_width = int(np.ceil(max_x - min_x)) + 1
        mask_height = int(np.ceil(max_y - min_y)) + 1
        offset_x = int(np.floor(min_x))
        offset_y = int(np.floor(min_y))

        # Create the mask
        mask = np.zeros((mask_height, mask_width), dtype=np.float32)
        translated_vertices = self.vertices - [offset_x, offset_y]
        translated_vertices_cv = np.round(translated_vertices).astype(np.int32)
        cv2.fillPoly(mask, [translated_vertices_cv], 1.0, lineType=cv2.LINE_AA)

        # Get coordinates and mask values of non-black pixels
        coords_y, coords_x = np.where(mask > 0)
        mask_values = mask[coords_y, coords_x]

        # Adjust coordinates to the position in the main field
        global_coords_y = coords_y + offset_y
        global_coords_x = coords_x + offset_x

        # Perform out-of-bounds check here
        in_bounds = (global_coords_y >= 0) & (global_coords_y < rows) & \
                    (global_coords_x >= 0) & (global_coords_x < cols)

        valid_global_y = global_coords_y[in_bounds]
        valid_global_x = global_coords_x[in_bounds]
        valid_mask_values = mask_values[in_bounds]

        self._cached_coords = (cp.array(valid_global_y), cp.array(valid_global_x))
        self._cached_mask_values = cp.array(valid_mask_values, dtype=cp.float32)
        self._cached_field_shape = field_shape
        return self._cached_coords, self._cached_mask_values

    def render(self, field: cp.ndarray, wave_speed_field: cp.ndarray, dampening_field: cp.ndarray):
        coords, mask_values = self._create_polygon_data(wave_speed_field.shape)

        # Use advanced indexing to update the field and perform alpha blending
        bg_wave_speed = wave_speed_field[coords[0], coords[1]]
        wave_speed_field[coords[0], coords[1]] = (bg_wave_speed * (1.0 - mask_values) +
                                                  mask_values / self.refractive_index)

    def update_field(self, field: cp.ndarray, t):
        pass

    def render_visualization(self, image: np.ndarray):
        vertices = np.round(self.vertices).astype(np.int32)
        cv2.fillPoly(image, [vertices], (60, 60, 60), lineType=cv2.LINE_AA)


class StaticRefractiveIndexBox(StaticRefractiveIndexPolygon):
    """
    Draws a static rotated box with a given refractive index into the wave_speed_field by
    inheriting from StaticRefractiveIndexPolygon.
    """

    def __init__(self, center: tuple, box_size: tuple, box_angle_rad: float, refractive_index: float):
        """
        Initializes the StaticRefractiveIndexBox.

        Args:
            center (tuple): A tuple (center_x, center_y) representing the box's center.
            box_size (tuple): A tuple (width, height) representing the box's dimensions.
            box_angle_rad (float): The rotation angle of the box in radians (counter-clockwise).
            refractive_index (float): The refractive index of the box. Values are clamped to [0.9, 10.0].
        """
        self.center = center
        self.box_size = box_size
        self.box_angle_rad = box_angle_rad
        refractive_index = min(max(refractive_index, 0.9), 10.0)

        # Unpack center and box size
        center_x, center_y = self.center
        width, height = self.box_size

        # Calculate the vertices of the rotated box
        half_width = width / 2
        half_height = height / 2
        local_vertices = np.array([[-half_width, -half_height],
                                   [half_width, -half_height],
                                   [half_width, half_height],
                                   [-half_width, half_height]], dtype=np.float32)

        # Create the rotation matrix
        rotation_matrix = cv2.getRotationMatrix2D((0, 0), np.rad2deg(self.box_angle_rad), 1)

        # Rotate the local vertices
        rotated_vertices = cv2.transform(np.array([local_vertices]), rotation_matrix)[0]

        # Translate the rotated vertices to the center
        translated_vertices = rotated_vertices + [center_x, center_y]

        # Initialize the parent class (StaticRefractiveIndexPolygon) with the vertices
        super().__init__(translated_vertices, refractive_index)


================================================
FILE: wave_sim2d/scene_objects/strain_refractive_index.py
================================================
from wave_sim2d.wave_simulation import SceneObject
import cupy as cp
import cupyx.scipy.signal
import numpy as np

class StrainRefractiveIndex(SceneObject):
    """
    Implements a dynamic refractive index field that linearly depends on the strain of the current field.
    The refractive index within the entire domain is overwritten
    """

    def __init__(self, refractive_index_offset, coupling_constant):
        """
        Creates a strain refractive index field object
        :param coupling_constant: coupling constant between the strain and the refractive index
        """
        self.coupling_constant = coupling_constant
        self.refractive_index_offset = refractive_index_offset

        self.du_dx_kernel = cp.array([[-1, 0.0, 1]])
        self.du_dy_kernel = cp.array([[-1], [0.0], [1]])

        self.strain_field = None

    def render(self, field: cp.ndarray, wave_speed_field: cp.ndarray, dampening_field: cp.ndarray):
        # compute strain
        du_dx = cupyx.scipy.signal.convolve2d(field, self.du_dx_kernel, mode='same', boundary='fill')
        du_dy = cupyx.scipy.signal.convolve2d(field, self.du_dy_kernel, mode='same', boundary='fill')

        self.strain_field = cp.sqrt(du_dx**2 + du_dy**2)

        # compute refractive index from strain
        refractive_index_field = self.refractive_index_offset + self.strain_field*self.coupling_constant

        # assign wave speed using refractive index from above
        wave_speed_field[:] = 1.0/cp.clip(cp.array(refractive_index_field), 0.9, 10.0)

    def update_field(self, field: cp.ndarray, t):
        pass

    def render_visualization(self, image: np.ndarray):
        """ renders a visualization of the scene object to the image """
        pass


================================================
FILE: wave_sim2d/wave_simulation.py
================================================
import cupy
import numpy as np
import cupy as cp
import cupyx.scipy.signal
from abc import ABC, abstractmethod


class SceneObject(ABC):
    """
    Interface for simulation scene objects. A scene object is anything defining or modifying the simulation scene.
    For example: Light sources, Absorbers or regions with specific refractive index. Scene objects can change the
    simulated field and draw their contribution to the wave speed field and dampening field each frame """

    @abstractmethod
    def render(self, field: cupy.ndarray, wave_speed_field: cupy.ndarray, dampening_field: cupy.ndarray):
        """ renders the scene objects contribution to the wave speed field and dampening field """
        pass

    @abstractmethod
    def update_field(self, field: cupy.ndarray, t):
        """ performs updates to the field itself, e.g. for adding sources """
        pass

    @abstractmethod
    def render_visualization(self, image: np.ndarray):
        """ renders a visualization of the scene object to the image """
        pass


class WaveSimulator2D:
    """
    Simulates the 2D wave equation
    The system assumes units, where the wave speed is 1.0 pixel/timestep
    source frequency should be adjusted accordingly
    """
    def __init__(self, w, h, scene_objects, initial_field=None):
        """
        Initialize the 2D wave simulator.
        @param w: Width of the simulation grid.
        @param h: Height of the simulation grid.
        """
        self.global_dampening = 1.0
        self.c = cp.ones((h, w), dtype=cp.float32)                      # wave speed field (from refractive indices)
        self.d = cp.ones((h, w), dtype=cp.float32)                      # dampening field
        self.u = cp.zeros((h, w), dtype=cp.float32)                     # field values
        self.u_prev = cp.zeros((h, w), dtype=cp.float32)                # field values of prev frame

        if initial_field is not None:
            assert w == initial_field.shape[1] and h == initial_field.shape[2], 'width/height of initial field invalid'
            self.u[:] = initial_field
            self.u_prev[:] = initial_field

        # Define Laplacian kernel
        self.laplacian_kernel = cp.array([[0.066, 0.184, 0.066],
                                          [0.184, -1.0, 0.184],
                                          [0.066, 0.184, 0.066]])

        # self.laplacian_kernel = cp.array([[0.05, 0.2, 0.05],
        #                           [0.2, -1.0, 0.2],
        #                           [0.05, 0.2, 0.05]])

        # self.laplacian_kernel = cp.array([[0.103, 0.147, 0.103],
        #                                   [0.147, -1.0, 0.147],
        #                                   [0.103, 0.147, 0.103]])

        self.t = 0
        self.dt = 1.0

        self.scene_objects = scene_objects if scene_objects is not None else []

    def reset_time(self):
        """
        Reset the simulation time to zero.
        """
        self.t = 0.0

    def update_field(self):
        """
        Update the simulation field based on the wave equation.
        """
        # calculate laplacian using convolution
        laplacian = cupyx.scipy.signal.convolve2d(self.u, self.laplacian_kernel, mode='same', boundary='fill')

        # update field
        v = (self.u - self.u_prev) * self.d * self.global_dampening
        r = (self.u + v + laplacian * (self.c * self.dt)**2)

        self.u_prev[:] = self.u
        self.u[:] = r

        self.t += self.dt

    def update_scene(self):
        # clear wave speed field and dampening field
        self.c.fill(1.0)
        self.d.fill(1.0)

        for obj in self.scene_objects:
            obj.render(self.u, self.c, self.d)

        for obj in self.scene_objects:
            obj.update_field(self.u, self.t)

    def get_field(self):
        """
        Get the current state of the simulation field.
        @return: A 2D array representing the simulation field.
        """
        return self.u

    def render_visualization(self, image=None):
        # clear wave speed field and dampening field
        if image is None:
            image = np.zeros((self.c.shape[0], self.c.shape[1], 3), dtype=np.uint8)

        for obj in self.scene_objects:
            obj.render_visualization(image)

        return image




================================================
FILE: wave_sim2d/wave_visualizer.py
================================================
import numpy as np
import cupy as cp
import cv2
import matplotlib.pyplot

colormap_icefire = [[179, 224, 216], [178, 223, 216], [176, 222, 215], [175, 221, 215], [173, 219, 214], [171, 218, 214], [169, 217, 214], [167, 215, 213], [165, 214, 213], [162, 212, 212], [160, 210, 212], [157, 209, 211], [154, 207, 211], [151, 205, 210], [148, 203, 210], [146, 201, 209], [143, 199, 209], [140, 198, 208], [137, 196, 208], [134, 194, 208], [131, 192, 207], [128, 190, 207], [125, 188, 207], [122, 187, 207], [119, 185, 206], [116, 183, 206], [113, 181, 206], [110, 179, 206], [108, 177, 206], [105, 176, 205], [102, 174, 205], [99, 172, 205], [97, 170, 205], [94, 168, 205], [91, 166, 205], [89, 164, 205], [86, 162, 205], [84, 161, 205], [82, 159, 205], [79, 157, 205], [77, 155, 205], [75, 153, 206], [73, 151, 206], [71, 149, 206], [69, 147, 206], [68, 145, 206], [66, 143, 206], [65, 140, 206], [64, 138, 206], [63, 136, 206], [62, 134, 206], [61, 132, 206], [61, 130, 205], [61, 127, 205], [60, 125, 205], [60, 123, 204], [60, 121, 203], [60, 118, 203], [61, 116, 202], [61, 114, 201], [61, 112, 200], [62, 109, 198], [62, 107, 197], [63, 105, 195], [64, 103, 194], [65, 100, 192], [65, 98, 190], [66, 96, 187], [67, 94, 185], [67, 92, 183], [68, 90, 180], [68, 88, 177], [69, 86, 174], [69, 85, 171], [69, 83, 168], [70, 81, 165], [70, 79, 162], [70, 78, 158], [69, 76, 155], [69, 75, 151], [69, 73, 148], [68, 72, 144], [68, 70, 141], [67, 69, 137], [66, 67, 134], [66, 66, 130], [65, 65, 127], [64, 63, 123], [63, 62, 120], [62, 61, 116], [61, 60, 113], [60, 59, 109], [59, 57, 106], [58, 56, 103], [57, 55, 99], [55, 54, 96], [54, 53, 93], [53, 52, 90], [52, 50, 87], [51, 49, 84], [50, 48, 81], [48, 47, 78], [47, 46, 75], [46, 45, 72], [45, 44, 70], [44, 43, 67], [43, 42, 65], [42, 41, 62], [41, 40, 60], [40, 39, 57], [39, 38, 55], [38, 37, 53], [37, 37, 51], [37, 36, 49], [36, 35, 47], [35, 35, 45], [35, 34, 44], [34, 33, 42], [34, 33, 41], [33, 32, 39], [33, 32, 38], [33, 32, 37], [33, 31, 36], [33, 31, 35], [33, 31, 35], [34, 30, 34], [34, 30, 33], [34, 30, 33], [35, 30, 32], [36, 30, 32], [36, 30, 32], [37, 30, 32], [38, 30, 32], [39, 30, 32], [40, 30, 32], [41, 30, 32], [42, 30, 33], [44, 31, 33], [46, 31, 34], [47, 31, 34], [49, 31, 35], [51, 32, 35], [53, 32, 36], [55, 32, 37], [57, 33, 38], [59, 33, 38], [61, 33, 39], [63, 34, 40], [65, 34, 41], [67, 35, 42], [70, 35, 43], [72, 36, 44], [74, 36, 45], [77, 37, 46], [79, 37, 47], [82, 38, 48], [84, 38, 49], [87, 39, 50], [90, 39, 51], [92, 40, 52], [95, 40, 53], [98, 40, 54], [100, 41, 55], [103, 41, 56], [106, 42, 57], [109, 42, 58], [111, 42, 59], [114, 43, 60], [117, 43, 60], [120, 43, 61], [123, 44, 62], [126, 44, 63], [129, 44, 63], [131, 44, 64], [134, 45, 64], [137, 45, 65], [140, 45, 65], [143, 46, 65], [146, 46, 65], [149, 46, 66], [152, 47, 66], [155, 47, 66], [158, 48, 66], [160, 48, 66], [163, 49, 65], [166, 49, 65], [169, 50, 65], [172, 51, 64], [174, 52, 64], [177, 53, 63], [180, 54, 63], [182, 55, 62], [185, 56, 62], [187, 57, 61], [190, 58, 61], [192, 60, 60], [195, 61, 59], [197, 63, 59], [199, 65, 58], [201, 66, 57], [203, 68, 57], [206, 70, 56], [208, 72, 55], [209, 74, 55], [211, 76, 54], [213, 78, 54], [215, 81, 54], [217, 83, 53], [218, 85, 53], [220, 88, 53], [221, 90, 53], [223, 93, 54], [224, 95, 54], [225, 98, 55], [227, 101, 55], [228, 103, 56], [229, 106, 57], [230, 109, 58], [231, 111, 60], [232, 114, 61], [233, 117, 62], [234, 120, 64], [235, 123, 66], [236, 125, 68], [237, 128, 70], [237, 131, 73], [238, 134, 75], [239, 137, 78], [240, 139, 80], [240, 142, 83], [241, 145, 86], [242, 148, 89], [242, 151, 93], [243, 153, 96], [243, 156, 99], [244, 159, 103], [245, 162, 106], [245, 165, 110], [246, 167, 113], [246, 170, 117], [247, 173, 120], [247, 176, 124], [248, 178, 127], [248, 181, 131], [249, 184, 134], [249, 186, 138], [250, 188, 141], [250, 190, 144], [251, 192, 147], [251, 194, 149], [251, 196, 152], [252, 198, 154], [252, 200, 156], [252, 201, 158], [253, 203, 160]]
colormap_wave1 = [[255, 255, 255], [254, 254, 253], [254, 253, 252], [253, 252, 250], [253, 250, 248], [252, 249, 246], [252, 248, 244], [251, 246, 242], [251, 245, 240], [250, 243, 237], [250, 242, 235], [249, 240, 232], [248, 238, 230], [248, 237, 227], [247, 235, 224], [247, 233, 221], [246, 231, 218], [245, 229, 215], [245, 227, 212], [244, 225, 209], [243, 223, 206], [242, 221, 203], [242, 219, 200], [241, 217, 196], [240, 215, 193], [239, 213, 190], [239, 211, 186], [238, 208, 183], [237, 206, 179], [236, 204, 176], [235, 202, 172], [234, 199, 169], [233, 197, 165], [232, 195, 162], [231, 192, 158], [230, 190, 155], [230, 188, 151], [228, 185, 148], [227, 183, 144], [226, 181, 141], [225, 178, 137], [224, 176, 134], [223, 174, 130], [222, 171, 127], [221, 169, 124], [219, 167, 120], [218, 164, 117], [217, 162, 114], [216, 160, 111], [214, 157, 108], [213, 155, 105], [212, 153, 102], [210, 151, 99], [209, 149, 96], [207, 146, 93], [206, 144, 91], [204, 142, 88], [203, 140, 85], [201, 138, 83], [199, 136, 80], [198, 134, 78], [196, 132, 76], [194, 130, 74], [193, 128, 72], [191, 127, 70], [189, 125, 69], [187, 123, 67], [185, 121, 65], [183, 119, 63], [180, 117, 62], [178, 115, 60], [176, 113, 58], [173, 111, 56], [170, 109, 55], [168, 107, 53], [165, 105, 52], [162, 102, 50], [159, 100, 49], [156, 98, 47], [154, 96, 46], [151, 94, 44], [148, 92, 43], [144, 90, 41], [141, 87, 40], [138, 85, 39], [135, 83, 37], [132, 81, 36], [129, 79, 35], [125, 76, 34], [122, 74, 33], [119, 72, 31], [115, 70, 30], [112, 68, 29], [109, 66, 28], [105, 64, 27], [102, 62, 27], [99, 60, 26], [96, 58, 25], [93, 56, 25], [89, 54, 25], [86, 52, 25], [83, 51, 25], [80, 49, 25], [77, 47, 25], [74, 45, 25], [71, 44, 25], [68, 42, 25], [65, 41, 25], [62, 39, 25], [60, 38, 25], [57, 37, 25], [54, 35, 25], [52, 34, 25], [49, 33, 25], [47, 32, 25], [45, 31, 25], [43, 30, 25], [40, 29, 25], [39, 28, 25], [37, 28, 25], [35, 27, 25], [33, 27, 25], [32, 26, 25], [30, 26, 25], [29, 25, 25], [28, 25, 25], [27, 25, 25], [26, 25, 25], [26, 25, 26], [26, 26, 27], [26, 26, 28], [26, 26, 30], [26, 27, 31], [26, 27, 33], [26, 28, 34], [26, 29, 36], [26, 30, 38], [26, 31, 40], [26, 32, 42], [26, 33, 44], [26, 34, 47], [26, 35, 49], [26, 37, 51], [26, 38, 54], [26, 40, 56], [26, 41, 59], [26, 43, 62], [26, 44, 64], [26, 46, 67], [26, 48, 70], [26, 50, 73], [27, 51, 76], [28, 53, 79], [28, 55, 82], [29, 57, 85], [30, 59, 88], [31, 61, 91], [32, 64, 94], [33, 66, 97], [35, 68, 101], [36, 70, 104], [37, 72, 107], [38, 74, 110], [40, 77, 113], [41, 79, 117], [42, 81, 120], [44, 84, 123], [45, 86, 126], [47, 88, 130], [48, 91, 133], [50, 93, 136], [51, 95, 139], [53, 98, 142], [54, 100, 145], [56, 102, 148], [58, 104, 151], [59, 107, 154], [61, 109, 157], [63, 111, 160], [64, 114, 163], [66, 116, 165], [68, 118, 168], [70, 120, 171], [71, 122, 173], [73, 125, 176], [75, 127, 178], [77, 129, 181], [78, 131, 183], [80, 133, 185], [82, 135, 187], [84, 136, 189], [86, 138, 191], [87, 140, 193], [89, 142, 194], [91, 144, 196], [93, 146, 198], [96, 147, 199], [98, 149, 201], [100, 151, 203], [103, 153, 204], [105, 155, 206], [108, 157, 207], [110, 160, 209], [113, 162, 210], [116, 164, 212], [118, 166, 213], [121, 168, 214], [124, 170, 216], [127, 172, 217], [130, 174, 218], [133, 176, 219], [136, 178, 221], [139, 180, 222], [142, 183, 223], [145, 185, 224], [148, 187, 225], [152, 189, 226], [155, 191, 227], [158, 193, 228], [161, 195, 230], [164, 197, 231], [168, 200, 231], [171, 202, 232], [174, 204, 233], [177, 206, 234], [180, 208, 235], [183, 210, 236], [187, 212, 237], [190, 214, 238], [193, 216, 239], [196, 218, 239], [199, 220, 240], [202, 222, 241], [205, 223, 242], [208, 225, 242], [211, 227, 243], [214, 229, 244], [217, 231, 245], [219, 232, 245], [222, 234, 246], [225, 236, 247], [227, 237, 247], [230, 239, 248], [232, 240, 248], [234, 242, 249], [237, 243, 250], [239, 245, 250], [241, 246, 251], [243, 247, 251], [245, 249, 252], [247, 250, 252], [249, 251, 253], [251, 252, 253], [252, 253, 254], [254, 254, 254]]
colormap_wave2 = [[255, 255, 255], [253, 254, 254], [252, 254, 253], [250, 253, 252], [248, 253, 252], [246, 252, 251], [244, 252, 250], [242, 251, 249], [240, 251, 247], [237, 250, 246], [235, 250, 245], [232, 249, 244], [230, 248, 243], [227, 248, 242], [224, 247, 240], [221, 247, 239], [218, 246, 238], [215, 245, 236], [212, 245, 235], [209, 244, 234], [206, 243, 232], [203, 242, 231], [200, 242, 229], [196, 241, 228], [193, 240, 226], [190, 239, 225], [186, 239, 223], [183, 238, 222], [179, 237, 220], [176, 236, 218], [172, 235, 217], [169, 234, 215], [165, 233, 213], [162, 232, 212], [158, 231, 210], [155, 231, 208], [151, 230, 206], [148, 228, 205], [144, 227, 203], [141, 226, 201], [137, 225, 199], [134, 224, 198], [130, 223, 196], [127, 222, 194], [124, 221, 192], [120, 219, 190], [117, 218, 188], [114, 217, 187], [111, 216, 185], [108, 214, 183], [105, 213, 181], [102, 212, 179], [99, 210, 177], [96, 209, 176], [93, 207, 174], [91, 206, 172], [88, 204, 170], [85, 203, 168], [83, 201, 166], [80, 199, 164], [78, 198, 163], [76, 196, 161], [74, 194, 159], [72, 193, 157], [70, 191, 155], [69, 189, 153], [67, 187, 152], [65, 185, 149], [63, 183, 147], [62, 180, 145], [60, 178, 143], [59, 175, 141], [57, 173, 138], [56, 170, 136], [54, 167, 133], [53, 164, 131], [52, 161, 128], [50, 158, 125], [49, 155, 123], [48, 152, 120], [47, 149, 117], [46, 146, 115], [45, 143, 112], [43, 140, 109], [42, 136, 106], [41, 133, 103], [40, 130, 101], [39, 126, 98], [39, 123, 95], [38, 119, 92], [37, 116, 89], [36, 112, 86], [35, 109, 84], [35, 106, 81], [34, 102, 78], [33, 99, 75], [33, 95, 73], [32, 92, 70], [31, 89, 67], [31, 86, 65], [30, 82, 62], [30, 79, 60], [29, 76, 57], [29, 73, 55], [28, 70, 52], [28, 67, 50], [28, 64, 48], [27, 61, 46], [27, 58, 44], [27, 55, 42], [26, 53, 40], [26, 50, 38], [26, 48, 37], [26, 46, 35], [26, 43, 34], [26, 41, 32], [26, 39, 31], [26, 37, 30], [26, 35, 29], [26, 34, 28], [26, 32, 27], [26, 31, 26], [26, 29, 26], [26, 28, 25], [26, 27, 25], [26, 27, 25], [26, 26, 25], [26, 25, 25], [26, 25, 26], [26, 25, 26], [26, 25, 27], [26, 25, 27], [26, 25, 28], [27, 25, 30], [27, 25, 31], [27, 26, 32], [28, 27, 34], [28, 27, 35], [28, 28, 37], [29, 29, 39], [29, 30, 41], [30, 31, 43], [30, 32, 45], [31, 34, 48], [31, 35, 50], [32, 36, 53], [32, 38, 55], [33, 40, 58], [33, 41, 61], [34, 43, 64], [35, 45, 66], [36, 47, 69], [36, 49, 73], [37, 51, 76], [38, 53, 79], [39, 55, 82], [39, 57, 85], [40, 59, 89], [41, 61, 92], [42, 64, 95], [43, 66, 99], [44, 68, 102], [45, 71, 105], [46, 73, 109], [47, 76, 112], [48, 78, 116], [49, 80, 119], [50, 83, 122], [52, 86, 126], [53, 88, 129], [54, 91, 133], [55, 93, 136], [57, 96, 139], [58, 98, 143], [59, 100, 146], [60, 103, 149], [62, 105, 152], [63, 108, 155], [65, 110, 158], [66, 113, 161], [68, 115, 164], [69, 117, 167], [71, 120, 170], [72, 122, 173], [74, 124, 175], [75, 126, 178], [77, 128, 180], [79, 131, 183], [80, 133, 185], [82, 134, 187], [84, 136, 189], [86, 138, 191], [87, 140, 193], [89, 142, 194], [91, 144, 196], [93, 146, 198], [96, 147, 199], [98, 149, 201], [100, 151, 203], [103, 153, 204], [105, 155, 206], [108, 157, 207], [110, 160, 209], [113, 162, 210], [116, 164, 212], [118, 166, 213], [121, 168, 214], [124, 170, 216], [127, 172, 217], [130, 174, 218], [133, 176, 219], [136, 178, 221], [139, 180, 222], [142, 183, 223], [145, 185, 224], [148, 187, 225], [152, 189, 226], [155, 191, 227], [158, 193, 228], [161, 195, 230], [164, 197, 231], [168, 200, 231], [171, 202, 232], [174, 204, 233], [177, 206, 234], [180, 208, 235], [183, 210, 236], [187, 212, 237], [190, 214, 238], [193, 216, 239], [196, 218, 239], [199, 220, 240], [202, 222, 241], [205, 223, 242], [208, 225, 242], [211, 227, 243], [214, 229, 244], [217, 231, 245], [219, 232, 245], [222, 234, 246], [225, 236, 247], [227, 237, 247], [230, 239, 248], [232, 240, 248], [234, 242, 249], [237, 243, 250], [239, 245, 250], [241, 246, 251], [243, 247, 251], [245, 249, 252], [247, 250, 252], [249, 251, 253], [251, 252, 253], [252, 253, 254], [254, 254, 254]]
colormap_wave3 = [[253, 203, 160], [252, 201, 158], [252, 200, 156], [252, 198, 154], [251, 196, 152], [251, 194, 149], [251, 192, 147], [250, 190, 145], [250, 189, 142], [249, 187, 139], [249, 185, 135], [248, 182, 132], [248, 179, 129], [247, 177, 125], [247, 175, 122], [247, 172, 119], [246, 168, 115], [246, 166, 112], [245, 164, 108], [245, 161, 105], [244, 158, 102], [243, 155, 98], [243, 152, 95], [242, 150, 92], [242, 147, 88], [241, 145, 86], [240, 142, 83], [240, 139, 80], [239, 137, 78], [238, 134, 75], [237, 131, 73], [237, 129, 70], [236, 126, 68], [235, 124, 67], [234, 121, 65], [233, 118, 63], [232, 115, 61], [231, 112, 61], [230, 110, 59], [229, 107, 57], [228, 104, 56], [228, 102, 55], [226, 100, 55], [224, 97, 55], [224, 94, 54], [222, 92, 54], [221, 89, 53], [219, 87, 53], [218, 84, 53], [217, 83, 53], [215, 80, 54], [213, 78, 54], [211, 76, 54], [209, 74, 55], [208, 72, 55], [206, 70, 56], [203, 68, 57], [201, 66, 57], [199, 65, 58], [198, 64, 59], [196, 62, 59], [193, 60, 60], [191, 59, 61], [188, 57, 61], [186, 56, 62], [183, 55, 62], [181, 55, 63], [179, 54, 63], [176, 53, 63], [173, 52, 64], [171, 51, 64], [168, 50, 65], [165, 49, 65], [162, 49, 65], [159, 48, 66], [157, 48, 66], [155, 47, 66], [152, 47, 66], [149, 46, 66], [146, 46, 65], [143, 46, 65], [140, 45, 65], [138, 45, 65], [135, 45, 64], [132, 44, 64], [130, 44, 63], [127, 44, 63], [124, 44, 62], [121, 43, 61], [118, 43, 60], [115, 43, 60], [112, 42, 60], [110, 42, 59], [108, 42, 58], [105, 42, 57], [102, 41, 56], [99, 41, 55], [97, 40, 54], [94, 40, 53], [91, 40, 52], [89, 39, 51], [86, 39, 50], [84, 38, 49], [82, 38, 48], [79, 37, 47], [77, 37, 46], [74, 36, 45], [72, 36, 44], [70, 35, 43], [68, 35, 42], [65, 34, 41], [64, 34, 40], [62, 33, 39], [60, 33, 38], [58, 33, 38], [56, 32, 38], [54, 32, 36], [52, 32, 35], [50, 32, 35], [48, 31, 35], [47, 31, 34], [45, 31, 34], [43, 31, 33], [42, 30, 33], [41, 30, 32], [40, 30, 32], [39, 30, 32], [38, 30, 32], [37, 30, 32], [36, 30, 32], [36, 30, 32], [37, 30, 32], [38, 30, 32], [39, 30, 32], [40, 30, 32], [41, 30, 32], [42, 30, 33], [44, 31, 33], [46, 31, 34], [47, 31, 34], [49, 31, 35], [51, 32, 35], [53, 32, 36], [55, 32, 37], [57, 33, 38], [59, 33, 38], [61, 33, 39], [63, 34, 40], [65, 34, 41], [67, 35, 42], [70, 35, 43], [72, 36, 44], [74, 36, 45], [77, 37, 46], [79, 37, 47], [82, 38, 48], [84, 38, 49], [87, 39, 50], [90, 39, 51], [92, 40, 52], [95, 40, 53], [98, 40, 54], [100, 41, 55], [103, 41, 56], [106, 42, 57], [109, 42, 58], [111, 42, 59], [114, 43, 60], [117, 43, 60], [120, 43, 61], [123, 44, 62], [126, 44, 63], [129, 44, 63], [131, 44, 64], [134, 45, 64], [137, 45, 65], [140, 45, 65], [143, 46, 65], [146, 46, 65], [149, 46, 66], [152, 47, 66], [155, 47, 66], [158, 48, 66], [160, 48, 66], [163, 49, 65], [166, 49, 65], [169, 50, 65], [172, 51, 64], [174, 52, 64], [177, 53, 63], [180, 54, 63], [182, 55, 62], [185, 56, 62], [187, 57, 61], [190, 58, 61], [192, 60, 60], [195, 61, 59], [197, 63, 59], [199, 65, 58], [201, 66, 57], [203, 68, 57], [206, 70, 56], [208, 72, 55], [209, 74, 55], [211, 76, 54], [213, 78, 54], [215, 81, 54], [217, 83, 53], [218, 85, 53], [220, 88, 53], [221, 90, 53], [223, 93, 54], [224, 95, 54], [225, 98, 55], [227, 101, 55], [228, 103, 56], [229, 106, 57], [230, 109, 58], [231, 111, 60], [232, 114, 61], [233, 117, 62], [234, 120, 64], [235, 123, 66], [236, 125, 68], [237, 128, 70], [237, 131, 73], [238, 134, 75], [239, 137, 78], [240, 139, 80], [240, 142, 83], [241, 145, 86], [242, 148, 89], [242, 151, 93], [243, 153, 96], [243, 156, 99], [244, 159, 103], [245, 162, 106], [245, 165, 110], [246, 167, 113], [246, 170, 117], [247, 173, 120], [247, 176, 124], [248, 178, 127], [248, 181, 131], [249, 184, 134], [249, 186, 138], [250, 188, 141], [250, 190, 144], [251, 192, 147], [251, 194, 149], [251, 196, 152], [252, 198, 154], [252, 200, 156], [252, 201, 158], [253, 203, 160]]
colormap_wave4 = [[246, 230, 183], [246, 229, 182], [246, 227, 180], [246, 226, 178], [246, 224, 176], [245, 222, 173], [245, 219, 170], [244, 217, 167], [244, 214, 163], [244, 211, 160], [243, 209, 156], [243, 206, 152], [242, 203, 148], [242, 200, 144], [241, 196, 140], [241, 193, 136], [241, 190, 132], [240, 186, 128], [240, 183, 124], [239, 180, 120], [239, 176, 116], [238, 173, 112], [238, 170, 108], [237, 166, 104], [237, 163, 100], [236, 160, 97], [236, 156, 93], [236, 153, 90], [235, 150, 87], [235, 147, 84], [235, 144, 81], [234, 140, 78], [234, 137, 76], [234, 134, 74], [234, 131, 71], [233, 127, 69], [233, 124, 67], [233, 121, 65], [233, 118, 64], [232, 115, 62], [232, 112, 61], [232, 109, 60], [232, 106, 59], [232, 103, 58], [232, 101, 58], [232, 98, 57], [232, 95, 57], [231, 93, 57], [230, 90, 57], [230, 88, 57], [229, 85, 57], [227, 83, 57], [226, 81, 57], [224, 78, 57], [222, 76, 58], [220, 74, 58], [217, 72, 59], [215, 70, 59], [212, 67, 60], [210, 65, 60], [207, 63, 60], [204, 62, 61], [201, 60, 61], [199, 58, 61], [196, 56, 61], [193, 55, 61], [189, 53, 61], [186, 52, 62], [183, 50, 61], [180, 49, 61], [176, 48, 61], [173, 46, 61], [170, 45, 61], [166, 44, 61], [163, 43, 61], [159, 42, 60], [156, 40, 60], [152, 39, 60], [149, 39, 59], [146, 38, 58], [142, 37, 58], [139, 36, 57], [135, 35, 56], [132, 34, 55], [128, 34, 54], [125, 33, 53], [122, 32, 52], [118, 31, 51], [115, 31, 50], [111, 30, 49], [108, 29, 48], [105, 28, 46], [101, 28, 45], [98, 27, 44], [94, 26, 43], [91, 25, 41], [88, 24, 40], [85, 24, 38], [82, 23, 37], [78, 22, 35], [75, 21, 34], [72, 20, 32], [69, 19, 31], [66, 18, 29], [63, 18, 28], [60, 16, 26], [57, 16, 25], [55, 15, 24], [52, 14, 22], [49, 13, 21], [46, 12, 19], [44, 11, 18], [41, 11, 17], [39, 10, 16], [36, 9, 14], [34, 8, 13], [31, 8, 12], [29, 7, 11], [27, 7, 10], [25, 6, 10], [23, 5, 9], [21, 5, 8], [19, 4, 7], [17, 4, 7], [16, 4, 6], [14, 3, 6], [13, 3, 5], [13, 3, 5], [12, 3, 5], [12, 3, 5], [12, 3, 5], [12, 3, 5], [13, 3, 5], [14, 3, 5], [15, 3, 6], [16, 4, 6], [18, 4, 7], [20, 4, 7], [21, 5, 8], [23, 5, 9], [26, 6, 10], [28, 7, 11], [30, 7, 12], [33, 8, 13], [35, 9, 14], [38, 9, 15], [40, 10, 17], [43, 11, 18], [46, 12, 19], [48, 13, 21], [51, 14, 22], [54, 15, 24], [57, 16, 25], [60, 17, 27], [63, 18, 28], [67, 19, 30], [70, 20, 31], [73, 20, 33], [76, 21, 34], [80, 22, 36], [83, 23, 37], [86, 24, 39], [90, 25, 40], [93, 26, 42], [96, 27, 43], [100, 27, 44], [103, 28, 46], [107, 29, 47], [110, 30, 48], [114, 30, 50], [117, 31, 51], [121, 32, 52], [124, 33, 53], [128, 34, 54], [131, 34, 55], [135, 35, 56], [138, 36, 57], [142, 37, 57], [146, 38, 58], [149, 39, 59], [153, 39, 59], [156, 40, 60], [160, 42, 60], [164, 43, 61], [167, 44, 61], [171, 45, 61], [174, 46, 61], [178, 48, 62], [181, 49, 62], [184, 51, 62], [188, 52, 62], [191, 54, 62], [194, 56, 61], [197, 57, 61], [200, 59, 61], [203, 61, 61], [206, 63, 60], [209, 65, 60], [212, 67, 59], [215, 69, 59], [217, 72, 59], [220, 74, 58], [222, 76, 58], [224, 78, 57], [226, 81, 57], [227, 83, 57], [229, 86, 57], [230, 88, 57], [230, 91, 57], [231, 94, 57], [232, 97, 57], [232, 99, 57], [232, 102, 58], [232, 105, 59], [232, 108, 60], [232, 111, 61], [232, 114, 62], [232, 117, 63], [233, 120, 65], [233, 123, 66], [233, 127, 68], [233, 130, 71], [234, 133, 73], [234, 137, 75], [234, 140, 78], [235, 143, 81], [235, 147, 84], [235, 150, 87], [236, 153, 90], [236, 157, 94], [236, 160, 97], [237, 164, 101], [237, 167, 105], [238, 170, 109], [238, 174, 113], [239, 178, 117], [239, 181, 121], [240, 184, 126], [240, 188, 130], [241, 192, 134], [241, 195, 138], [242, 198, 142], [242, 202, 147], [243, 205, 151], [243, 208, 155], [244, 211, 159], [244, 214, 162], [244, 216, 166], [245, 219, 169], [245, 221, 173], [246, 224, 175], [246, 226, 178], [246, 227, 180], [246, 229, 182], [246, 230, 183]]


class WaveVisualizer:
    def __init__(self, field_colormap, intensity_colormap):
        self.field_colormap = field_colormap
        self.intensity_colormap = intensity_colormap
        self.intensity = None
        self.intensity_exp_average_factor = 0.98
        self.field = None
        self.visualization_image = None

    def update(self, wave_sim):
        self.field = wave_sim.get_field()

        if self.intensity is None:
            self.intensity = cp.zeros_like(self.field)

        t = self.intensity_exp_average_factor
        self.intensity = self.intensity*t + (self.field**2)*(1.0-t)
        self.visualization_image = wave_sim.render_visualization()

    def render_intensity(self, brightness_scale=1.0, exp=0.5, overlay_visualization=True):
        gray = (cp.clip((self.intensity**exp)*brightness_scale, 0.0, 1.0) * 254.0).astype(np.uint8)
        img = self.intensity_colormap[gray].get() if self.intensity_colormap is not None else gray.get()
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        if overlay_visualization:
            img = cv2.add(img, self.visualization_image)
        return img

    def render_field(self, brightness_scale=1.0, overlay_visualization=True):
        gray = (cp.clip(self.field*brightness_scale, -1.0, 1.0) * 127 + 127).astype(np.uint8)
        img = self.field_colormap[gray].get() if self.field_colormap is not None else gray.get()
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        if overlay_visualization:
            img = cv2.add(img, self.visualization_image)
        return img


def get_colormap_lut(name, invert, black_level=0.0, make_symmetric=False):
    if name == 'icefire': color_values = np.array(colormap_icefire)/255
    elif name == 'colormap_wave1': color_values = np.array(colormap_wave1)/255
    elif name == 'colormap_wave2':color_values = np.array(colormap_wave2) / 255
    elif name == 'colormap_wave3': color_values = np.array(colormap_wave3) / 255
    elif name == 'colormap_wave4': color_values = np.array(colormap_wave4) / 255
    else:
        colormap = matplotlib.pyplot.get_cmap(name)
        color_values = colormap(np.linspace(0, 1, 255))

    if invert:
        color_values = 1.0-color_values

    if make_symmetric:
        src = color_values.copy()
        color_values[255:126:-1, :] = src[0:255:2, :]
        color_values[0:128, :] = src[0:255:2, :]

    color_values = np.clip(color_values*(1.0-black_level)+black_level, 0, 255)

    return cp.asarray((color_values*255).astype(np.uint8))