Source code for simvx.core.gpu_particles

"""GPU-accelerated particle emitter nodes (2D and 3D).

Scene-node wrappers for GPU compute-shader particle simulation. Particle
state lives entirely on the GPU -- positions, velocities, colours, and
lifetimes are updated each frame by a compute dispatch, avoiding per-frame
CPU-to-GPU uploads.

Both ``GPUParticles2D`` and ``GPUParticles3D`` expose ``Property`` descriptors
so the editor inspector can tweak parameters at design time. At runtime the
graphics backend reads ``emitter_config`` each frame and feeds it into the
compute shader push constants.
"""

import logging
import math

import numpy as np

from .colour import srgb_to_linear_rgb
from .descriptors import Property
from .nodes_2d.node2d import Node2D
from .nodes_3d.node3d import Node3D
from .properties import Colour
from .signals import Signal

log = logging.getLogger(__name__)

__all__ = ["GPUParticles2D", "GPUParticles3D"]

# ============================================================================
# Shared mixin -- common Property declarations and helpers
# ============================================================================


class _GPUParticlesBase:
    """Mixin providing shared GPU particle properties and helpers.

    Not instantiated directly -- mixed into ``GPUParticles2D`` and
    ``GPUParticles3D`` alongside their respective spatial base classes.
    """

    amount = Property(1024, range=(1, 1_000_000), hint="Maximum particle count (GPU buffer size)", group="Emission")
    lifetime = Property(2.0, range=(0.01, 60.0), hint="Particle lifetime in seconds", group="Emission")
    emitting = Property(True, hint="Whether the emitter is actively spawning particles", group="Emission")
    one_shot = Property(False, hint="Stop emitting after one full cycle", group="Emission")
    # speed/speed_variance are magnitudes in the emitter's space -- world units/sec
    # in 3D, screen pixels/sec in 2D -- so there is no single hard ceiling. The range
    # is a soft editor-slider hint (clamp=False): a 2D fountain may legitimately run
    # at several hundred px/sec, well past a 3D-sized cap.
    speed = Property(5.0, range=(0.0, 1000.0), clamp=False, hint="Initial speed magnitude", group="Movement")
    speed_variance = Property(0.0, range=(0.0, 1000.0), clamp=False, hint="Random speed variation", group="Movement")
    direction = Property((0.0, 1.0, 0.0), hint="Emission direction (normalised internally)", group="Movement")
    spread = Property(0.3, range=(0.0, 10.0), hint="Velocity spread (randomisation)", group="Movement")
    # How the launch velocity is sampled within ``spread``:
    #   box      - independent uniform per axis (a square/cube envelope; legacy default)
    #   disc     - uniform circular cone of half-angle ``spread`` (radians, 2D)
    #   gaussian - normal-distributed cone + slight speed variation (a soft circle)
    #   ring     - omnidirectional, uniform on a circle (ignores ``direction``)
    #   star     - omnidirectional, speed follows an N-point star curve (``spread_points``)
    # ring/star are radial bursts (use the launch speed, ignore the aim direction).
    spread_pattern = Property(
        "box",
        enum=["box", "disc", "gaussian", "ring", "star"],
        hint="Velocity sampling shape",
        group="Movement",
    )
    spread_points = Property(5, range=(2, 12), hint="Point count for ring/star patterns", group="Movement")
    gravity = Property((0.0, -9.8, 0.0), hint="Gravity vector applied each frame", group="Movement")
    damping = Property(0.0, range=(0.0, 10.0), hint="Velocity damping per second", group="Movement")
    emission_shape = Property("point", enum=["point", "sphere", "box"], hint="Emission shape", group="Emission")
    emission_radius = Property(1.0, range=(0.0, 100.0), hint="Sphere emission radius", group="Emission")
    emission_box = Property((1.0, 1.0, 1.0), hint="Box emission half-extents", group="Emission")
    start_colour = Colour((1.0, 1.0, 1.0, 1.0), group="Appearance")
    end_colour = Colour((1.0, 1.0, 1.0, 0.0), group="Appearance")
    start_scale = Property(1.0, range=(0.0, 50.0), hint="Particle scale at birth", group="Appearance")
    end_scale = Property(0.0, range=(0.0, 50.0), hint="Particle scale at death", group="Appearance")
    # Velocity-stretched (streak) billboards. 0.0 (default) keeps the round
    # camera-facing sprite every emitter has always drawn; a positive value
    # stretches each sprite along its screen-space velocity by that many
    # seconds of travel, so a fast-moving grain reads as a motion streak rather
    # than a round blob (used by Rain3D to draw rain as streaks, not snow).
    streak = Property(
        0.0,
        range=(0.0, 1.0),
        clamp=False,
        hint="Velocity stretch: 0 = round sprite, >0 = streak length in seconds of motion",
        group="Appearance",
    )
    # Camera-proximity shrink. 0.0 (default) draws every sprite at its true
    # perspective size. A positive value is a distance in metres within which a
    # sprite shrinks toward the camera, capping its on-screen size so grains
    # passing close to the lens do not balloon into big blobs (used by Rain3D so
    # near-camera drops stay slim). Beyond the distance sprites are untouched.
    near_fade = Property(
        0.0,
        range=(0.0, 50.0),
        clamp=False,
        hint="Camera-proximity shrink distance in metres: 0 = off, >0 caps near-camera sprite size",
        group="Appearance",
    )
    explosiveness = Property(0.0, range=(0.0, 1.0), hint="0 = steady stream, 1 = all at once", group="Emission")
    randomness = Property(0.0, range=(0.0, 1.0), hint="Lifetime randomness factor", group="Emission")
    fixed_fps = Property(0, range=(0, 120), hint="Lock simulation to N FPS (0 = unlocked)")
    preprocess = Property(0.0, range=(0.0, 10.0), hint="Seconds to pre-simulate on ready")
    local_coords = Property(False, hint="Simulate in local space (True) or world space (False)")

    finished = Signal()  # Emitted when one_shot completes

    _emitter_dims = 3  # 3D by default; GPUParticles2D overrides to 2

    def _gpu_particles_init(self):
        """Initialise internal state (called from ``__init__``)."""
        self._gpu_ready = False
        self._elapsed = 0.0
        self._cycle_complete = False

    @property
    def emitter_config(self) -> dict:
        """Build the config dict consumed by ``ParticleCompute.dispatch()``.

        The graphics backend reads this every frame. Keys match the push
        constant layout defined in ``particle_sim.comp``.
        """
        # Resolve emitter position from the spatial node
        pos = tuple(float(v) for v in self.world_position)  # type: ignore[attr-defined]
        # Pad 2D positions to 3D. ``dims`` (2 for GPUParticles2D) tells the sampler
        # to keep bursts in the XY plane -- a z-component would push grains off-screen.
        dims = self._emitter_dims
        if len(pos) == 2:
            pos = (pos[0], pos[1], 0.0)

        _PATTERNS = ("box", "disc", "gaussian", "ring", "star")
        spread_pattern = _PATTERNS.index(self.spread_pattern)

        # Build initial_velocity from direction + speed
        d = tuple(self.direction)
        dx = float(d[0])
        dy = float(d[1])
        dz = float(d[2]) if len(d) > 2 else 0.0
        mag = (dx * dx + dy * dy + dz * dz) ** 0.5
        if mag > 1e-6:
            dx, dy, dz = dx / mag, dy / mag, dz / mag
        spd = float(self.speed)
        initial_velocity = (dx * spd, dy * spd, dz * spd)

        return {
            "emitter_pos": pos,
            "gravity": tuple(float(v) for v in self.gravity),
            "damping": float(self.damping),
            "initial_velocity": initial_velocity,
            "velocity_spread": float(self.spread),
            # This dict is GPU-shaped by declaration: its keys are the
            # particle_sim.comp push constants, so it is the sink for both
            # backends. Alpha is coverage and passes through.
            "start_colour": srgb_to_linear_rgb(tuple(float(v) for v in self.start_colour)),
            "end_colour": srgb_to_linear_rgb(tuple(float(v) for v in self.end_colour)),
            "start_scale": float(self.start_scale),
            "end_scale": float(self.end_scale),
            # Render-stage only: the compute sim ignores it, the billboard
            # vertex shader stretches the quad along velocity when > 0.
            "streak": float(self.streak),
            # Render-stage only: the billboard vertex shader shrinks the quad
            # toward the camera within this distance so near sprites do not balloon.
            "near_fade": float(self.near_fade),
            "emission_radius": float(self.emission_radius),
            "max_particles": int(self.amount),
            "spread_pattern": spread_pattern,
            "spread_points": int(self.spread_points),
            "dims": dims,
            "emitting": bool(self.emitting),
            # Lifetime drives both the age clock and the initial-phase stagger:
            # the shader seeds each grain partway through a randomised life so a
            # continuous emitter fills with a full spread of ages at once (a
            # stream), while explosiveness=1 spawns them together (a burst).
            "lifetime": float(self.lifetime),
            "randomness": float(self.randomness),
            "explosiveness": float(self.explosiveness),
            # Per-grain launch-speed jitter (added to the aim magnitude before the
            # spread pattern is sampled) so a fountain shows a spread of arc heights.
            "speed_variance": float(self.speed_variance),
        }

    def restart(self):
        """Reset the emitter cycle (re-arms one_shot, resets elapsed time)."""
        self._elapsed = 0.0
        self._cycle_complete = False
        self.emitting = True

    def _gpu_process(self, dt: float):
        """Per-frame bookkeeping (one_shot tracking, fixed_fps accumulation)."""
        if not self.emitting:
            return
        self._elapsed += dt
        if self.one_shot and self._elapsed >= float(self.lifetime):
            if not self._cycle_complete:
                self._cycle_complete = True
                self.emitting = False
                self.finished.emit()


# ============================================================================
# GPUParticles2D
# ============================================================================


def _merge_properties(cls):
    """Merge __properties__ from all bases in MRO (handles diamond/mixin inheritance)."""
    merged = {}
    for base in reversed(cls.__mro__):
        if "__properties__" in base.__dict__:
            merged.update(base.__properties__)
    cls.__properties__ = merged
    return cls


[docs] @_merge_properties class GPUParticles2D(_GPUParticlesBase, Node2D): """GPU-accelerated 2D particle emitter. Particle simulation runs entirely on the GPU via a compute shader. The node exposes editor-visible ``Property`` descriptors for all emitter parameters (amount, lifetime, speed, colours, etc.). The graphics backend collects ``GPUParticles2D`` nodes during scene traversal, reads ``emitter_config``, and dispatches the compute shader each frame. Example:: particles = GPUParticles2D( amount=2048, lifetime=1.5, speed=8.0, start_colour=(1.0, 0.8, 0.2, 1.0), end_colour=(1.0, 0.0, 0.0, 0.0), ) scene.add_child(particles) """ _emitter_dims = 2 # keeps spread patterns in the XY plane (see emitter_config)
[docs] @property def emitter_config(self) -> dict: """The shared config, carried through the node's transform. ``direction``, the emission radius and the particle size are given in the emitter's own space, so the launch velocity turns with its world rotation and the radius and size grow with its world scale (to the larger axis: grains are round). """ config = super().emitter_config vx, vy, vz = config["initial_velocity"] c, s = math.cos(self.world_rotation), math.sin(self.world_rotation) config["initial_velocity"] = (vx * c - vy * s, vx * s + vy * c, vz) sx, sy = self.world_scale k = max(abs(float(sx)), abs(float(sy))) config["emission_radius"] *= k config["start_scale"] *= k config["end_scale"] *= k return config
def __init__(self, **kwargs): super().__init__(**kwargs) self._gpu_particles_init() @property def world_position(self): """Return 3D position for the compute shader (z=0 for 2D).""" pos2d = super().world_position return np.array([float(pos2d[0]), float(pos2d[1]), 0.0], dtype=np.float32)
[docs] @world_position.setter def world_position(self, v): """Move the emitter so its WORLD position is ``v``, as on any other ``Node2D``. The getter above widens the position to three components for the compute shader, so the setter narrows again and accepts either width: reading the value and writing it back unchanged is a no-op. Any third component is dropped, because a ``GPUParticles2D`` has no z to place itself at. """ Node2D.world_position.fset(self, (float(v[0]), float(v[1]))) # type: ignore[attr-defined]
[docs] def on_update(self, dt: float): self._gpu_process(dt)
# ============================================================================ # GPUParticles3D # ============================================================================
[docs] @_merge_properties class GPUParticles3D(_GPUParticlesBase, Node3D): """GPU-accelerated 3D particle emitter. Particle simulation runs entirely on the GPU via a compute shader. The node exposes editor-visible ``Property`` descriptors for all emitter parameters (amount, lifetime, speed, colours, etc.). The graphics backend collects ``GPUParticles3D`` nodes during scene traversal, reads ``emitter_config``, and dispatches the compute shader each frame. Example:: particles = GPUParticles3D( amount=4096, lifetime=2.0, speed=10.0, gravity=(0.0, -9.8, 0.0), start_colour=(0.2, 0.6, 1.0, 1.0), end_colour=(0.0, 0.2, 0.8, 0.0), ) scene.add_child(particles) """ def __init__(self, **kwargs): super().__init__(**kwargs) self._gpu_particles_init()
[docs] def on_update(self, dt: float): self._gpu_process(dt)