test/effects/displacement.py

# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "opencv-python"]
# ///
"""
@effect displacement
@version 1.0.0
@author artdag

@description
Displacement effect. Warps the image based on a pattern (sine waves,
noise, or radial). Creates flowing, liquid-like distortions.

@param amount float
  @range 0 100
  @default 20
  Displacement strength in pixels.

@param pattern string
  @enum sine noise radial turbulence
  @default sine
  Displacement pattern type:
  - sine: smooth sine wave
  - noise: random displacement
  - radial: swirl from center
  - turbulence: multi-scale noise

@param frequency float
  @range 1 50
  @default 10
  Pattern frequency (waves per frame width).

@param speed float
  @range 0 10
  @default 1
  Animation speed.

@param direction string
  @enum horizontal vertical both
  @default both
  Displacement direction.

@param seed int
  @default 42
  Random seed for noise patterns.

@state rng DeterministicRNG
  Random number generator.

@example
  (effect displacement :amount 30 :pattern "sine" :frequency 5)

@example
  ;; Reactive turbulence
  (effect displacement :amount (bind energy :range [10 50]) :pattern "turbulence")
"""

import numpy as np
import cv2
from pathlib import Path
import sys

# Import DeterministicRNG from same directory
_effects_dir = Path(__file__).parent
if str(_effects_dir) not in sys.path:
    sys.path.insert(0, str(_effects_dir))
from random import DeterministicRNG


def process_frame(frame: np.ndarray, params: dict, state: dict) -> tuple:
    """
    Apply displacement effect to a video frame.

    Args:
        frame: Input frame as numpy array (H, W, 3) RGB uint8
        params: Effect parameters
        state: Persistent state dict

    Returns:
        Tuple of (processed_frame, new_state)
    """
    amount = params.get("amount", 20)
    pattern = params.get("pattern", "sine")
    frequency = max(1, params.get("frequency", 10))
    speed = params.get("speed", 1)
    direction = params.get("direction", "both")
    seed = int(params.get("seed", 42))
    t = params.get("_time", 0)

    if state is None:
        state = {}

    if amount == 0:
        return frame, state

    # Initialize RNG
    if "rng" not in state:
        state["rng"] = DeterministicRNG(seed)

    h, w = frame.shape[:2]

    # Create base coordinate maps
    map_x = np.tile(np.arange(w, dtype=np.float32), (h, 1))
    map_y = np.tile(np.arange(h, dtype=np.float32).reshape(-1, 1), (1, w))

    # Generate displacement based on pattern
    if pattern == "sine":
        # Sine wave displacement
        phase = t * speed * 2 * np.pi
        if direction in ["horizontal", "both"]:
            map_x = map_x + amount * np.sin(2 * np.pi * map_y / h * frequency + phase)
        if direction in ["vertical", "both"]:
            map_y = map_y + amount * np.sin(2 * np.pi * map_x / w * frequency + phase)

    elif pattern == "noise":
        # Generate noise displacement
        rng = state["rng"]
        if "noise_x" not in state or state.get("noise_size") != (h, w):
            state["noise_x"] = np.array([[rng.uniform(-1, 1) for _ in range(w)] for _ in range(h)], dtype=np.float32)
            state["noise_y"] = np.array([[rng.uniform(-1, 1) for _ in range(w)] for _ in range(h)], dtype=np.float32)
            state["noise_size"] = (h, w)

        if direction in ["horizontal", "both"]:
            map_x = map_x + amount * state["noise_x"]
        if direction in ["vertical", "both"]:
            map_y = map_y + amount * state["noise_y"]

    elif pattern == "radial":
        # Radial/swirl displacement
        cx, cy = w / 2, h / 2
        y_coords, x_coords = np.mgrid[0:h, 0:w].astype(np.float32)
        dx = x_coords - cx
        dy = y_coords - cy
        dist = np.sqrt(dx**2 + dy**2) + 1e-6
        angle = np.arctan2(dy, dx)

        # Swirl amount varies with distance and time
        swirl = amount * 0.01 * np.sin(dist / (w / frequency) + t * speed * 2 * np.pi)

        new_angle = angle + swirl
        if direction in ["horizontal", "both"]:
            map_x = cx + dist * np.cos(new_angle)
        if direction in ["vertical", "both"]:
            map_y = cy + dist * np.sin(new_angle)

    elif pattern == "turbulence":
        # Multi-scale noise
        rng = state["rng"]
        disp_x = np.zeros((h, w), dtype=np.float32)
        disp_y = np.zeros((h, w), dtype=np.float32)

        for scale in [1, 2, 4]:
            sh, sw = h // scale, w // scale
            noise_x = np.array([[rng.uniform(-1, 1) for _ in range(sw)] for _ in range(sh)], dtype=np.float32)
            noise_y = np.array([[rng.uniform(-1, 1) for _ in range(sw)] for _ in range(sh)], dtype=np.float32)
            if scale > 1:
                noise_x = cv2.resize(noise_x, (w, h))
                noise_y = cv2.resize(noise_y, (w, h))
            disp_x += noise_x / scale
            disp_y += noise_y / scale

        if direction in ["horizontal", "both"]:
            map_x = map_x + amount * disp_x
        if direction in ["vertical", "both"]:
            map_y = map_y + amount * disp_y

    # Apply remapping
    result = cv2.remap(frame, map_x, map_y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)

    return result, state