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.
"""


from __future__ import annotations

import logging

import numpy as np

from .descriptors import Property, Signal
from .nodes_2d.node2d import Node2D
from .nodes_3d.node3d import Node3D

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 = Property(5.0, range=(0.0, 200.0), hint="Initial speed magnitude", group="Movement")
    speed_variance = Property(0.0, range=(0.0, 200.0), 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")
    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 = Property((1.0, 1.0, 1.0, 1.0), hint="Particle colour at birth (RGBA)", group="Appearance")
    end_colour = Property((1.0, 1.0, 1.0, 0.0), hint="Particle colour at death (RGBA)", 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")
    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

    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
        if len(pos) == 2:
            pos = (pos[0], pos[1], 0.0)

        # 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),
            "start_colour": tuple(float(v) for v in self.start_colour),
            "end_colour": tuple(float(v) for v in self.end_colour),
            "start_scale": float(self.start_scale),
            "end_scale": float(self.end_scale),
            "emission_radius": float(self.emission_radius),
            "max_particles": int(self.amount),
        }

    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) """ 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] def process(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 process(self, dt: float): self._gpu_process(dt)