hyperdrone 0.1.0__py3-none-any.whl

hyperdrone/_vendor/rl-tools/src/rendering/raytracing/yaw_prediction/yaw_prediction_viewer.cu

@@ -0,0 +1,445 @@
+#define RL_TOOLS_OPERATIONS_CPU_MUX_INCLUDE_CUDA
+#include <rl_tools/operations/cpu_mux.h>
+
+// Layer operations
+#include <rl_tools/nn/layers/dense/operations_generic.h>
+#include <rl_tools/nn/layers/conv2d/operations_generic.h>
+#include <rl_tools/nn/layers/avg_pool2d/operations_generic.h>
+#include <rl_tools/nn/layers/dynamic_conv2d/operations_generic.h>
+#include <rl_tools/nn/layers/dynamic_conv2d/operations_cuda.h>
+
+// Model operations
+#include <rl_tools/nn_models/sequential/operations_generic.h>
+
+// CUDA-specific operations
+#include <rl_tools/nn/operations_cuda.h>
+#include <rl_tools/nn_models/operations_generic.h>
+
+// Tensor operations
+#include <rl_tools/containers/tensor/operations_generic.h>
+#include <rl_tools/containers/tensor/operations_cpu.h>
+
+// Persist (HDF5)
+#include <rl_tools/persist/backends/hdf5/operations_cpu.h>
+#include <rl_tools/nn/layers/dense/persist.h>
+#include <rl_tools/nn/layers/conv2d/persist.h>
+#include <rl_tools/nn/layers/avg_pool2d/persist.h>
+#include <rl_tools/nn/layers/dynamic_conv2d/persist.h>
+#include <rl_tools/nn_models/sequential/persist.h>
+#include <rl_tools/nn_models/parallel/persist.h>
+
+// Model definition
+#include "model.h"
+#include "model_operations.h"
+#include "model_persist.h"
+#include "model_forward_cuda.h"
+
+// Scene management
+#include "scene.h"
+
+#include <iostream>
+#include <vector>
+#include <cmath>
+#include <cstring>
+#include <sstream>
+#include <iomanip>
+#include <algorithm>
+#include <filesystem>
+#include <cuda_runtime.h>
+#include <cstdio>
+
+namespace rlt = rl_tools;
+namespace fs = std::filesystem;
+
+// ---- Configuration ----
+using T = float;
+using DEVICE_CPU = rlt::devices::DEVICE_FACTORY<>;
+using DEVICE_CUDA = rlt::devices::DEVICE_FACTORY_CUDA<>;
+using TI = DEVICE_CPU::index_t;
+using TI_CUDA = DEVICE_CUDA::index_t;
+
+namespace yp = rlt::rendering::raytracing::yaw_prediction;
+static constexpr TI_CUDA BATCH_SIZE = yp::SCENE_NUM_CAMERAS / 2; // 128
+static constexpr TI_CUDA NUM_CAMERAS = yp::SCENE_NUM_CAMERAS;    // 256
+static constexpr TI CAM_WIDTH = yp::SCENE_CAM_WIDTH;
+static constexpr TI CAM_HEIGHT = yp::SCENE_CAM_HEIGHT;
+
+static constexpr float PI = 3.14159265358979323846f;
+static constexpr float MAX_ANGLE = PI / 6.0f; // 30 degrees
+static constexpr float VIEWER_COS_FOV = 0.66f; // Fixed FOV for viewer
+static constexpr float VIEWER_ASPECT = static_cast<float>(CAM_WIDTH) / static_cast<float>(CAM_HEIGHT);
+static constexpr TI NUM_GRID_CAMERAS = 64;
+static constexpr TI GRID_DIM = 8; // 8x8 grid
+
+using TYPE_POLICY = rlt::numeric_types::Policy<float>;
+
+// Inference-only model
+using GPU_CAPABILITY = rlt::nn::capability::Forward<>;
+using GPU_MODEL = yp::MODEL<GPU_CAPABILITY, TYPE_POLICY, TI_CUDA, BATCH_SIZE, CAM_HEIGHT, CAM_WIDTH>;
+using CPU_CAPABILITY = rlt::nn::capability::Forward<>;
+using CPU_MODEL = yp::MODEL<CPU_CAPABILITY, TYPE_POLICY, TI, BATCH_SIZE, CAM_HEIGHT, CAM_WIDTH>;
+
+using GPU_INPUT_SHAPE = rlt::tensor::Shape<TI_CUDA, BATCH_SIZE, CAM_HEIGHT, CAM_WIDTH, 3>;
+using GPU_INPUT_SPEC = rlt::tensor::Specification<T, TI_CUDA, GPU_INPUT_SHAPE>;
+
+// Head output tensor: [BATCH_SIZE, 3]
+using GPU_OUTPUT_SHAPE = typename GPU_MODEL::OUTPUT_SHAPE;
+using GPU_OUTPUT_SPEC = rlt::tensor::Specification<T, TI_CUDA, GPU_OUTPUT_SHAPE>;
+
+// ---- Progress bead drawing ----
+void draw_bead(std::vector<uint8_t>& image, int image_width, int tile_x, int tile_y,
+               float angle, float max_angle, int bead_y, uint8_t r, uint8_t g, uint8_t b) {
+    float normalized = (angle + max_angle) / (2.0f * max_angle);
+    normalized = std::max(0.0f, std::min(1.0f, normalized));
+    int bead_x = 2 + static_cast<int>(normalized * 59.0f);
+    for (int dy = 0; dy < 3; dy++) {
+        for (int dx = 0; dx < 5; dx++) {
+            int px = tile_x + bead_x + dx;
+            int py = tile_y + bead_y + dy;
+            int idx = (py * image_width + px) * 3;
+            image[idx + 0] = r;
+            image[idx + 1] = g;
+            image[idx + 2] = b;
+        }
+    }
+}
+
+int main(int argc, char** argv) {
+    std::string scene_dir;
+    std::string single_scene;
+    const char* checkpoint_path = nullptr;
+    const char* output_path = "yaw_prediction_viewer.mp4";
+    int num_frames = 120;
+    TI num_scenes = 64;
+
+    for (int i = 1; i < argc; i++) {
+        std::string arg = argv[i];
+        if (arg == "--scene-dir" && i + 1 < argc) {
+            scene_dir = argv[++i];
+        } else if (arg == "--scene" && i + 1 < argc) {
+            single_scene = argv[++i];
+        } else if (arg == "--checkpoint" && i + 1 < argc) {
+            checkpoint_path = argv[++i];
+        } else if (arg == "--output" && i + 1 < argc) {
+            output_path = argv[++i];
+        } else if (arg == "--num-frames" && i + 1 < argc) {
+            num_frames = std::atoi(argv[++i]);
+        } else if (arg == "--num-scenes" && i + 1 < argc) {
+            num_scenes = std::atoi(argv[++i]);
+        } else {
+            std::cerr << "Usage: " << argv[0]
+                      << " [--scene-dir <dir> | --scene <path.glb>]"
+                      << " --checkpoint <path.h5>"
+                      << " [--output <path.mp4>] [--num-frames <N>]"
+                      << " [--num-scenes <N>]" << std::endl;
+            return 1;
+        }
+    }
+
+    if (!checkpoint_path) {
+        std::cerr << "Error: --checkpoint is required" << std::endl;
+        return 1;
+    }
+
+    // Collect scene paths
+    std::vector<std::string> all_scene_paths;
+    if (!scene_dir.empty()) {
+        for (auto& entry : fs::directory_iterator(scene_dir)) {
+            if (entry.path().extension() == ".glb") {
+                all_scene_paths.push_back(entry.path().string());
+            }
+        }
+        std::sort(all_scene_paths.begin(), all_scene_paths.end());
+        if (all_scene_paths.empty()) {
+            std::cerr << "No .glb files found in " << scene_dir << std::endl;
+            return 1;
+        }
+        if (num_scenes < all_scene_paths.size()) {
+            all_scene_paths.resize(num_scenes);
+        }
+    } else if (!single_scene.empty()) {
+        all_scene_paths.push_back(single_scene);
+    } else {
+        std::cerr << "Must specify --scene-dir or --scene" << std::endl;
+        return 1;
+    }
+
+    // Cap to 64 scenes (one per grid tile max)
+    if (all_scene_paths.size() > NUM_GRID_CAMERAS) {
+        all_scene_paths.resize(NUM_GRID_CAMERAS);
+    }
+
+    std::cout << "Yaw prediction viewer (multi-scene)" << std::endl;
+    std::cout << "  Scenes: " << all_scene_paths.size() << std::endl;
+    std::cout << "  Checkpoint: " << checkpoint_path << std::endl;
+    std::cout << "  Output: " << output_path << std::endl;
+    std::cout << "  Num frames: " << num_frames << std::endl;
+
+    // ---- Load scenes ----
+    struct LoadedScene {
+        yp::SceneHandle* handle = nullptr;
+        std::string path;
+    };
+    std::vector<LoadedScene> loaded_scenes(all_scene_paths.size());
+    for (TI s = 0; s < all_scene_paths.size(); s++) {
+        loaded_scenes[s].path = all_scene_paths[s];
+        std::cout << "Loading scene " << s << "/" << all_scene_paths.size() << ": "
+                  << fs::path(loaded_scenes[s].path).filename().string() << std::endl;
+        loaded_scenes[s].handle = yp::create_scene(loaded_scenes[s].path.c_str());
+        std::cout << "  Indoor states: " << yp::get_num_indoor_states(loaded_scenes[s].handle) << std::endl;
+    }
+
+    // Assign each grid tile to a scene (round-robin)
+    std::vector<TI> tile_scene_idx(NUM_GRID_CAMERAS);
+    for (TI i = 0; i < NUM_GRID_CAMERAS; i++) {
+        tile_scene_idx[i] = i % loaded_scenes.size();
+    }
+
+    // Group tiles by scene
+    std::vector<std::vector<TI>> scene_tiles(loaded_scenes.size());
+    for (TI i = 0; i < NUM_GRID_CAMERAS; i++) {
+        scene_tiles[tile_scene_idx[i]].push_back(i);
+    }
+
+    // ---- Device setup ----
+    DEVICE_CPU device_cpu;
+    DEVICE_CUDA device_cuda;
+    rlt::init(device_cpu);
+    rlt::init(device_cuda);
+
+    // ---- Load model ----
+    CPU_MODEL model_cpu;
+    rlt::malloc(device_cpu, model_cpu);
+    {
+        auto file = rl_tools::persist::backends::hdf5::File(checkpoint_path, rl_tools::persist::backends::hdf5::Mode::READ);
+        auto mg = rlt::get_group(device_cpu, file, "model");
+        bool success = rlt::load(device_cpu, model_cpu, mg);
+        if (!success) {
+            std::cerr << "Failed to load checkpoint" << std::endl;
+            return 1;
+        }
+    }
+    std::cout << "Checkpoint loaded" << std::endl;
+
+    GPU_MODEL model;
+    typename GPU_MODEL::template Buffer<true> model_buffer;
+    rlt::malloc(device_cuda, model);
+    rlt::malloc(device_cuda, model_buffer);
+    rlt::copy(device_cpu, device_cuda, model_cpu, model);
+
+    // ---- GPU tensors ----
+    rlt::Tensor<GPU_INPUT_SPEC> gpu_input_a, gpu_input_b;
+    rlt::malloc(device_cuda, gpu_input_a);
+    rlt::malloc(device_cuda, gpu_input_b);
+
+    rlt::Tensor<GPU_OUTPUT_SPEC> gpu_output;
+    rlt::malloc(device_cuda, gpu_output);
+
+    DEVICE_CUDA::SPEC::RANDOM::ENGINE<> rng_cuda;
+    rlt::malloc(device_cuda, rng_cuda);
+    rlt::init(device_cuda, rng_cuda, 42);
+
+    // ---- Sample base cameras per scene ----
+    std::vector<rlt::rendering::raytracing::Camera<float>> base_cameras(NUM_GRID_CAMERAS);
+    {
+        std::vector<rlt::rendering::raytracing::Camera<float>> setup_cameras(NUM_CAMERAS);
+        std::vector<float> setup_targets(BATCH_SIZE * 3);
+
+        for (TI s = 0; s < loaded_scenes.size(); s++) {
+            auto& tiles = scene_tiles[s];
+            if (tiles.empty()) continue;
+
+            yp::sample_camera_batch(loaded_scenes[s].handle, setup_cameras.data(),
+                                    setup_targets.data(),
+                                    BATCH_SIZE, 0.0f,
+                                    VIEWER_COS_FOV, VIEWER_COS_FOV);
+
+            for (TI t = 0; t < tiles.size(); t++) {
+                base_cameras[tiles[t]] = setup_cameras[t % BATCH_SIZE];
+            }
+        }
+    }
+
+    // Rotate a camera's view direction by delta_yaw around the Y axis
+    using YPCameraData = rlt::rendering::raytracing::Camera<float>;
+    auto rotate_camera_yaw = [](const YPCameraData& cam, float delta_yaw) -> YPCameraData {
+        YPCameraData rotated = cam;
+        float cy = std::cos(delta_yaw);
+        float sy = std::sin(delta_yaw);
+        auto rotate_y = [cy, sy](const float v[3], float out[3]) {
+            out[0] = cy * v[0] - sy * v[2];
+            out[1] = v[1];
+            out[2] = sy * v[0] + cy * v[2];
+        };
+        rotate_y(cam.dir_00, rotated.dir_00);
+        rotate_y(cam.dir_du, rotated.dir_du);
+        rotate_y(cam.dir_dv, rotated.dir_dv);
+        return rotated;
+    };
+
+    // ---- Open ffmpeg pipe ----
+    static constexpr TI IMAGE_WIDTH = GRID_DIM * CAM_WIDTH;
+    static constexpr TI IMAGE_HEIGHT = GRID_DIM * CAM_HEIGHT;
+    std::ostringstream ffmpeg_cmd;
+    ffmpeg_cmd
+        << "ffmpeg -y -f rawvideo -pixel_format rgb24 "
+        << "-video_size " << IMAGE_WIDTH << "x" << IMAGE_HEIGHT << " "
+        << "-framerate 30 -i - "
+        << "-an -c:v libx264 -pix_fmt yuv420p "
+        << output_path;
+    FILE* mp4_pipe = popen(ffmpeg_cmd.str().c_str(), "w");
+    if (!mp4_pipe) {
+        std::cerr << "Failed to start ffmpeg" << std::endl;
+        return 1;
+    }
+
+    // ---- Frame generation loop ----
+    static constexpr TI CAM_PIXELS = CAM_WIDTH * CAM_HEIGHT;
+    std::vector<uint32_t> host_fb(NUM_CAMERAS * CAM_PIXELS);
+    std::vector<uint8_t> output_image(IMAGE_WIDTH * IMAGE_HEIGHT * 3);
+    static constexpr TI OUTPUT_DIM = 3;
+    std::vector<float> pred_buf(BATCH_SIZE * OUTPUT_DIM);
+    std::vector<rlt::rendering::raytracing::Camera<float>> cameras(NUM_CAMERAS);
+
+    // Per-tile host framebuffer for swept views (read back per scene)
+    std::vector<std::vector<uint32_t>> tile_pixels(NUM_GRID_CAMERAS, std::vector<uint32_t>(CAM_PIXELS));
+
+    auto eval_mode = rlt::Mode<rlt::mode::Default<>>{};
+
+    for (int frame = 0; frame < num_frames; frame++) {
+        float t = (num_frames > 1) ? static_cast<float>(frame) / static_cast<float>(num_frames - 1) : 0.5f;
+        float delta_yaw = -MAX_ANGLE + t * 2.0f * MAX_ANGLE;
+
+        // ---- Render each scene and fill GPU input tensors ----
+        for (TI s = 0; s < loaded_scenes.size(); s++) {
+            auto& tiles = scene_tiles[s];
+            if (tiles.empty()) continue;
+
+            TI num_pairs = tiles.size();
+
+            for (TI t_idx = 0; t_idx < num_pairs; t_idx++) {
+                cameras[t_idx] = base_cameras[tiles[t_idx]];
+                cameras[num_pairs + t_idx] = rotate_camera_yaw(base_cameras[tiles[t_idx]], delta_yaw);
+            }
+            for (TI i = num_pairs * 2; i < NUM_CAMERAS; i++) {
+                cameras[i] = cameras[i % (num_pairs * 2)];
+            }
+
+            yp::render_batch<false>(loaded_scenes[s].handle, cameras.data());
+            uint32_t* device_fb = yp::get_framebuffer_device_ptr(loaded_scenes[s].handle);
+
+            for (TI t_idx = 0; t_idx < num_pairs; t_idx++) {
+                TI batch_pos = tiles[t_idx];
+                yp::cuda_kernels::rgba_to_activation_kernel<<<1, 256, 0, device_cuda.stream>>>(
+                    device_fb, gpu_input_a._data, CAM_WIDTH, CAM_HEIGHT, t_idx, batch_pos
+                );
+            }
+            for (TI t_idx = 0; t_idx < num_pairs; t_idx++) {
+                TI batch_pos = tiles[t_idx];
+                yp::cuda_kernels::rgba_to_activation_kernel<<<1, 256, 0, device_cuda.stream>>>(
+                    device_fb, gpu_input_b._data, CAM_WIDTH, CAM_HEIGHT, num_pairs + t_idx, batch_pos
+                );
+            }
+
+            cudaStreamSynchronize(device_cuda.stream);
+            for (TI t_idx = 0; t_idx < num_pairs; t_idx++) {
+                TI cam_idx = num_pairs + t_idx;
+                cudaMemcpy(tile_pixels[tiles[t_idx]].data(),
+                           device_fb + cam_idx * CAM_PIXELS,
+                           CAM_PIXELS * sizeof(uint32_t),
+                           cudaMemcpyDeviceToHost);
+            }
+        }
+
+        // Pad batch positions [64..127]
+        for (TI i = NUM_GRID_CAMERAS; i < BATCH_SIZE; i++) {
+            TI src_pos = i % NUM_GRID_CAMERAS;
+            cudaMemcpyAsync(
+                gpu_input_a._data + i * CAM_HEIGHT * CAM_WIDTH * 3,
+                gpu_input_a._data + src_pos * CAM_HEIGHT * CAM_WIDTH * 3,
+                CAM_HEIGHT * CAM_WIDTH * 3 * sizeof(float),
+                cudaMemcpyDeviceToDevice, device_cuda.stream
+            );
+            cudaMemcpyAsync(
+                gpu_input_b._data + i * CAM_HEIGHT * CAM_WIDTH * 3,
+                gpu_input_b._data + src_pos * CAM_HEIGHT * CAM_WIDTH * 3,
+                CAM_HEIGHT * CAM_WIDTH * 3 * sizeof(float),
+                cudaMemcpyDeviceToDevice, device_cuda.stream
+            );
+        }
+
+        // ---- Model forward pass ----
+        rlt::evaluate_cuda(device_cuda, model, gpu_input_a, gpu_input_b, gpu_output, model_buffer, rng_cuda, eval_mode);
+        cudaStreamSynchronize(device_cuda.stream);
+
+        cudaMemcpy(pred_buf.data(), gpu_output._data, BATCH_SIZE * OUTPUT_DIM * sizeof(float), cudaMemcpyDeviceToHost);
+
+        const float half_hfov = std::atan(0.5f * VIEWER_COS_FOV * VIEWER_ASPECT);
+        const float tan_half_hfov = std::tan(half_hfov);
+
+        // ---- Composite 8x8 grid ----
+        std::memset(output_image.data(), 0, output_image.size());
+        for (TI i = 0; i < NUM_GRID_CAMERAS; i++) {
+            TI grid_col = i % GRID_DIM;
+            TI grid_row = i / GRID_DIM;
+            TI tile_x = grid_col * CAM_WIDTH;
+            TI tile_y = grid_row * CAM_HEIGHT;
+
+            const uint32_t* cam_pixels_ptr = tile_pixels[i].data();
+            for (TI y = 0; y < CAM_HEIGHT; y++) {
+                for (TI x = 0; x < CAM_WIDTH; x++) {
+                    uint32_t rgba = cam_pixels_ptr[y * CAM_WIDTH + x];
+                    TI px = tile_x + x;
+                    TI py = tile_y + y;
+                    TI idx = (py * IMAGE_WIDTH + px) * 3;
+                    output_image[idx + 0] = (rgba >>  0) & 0xFF;
+                    output_image[idx + 1] = (rgba >>  8) & 0xFF;
+                    output_image[idx + 2] = (rgba >> 16) & 0xFF;
+                }
+            }
+
+            float pred_angle = std::atan(pred_buf[i * OUTPUT_DIM + 0] * tan_half_hfov);
+
+            draw_bead(output_image, IMAGE_WIDTH, tile_x, tile_y, delta_yaw, MAX_ANGLE, 2, 0, 200, 0);
+            draw_bead(output_image, IMAGE_WIDTH, tile_x, tile_y, pred_angle, MAX_ANGLE, 5, 200, 0, 0);
+        }
+
+        // Write frame to ffmpeg pipe
+        size_t frame_bytes = output_image.size();
+        size_t written = fwrite(output_image.data(), 1, frame_bytes, mp4_pipe);
+        if (written != frame_bytes) {
+            std::cerr << "Failed writing frame " << frame << " to ffmpeg" << std::endl;
+            pclose(mp4_pipe);
+            mp4_pipe = nullptr;
+            break;
+        }
+
+        if (frame % 10 == 0 || frame == num_frames - 1) {
+            std::cout << "Frame " << frame << "/" << num_frames
+                      << "  delta_yaw=" << delta_yaw * 180.0f / PI << " deg" << std::endl;
+        }
+    }
+
+    if (mp4_pipe) {
+        int ffmpeg_status = pclose(mp4_pipe);
+        if (ffmpeg_status != 0) {
+            std::cerr << "ffmpeg exited with status " << ffmpeg_status << std::endl;
+        } else {
+            std::cout << "Video written: " << output_path << std::endl;
+        }
+    }
+
+    // ---- Cleanup ----
+    rlt::free(device_cuda, gpu_output);
+    rlt::free(device_cuda, gpu_input_a);
+    rlt::free(device_cuda, gpu_input_b);
+    rlt::free(device_cuda, model_buffer);
+    rlt::free(device_cuda, model);
+    rlt::free(device_cuda, rng_cuda);
+    rlt::free(device_cpu, model_cpu);
+    for (auto& s : loaded_scenes) {
+        yp::destroy_scene(s.handle);
+    }
+
+    return 0;
+}

hyperdrone/cuda/__init__.py

@@ -0,0 +1,78 @@
+"""hyperdrone.cuda — CUDA staging helpers shared by render and dynamics.
+
+The native CudaBuffer class is compiled into every CUDA-capable component core
+(render on OptiX, dynamics on CUDA). This package defines the Python-side views and
+resolves a provider from whichever domain package registered one at import time, keeping
+the dependency direction render/dynamics -> cuda.
+"""
+import numpy as np
+
+_providers = []
+
+
+def register_provider(load_core):
+    """Called by domain packages (render, dynamics) at import time; load_core is a
+    zero-argument callable returning a core module that may expose CudaBuffer."""
+    if load_core not in _providers:
+        _providers.append(load_core)
+
+
+def _core():
+    errors = []
+    for load_core in _providers:
+        try:
+            core = load_core()
+        except Exception as error:  # provider may fail to build on this machine
+            errors.append(str(error))
+            continue
+        if getattr(core, "HAS_CUDA", False):
+            return core
+    detail = ("; ".join(errors)) if errors else "no CUDA-capable component core is available"
+    raise RuntimeError(
+        "hyperdrone: CUDA staging needs a CUDA-capable component (render on the OptiX "
+        f"backend, or dynamics on CUDA); import one first ({detail})"
+    )
+
+
+class CudaView:
+    """DLPack producer over one slice of an uploaded CUDA buffer. Fresh capsule per
+    __dlpack__ call; holds the parent buffer alive."""
+
+    def __init__(self, buffer, index):
+        self._buffer = buffer
+        self._index = index
+
+    def __dlpack__(self, stream=None, **_kwargs):
+        if self._index is None:
+            return self._buffer.dlpack()
+        return self._buffer.slice_dlpack(self._index)
+
+    def __dlpack_device__(self):
+        return (2, 0)  # kDLCUDA
+
+
+class TensorSet:
+    """A host float32 array uploaded to CUDA device memory once; view(i) exposes slice i
+    along axis 0 as a zero-copy DLPack CUDA tensor (e.g. one camera set per slice)."""
+
+    def __init__(self, array):
+        self._buffer = _core().CudaBuffer(np.ascontiguousarray(array, dtype=np.float32))
+
+    @property
+    def shape(self):
+        return tuple(self._buffer.shape)
+
+    @property
+    def num_sets(self):
+        return self.shape[0]
+
+    @property
+    def raw(self):
+        return self._buffer
+
+    def view(self, index=None):
+        return CudaView(self._buffer, index)
+
+
+def upload(array):
+    return TensorSet(array)

hyperdrone/dynamics/__init__.py

@@ -0,0 +1,22 @@
+"""hyperdrone.dynamics — vectorized, stateful L2F multirotor simulator.
+
+Self-contained peer of hyperdrone.render (never imports it): the coupling surface is a
+plain DLPack tensor of packed camera bases produced by Sim.camera_bases(), which the
+renderer consumes without knowing where it came from. Two JIT variants of one component:
+cpu (portable reference) and cuda (batch-parallel kernels, device-resident buffers);
+select via device= or HYPERDRONE_DYNAMICS_DEVICE (AUTO prefers cuda when available).
+"""
+from .. import cuda as _cuda
+from ._component import component, load_core, resolve_device
+from ._config import MODELS, STATE_COMPONENTS, SimConfig
+from ._sim import IDENTITY_MOUNT, Sim
+
+__all__ = [
+    "IDENTITY_MOUNT",
+    "MODELS",
+    "Sim",
+    "SimConfig",
+    "resolve_device",
+]
+
+_cuda.register_provider(load_core)

hyperdrone/dynamics/_component.py

@@ -0,0 +1,40 @@
+import os
+
+from .. import jit
+
+DEVICES = ("CPU", "CUDA")
+
+
+def resolve_device(device="auto"):
+    value = (device or "auto").upper()
+    if value == "AUTO":
+        value = os.environ.get("HYPERDRONE_DYNAMICS_DEVICE", "AUTO").upper()
+    if value == "AUTO":
+        value = "CUDA" if _cuda_available() else "CPU"
+    if value not in DEVICES:
+        raise jit.BuildError(f"hyperdrone: invalid dynamics device {value} (CPU|CUDA|AUTO)")
+    return value
+
+
+def _cuda_available():
+    from pathlib import Path
+    import shutil
+    return Path("/usr/local/cuda/bin/nvcc").exists() or shutil.which("nvcc") is not None
+
+
+def component(device="auto"):
+    resolved = resolve_device(device)
+    return jit.Component(
+        name="dynamics",
+        variant=resolved.lower(),
+        core_target="hyperdrone_dynamics_core",
+        needs_cuda=resolved == "CUDA",
+    )
+
+
+def load_core(device="auto"):
+    return jit.core(component(device))
+
+
+def ensure_sim_library(config, device="auto"):
+    return jit.ensure(component(device), config)

hyperdrone/dynamics/_config.py

@@ -0,0 +1,47 @@
+from dataclasses import dataclass
+
+from .. import jit
+
+MODELS = (
+    "crazyflie",
+    "crazyflie_openmv",
+    "mrs",
+    "x500_real",
+    "x500_sim",
+    "arpl",
+    "fs_base",
+    "flightmare",
+    "soft",
+    "soft_rigid",
+)
+
+STATE_COMPONENTS = {
+    "position": (0, 3),
+    "orientation": (1, 4),
+    "linear_velocity": (2, 3),
+    "angular_velocity": (3, 3),
+    "rpm": (4, 4),
+}
+
+
+@dataclass(frozen=True)
+class SimConfig:
+    num_drones: int
+    domain_randomization: bool
+
+    def __post_init__(self):
+        if self.num_drones < 1:
+            raise ValueError("num_drones must be >= 1")
+
+    def canonical(self):
+        # must match build_config_string() in _native/dynamics/impl.cpp exactly
+        return f"n={self.num_drones};dr={1 if self.domain_randomization else 0}"
+
+    def key(self):
+        return jit.canonical_key(self.canonical())
+
+    def defines(self):
+        return {
+            "HYPERDRONE_DYNAMICS_NUM_DRONES": self.num_drones,
+            "HYPERDRONE_DYNAMICS_DOMAIN_RANDOMIZATION": 1 if self.domain_randomization else 0,
+        }