Add ability to generate python raytracing renderering

crates/nexus_python3d/Cargo.toml

@@ -36,6 +36,7 @@ extension-module = ["pyo3/extension-module"]
 
 [dependencies]
 pyo3 = { version = "0.29", features = ["abi3-py39"] }
+numpy = "0.29"
 nexus3d = { workspace = true, features = ["rbd"] }
 nexus_viewer3d = { workspace = true }
 rapier3d = { workspace = true, features = ["default"] }

crates/nexus_python3d/src/viewer.rs

@@ -10,6 +10,7 @@ use crate::nexus::{GpuTimestamps, NexusState};
 use crate::rbd::{RigidBodyHandle, SharedShape};
 use khal::backend::GpuBackend;
 use nexus_viewer3d::NexusViewer as RViewer;
+use numpy::{IntoPyArray, PyArray3, PyArrayMethods};
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 
@@ -57,10 +58,17 @@ impl NexusViewer {
 impl NexusViewer {
     /// Opens a window and probes the GPU. Blocks on the async setup.
     ///
-    /// Must be called on the main thread (required by the OS windowing system).
+    /// `width`/`height` set the window and render-target resolution
+    /// (default 1200x900). Must be called on the main thread (required by the
+    /// OS windowing system).
     #[new]
-    fn new() -> Self {
-        NexusViewer(Some(pollster::block_on(RViewer::new(Vec::new()))))
+    #[pyo3(signature = (width=1200, height=900))]
+    fn new(width: u32, height: u32) -> Self {
+        NexusViewer(Some(pollster::block_on(RViewer::new_with_size(
+            Vec::new(),
+            width,
+            height,
+        ))))
     }
 
     // --- backend selection (fluent) --------------------------------------
@@ -157,11 +165,59 @@ impl NexusViewer {
         pollster::block_on(self.inner_mut().render_frame())
     }
 
+    /// Whether `render_frame` draws the built-in egui panel (default `True`).
+    /// Disable for clean frame capture with `render`.
+    fn set_draw_ui(&mut self, enabled: bool) {
+        self.inner_mut().set_draw_ui(enabled);
+    }
+
+    /// Renders one path-traced frame with kiss3d's GPU path tracer instead of
+    /// the rasterizer. Samples accumulate across calls while the scene is
+    /// static; call it several times (or raise `set_raytracer_samples_per_frame`)
+    /// before `render` to converge. Returns `False` when the window is closed.
+    fn raytrace_frame(&mut self) -> bool {
+        pollster::block_on(self.inner_mut().raytrace_frame())
+    }
+
+    /// Number of path-tracing samples accumulated per `raytrace_frame` call.
+    fn set_raytracer_samples_per_frame(&mut self, samples: u32) {
+        self.inner_mut().set_raytracer_samples_per_frame(samples);
+    }
+
+    /// Maximum path-tracing bounce depth.
+    fn set_raytracer_max_bounces(&mut self, bounces: u32) {
+        self.inner_mut().set_raytracer_max_bounces(bounces);
+    }
+
+    /// Enables/disables the path tracer's denoiser.
+    fn set_raytracer_denoise(&mut self, enabled: bool) {
+        self.inner_mut().set_raytracer_denoise(enabled);
+    }
+
+    /// Which intersection backend the path tracer uses: `"hardware"` (RT-core
+    /// ray queries) or `"software"` (portable compute-shader BVH fallback).
+    fn raytracer_backend(&mut self) -> &'static str {
+        self.inner_mut().raytracer_backend_name()
+    }
+
     /// Whether the simulation should advance this frame (honors play/pause/step).
     fn simulating(&mut self) -> bool {
         self.inner_mut().simulating()
     }
 
+    /// Returns the last rendered frame as an `(H, W, 3)` `uint8` NumPy array
+    /// (row-major, top-to-bottom, RGB), like `mujoco.Renderer.render()`.
+    ///
+    /// Call once per frame after [`render_frame`][Self::render_frame] to export
+    /// frames off-screen (e.g. to encode a video) instead of only presenting to
+    /// the window.
+    fn snap_rgb<'py>(&mut self, py: Python<'py>) -> PyResult<Bound<'py, PyArray3<u8>>> {
+        let (w, h, rgb) = self.inner_mut().snap_rgb();
+        rgb.into_pyarray(py)
+            .reshape([h as usize, w as usize, 3])
+            .map_err(|e| PyRuntimeError::new_err(format!("{e:?}")))
+    }
+
     /// Reads GPU state back into the renderer. Call once per frame after
     /// `simulate`. Blocks on the async readback.
     #[pyo3(signature = (state, timestamps=None))]

src_viewer/viewer.rs

@@ -23,6 +23,8 @@ use khal::re_exports::wgpu::{Features, Limits};
 use std::time::Duration;
 
 use kiss3d::prelude::Color;
+#[cfg(feature = "dim3")]
+use kiss3d::renderer::RayTracer;
 use kiss3d::scene::{SceneNode2d, SceneNode3d};
 use kiss3d::window::{NumSamples, Window};
 
@@ -145,6 +147,13 @@ pub struct NexusViewer {
     last_gpu_pass_times: Vec<(String, f64)>,
     /// Total GPU time matching [`Self::last_gpu_pass_times`].
     last_gpu_total_time_ms: f64,
+    /// Whether [`Self::render_frame`] draws the built-in egui panel. Disable
+    /// for clean frame capture ([`Self::snap_rgb`]).
+    draw_ui: bool,
+    /// Lazily-created path tracer for [`Self::raytrace_frame`]. Kept across
+    /// frames so samples keep accumulating while the scene is static.
+    #[cfg(feature = "dim3")]
+    raytracer: Option<RayTracer>,
     pub ui: UiState,
 }
 
@@ -154,13 +163,18 @@ impl NexusViewer {
     /// `demos` is the list of `(name, kind)` shown in the demo picker; pass an
     /// empty vec for a standalone single-example viewer (no picker).
     pub async fn new(demos: Vec<(String, DemoKind)>) -> Self {
+        Self::new_with_size(demos, 1200, 900).await
+    }
+
+    /// Like [`Self::new`] but with an explicit window/render resolution.
+    pub async fn new_with_size(demos: Vec<(String, DemoKind)>, width: u32, height: u32) -> Self {
         // NOTE: PASSTHROUGH_SHADERS is required for compute shaders with spirv_passthrough on
         //       platforms running vulkan (to work around some naga limitations).
         let setup = kiss3d::window::CanvasSetup {
             required_features: Features::PASSTHROUGH_SHADERS,
             ..Default::default()
         };
-        let mut window = Window::new_with_setup("nexus demos", 1200, 900, setup).await;
+        let mut window = Window::new_with_setup("nexus demos", width, height, setup).await;
         window.set_background_color(Color::new(245.0 / 255.0, 245.0 / 255.0, 236.0 / 255.0, 1.0));
         // Disable MSAA, this puts extra load on the GPU that ends up
         // falsifying the gpu physics timestamps.
@@ -213,6 +227,9 @@ impl NexusViewer {
             render_resources_backend: None,
             last_gpu_pass_times: Vec::new(),
             last_gpu_total_time_ms: 0.0,
+            draw_ui: true,
+            #[cfg(feature = "dim3")]
+            raytracer: None,
             ui: UiState {
                 run_state: RunState::Paused,
                 run_stats: RunStats::default(),
@@ -323,10 +340,26 @@ impl NexusViewer {
     /// falls back to the GPU→CPU readback path.
     fn direct_render_path(&self) -> bool {
         self.webgpu_shared
+            && !self.rt_active()
             && self.ui.backend_type == BackendType::Gpu
             && matches!(self.webgpu, Some(KhalGpuBackend::WebGpu(_)))
     }
 
+    /// Whether the path tracer has been used/configured. The tracer rebuilds its
+    /// acceleration structure from the CPU-side instance buffers, so while it is
+    /// active `sync` must take the readback path — the zero-readback kernel only
+    /// updates kiss3d's GPU instance buffers, which the tracer never reads.
+    fn rt_active(&self) -> bool {
+        #[cfg(feature = "dim3")]
+        {
+            self.raytracer.is_some()
+        }
+        #[cfg(not(feature = "dim3"))]
+        {
+            false
+        }
+    }
+
     #[cfg(feature = "cuda")]
     fn init_cuda(&mut self) -> Option<KhalGpuBackend> {
         match khal::backend::cuda::Cuda::new(0) {
@@ -769,14 +802,100 @@ impl NexusViewer {
         let gpu_available = self.webgpu.is_some();
         // Disjoint closure capture (edition 2024): the closure borrows `self.ui`
         // and `scene_ui` while `self.window` is the receiver.
-        self.window.draw_ui(|ctx| {
-            crate::ui::setup_custom_theme(ctx);
-            crate::ui::main_panel(ctx, &mut self.ui, gpu_available);
-        });
+        if self.draw_ui {
+            self.window.draw_ui(|ctx| {
+                crate::ui::setup_custom_theme(ctx);
+                crate::ui::main_panel(ctx, &mut self.ui, gpu_available);
+            });
+        }
 
         self.ui.transition.is_none()
     }
 
+    /// Captures the last rendered frame as `(width, height, rgb)`, where `rgb`
+    /// is row-major, top-to-bottom, 3 bytes (R, G, B) per pixel.
+    ///
+    /// This is the off-screen counterpart of the on-window presentation done by
+    /// [`Self::render_frame`]: call `render_frame` to draw the scene, then this
+    /// to read the framebuffer back to the CPU (e.g. to export a video).
+    pub fn snap_rgb(&mut self) -> (u32, u32, Vec<u8>) {
+        let img = self.window.snap_image();
+        (img.width(), img.height(), img.into_raw())
+    }
+
+    /// Whether [`Self::render_frame`] draws the built-in egui panel (default
+    /// `true`). Disable for clean frame capture.
+    pub fn set_draw_ui(&mut self, enabled: bool) {
+        self.draw_ui = enabled;
+    }
+
+    /// Renders one path-traced frame of the viewer-owned scene with kiss3d's
+    /// GPU path tracer, instead of the rasterizer used by
+    /// [`Self::render_frame`].
+    ///
+    /// The tracer accumulates samples across calls while the scene and camera
+    /// are static, and restarts accumulation automatically when either changes.
+    /// Call it several times per video frame (or raise
+    /// [`Self::set_raytracer_samples_per_frame`]) to converge before grabbing
+    /// the image with [`Self::snap_rgb`]. Returns `false` when the loop
+    /// should end (window closed).
+    #[cfg(feature = "dim3")]
+    pub async fn raytrace_frame(&mut self) -> bool {
+        let raytracer = self.raytracer.get_or_insert_with(RayTracer::new);
+        let cont = self
+            .window
+            .raytrace_3d(&mut self.scene3d, &mut self.camera3d, raytracer)
+            .await;
+        if !cont {
+            self.ui.transition = Some(Transition::Quit);
+        }
+        cont
+    }
+
+    /// Number of path-tracing samples accumulated per [`Self::raytrace_frame`]
+    /// call (default: chosen by kiss3d from the GPU tier).
+    #[cfg(feature = "dim3")]
+    pub fn set_raytracer_samples_per_frame(&mut self, samples: u32) {
+        self.raytracer
+            .get_or_insert_with(RayTracer::new)
+            .set_samples_per_frame(samples);
+    }
+
+    /// Maximum path-tracing bounce depth.
+    #[cfg(feature = "dim3")]
+    pub fn set_raytracer_max_bounces(&mut self, bounces: u32) {
+        self.raytracer
+            .get_or_insert_with(RayTracer::new)
+            .set_max_bounces(bounces);
+    }
+
+    /// Which intersection backend the path tracer uses: hardware ray queries
+    /// (RT cores, when the device supports them) or the portable
+    /// compute-shader BVH fallback.
+    #[cfg(feature = "dim3")]
+    pub fn raytracer_backend(&mut self) -> kiss3d::renderer::RayBackend {
+        self.raytracer.get_or_insert_with(RayTracer::new).backend()
+    }
+
+    /// [`Self::raytracer_backend`] as a string (`"hardware"` / `"software"`),
+    /// for consumers that don't depend on kiss3d directly (e.g. the Python
+    /// bindings).
+    #[cfg(feature = "dim3")]
+    pub fn raytracer_backend_name(&mut self) -> &'static str {
+        match self.raytracer_backend() {
+            kiss3d::renderer::RayBackend::Hardware => "hardware",
+            kiss3d::renderer::RayBackend::Software => "software",
+        }
+    }
+
+    /// Enables/disables the path tracer's denoiser.
+    #[cfg(feature = "dim3")]
+    pub fn set_raytracer_denoise(&mut self, enabled: bool) {
+        self.raytracer
+            .get_or_insert_with(RayTracer::new)
+            .set_denoise(enabled);
+    }
+
     /// Draws example-specific egui widgets into the current frame's UI pass.
     ///
     /// Call this once per frame, after [`Self::render_frame`], to overlay a