Expose sequential tracing and fixed seeding to Python bindings

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rayx"
-version = "0.4.5"
+version = "0.4.6"
 description = "Python bindings for RAYX"
 readme = "README.md"
 license = { file = "LICENSE" }

rayx/main.cpp

@@ -1,12 +1,14 @@
 #include <Beamline/StringConversion.h>
 #include <Core.h>
+#include <Random.h>
 #include <Rml/Importer.h>
 #include <Rml/Locate.h>
 #include <Tracer/Tracer.h>
 #include <Variant.h>
 #include <nanobind/nanobind.h>
 #include <nanobind/ndarray.h>
 #include <nanobind/stl/array.h>
+#include <nanobind/stl/optional.h>
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/vector.h>
 
@@ -22,6 +24,75 @@ std::filesystem::path getModulePath(py::module_ m) {
     return std::filesystem::path(py::cast<std::string>(m.attr("__file__"))).parent_path();
 }
 
+namespace nanobind::detail {
+
+// nanobind caster for glm::dmat4x4, used for the `orientation` property of DesignElement / DesignSource.
+//
+// Convention: Python is row-major, indexed [row][col] (numpy / nested lists); glm is column-major
+// (m[col] is a column), so we map python[row][col] <-> glm m[col][row]. np.asarray(c.orientation)
+// is thus a faithful (4, 4) view and assigning it straight back is a round-trip identity.
+//
+// Caveat: rayx-core only persists the 3 direction vectors (the 3x3 rotation block). DesignSource
+// also sets the homogeneous 4th column to (0, 0, 0, 1), but DesignElement leaves it unset, so an
+// element's orientation reads back with arbitrary values in the last numpy column. setOrientation
+// ignores that column, so tracing and the 3x3 round-trip are unaffected.
+template <>
+struct type_caster<glm::dmat4x4> {
+    NB_TYPE_CASTER(glm::dmat4x4, const_name("numpy.ndarray[dtype=float64, shape=(4, 4)]"))
+
+    // Accepts a (4, 4) numpy array or a nested 4x4 sequence; rejects anything not exactly 4x4.
+    bool from_python(handle src, uint8_t /*flags*/, cleanup_list* /*cleanup*/) noexcept {
+        PyObject* obj = src.ptr();
+        if (PySequence_Length(obj) != 4) {
+            PyErr_Clear();
+            return false;
+        }
+        glm::dmat4x4 result;
+        for (int row = 0; row < 4; ++row) {
+            PyObject* row_obj = PySequence_GetItem(obj, row);
+            if (!row_obj) {
+                PyErr_Clear();
+                return false;
+            }
+            if (PySequence_Length(row_obj) != 4) {
+                PyErr_Clear();
+                Py_DECREF(row_obj);
+                return false;
+            }
+            for (int col = 0; col < 4; ++col) {
+                PyObject* elem = PySequence_GetItem(row_obj, col);
+                if (!elem) {
+                    PyErr_Clear();
+                    Py_DECREF(row_obj);
+                    return false;
+                }
+                double v = PyFloat_AsDouble(elem);
+                Py_DECREF(elem);
+                if (v == -1.0 && PyErr_Occurred()) {
+                    PyErr_Clear();
+                    Py_DECREF(row_obj);
+                    return false;
+                }
+                result[col][row] = v;
+            }
+            Py_DECREF(row_obj);
+        }
+        value = result;
+        return true;
+    }
+
+    static handle from_cpp(const glm::dmat4x4& mat, rv_policy /*policy*/, cleanup_list* cleanup) noexcept {
+        double* data = new double[16];
+        for (int row = 0; row < 4; ++row)
+            for (int col = 0; col < 4; ++col) data[row * 4 + col] = mat[col][row];
+        capsule owner(data, [](void* p) noexcept { delete[] static_cast<double*>(p); });
+        ndarray<numpy, double, shape<4, 4>> arr(data, {4, 4}, owner);
+        return type_caster<decltype(arr)>::from_cpp(arr, rv_policy::reference, cleanup);
+    }
+};
+
+}  // namespace nanobind::detail
+
 namespace reflect {
 
 template <>
@@ -269,7 +340,6 @@ static_assert(Structure<rayx::DesignSource>);
 
 }  // namespace reflect
 
-// TODO: dmat4x4
 // TODO: rays struct
 // TODO: LayerCoating
 // TODO: CurvatureType
@@ -414,6 +484,15 @@ NB_MODULE(core, m) {
         .value("BEYOND_HORIZON", rayx::EventType::BeyondHorizon)
         .value("TOO_MANY_EVENTS", rayx::EventType::TooManyEvents);
 
+    // Device classes that tracing can run on. Use these to restrict device selection in trace() and list_devices().
+    py::enum_<rayx::DeviceConfig::DeviceType>(m, "DeviceType")
+        .value("CpuSerial", rayx::DeviceConfig::DeviceType::CpuSerial)
+        .value("CpuParallel", rayx::DeviceConfig::DeviceType::CpuParallel)
+        .value("Cpu", rayx::DeviceConfig::DeviceType::Cpu)
+        .value("GpuCuda", rayx::DeviceConfig::DeviceType::GpuCuda)
+        .value("Gpu", rayx::DeviceConfig::DeviceType::Gpu)
+        .value("All", rayx::DeviceConfig::DeviceType::All);
+
     py::class_<rayx::Rays>(m, "Rays")
         .def_prop_ro("path_id", [](rayx::Rays& rays) { return to_numpy(rays.path_id); })
         .def_prop_ro("path_event_id", [](rayx::Rays& rays) { return to_numpy(rays.path_event_id); })
@@ -451,14 +530,48 @@ NB_MODULE(core, m) {
         .def_prop_ro("elements", &rayx::Beamline::getElements)
         .def_prop_ro("sources", &rayx::Beamline::getSources)
         .def("trace",
-             [](rayx::Beamline& bl) {
-                 rayx::DeviceConfig deviceConfig = rayx::DeviceConfig().enableBestDevice();
+             [](rayx::Beamline& bl, bool sequential, std::optional<uint32_t> seed, std::optional<int> max_events,
+                std::optional<int> device_index, rayx::DeviceConfig::DeviceType device_type) {
+                 // Seed the RNG: a given seed yields deterministic results, otherwise the seed is derived from system time.
+                 if (seed)
+                     rayx::fixSeed(*seed);
+                 else
+                     rayx::randomSeed();
+
+                 // We validate device_index here rather than let enableDeviceByIndex() call RAYX_EXIT,
+                 // which would terminate the Python interpreter.
+                 rayx::DeviceConfig deviceConfig(device_type);
+                 if (deviceConfig.devices.empty())
+                     throw std::runtime_error(
+                         "No compute device available for the requested device_type. Use list_devices() to see the available devices "
+                         "(note: GPU tracing requires a CUDA-enabled build and an NVIDIA GPU).");
+                 if (device_index) {
+                     if (*device_index < 0 || static_cast<size_t>(*device_index) >= deviceConfig.devices.size())
+                         throw std::out_of_range("device_index " + std::to_string(*device_index) +
+                                                 " is out of range; use list_devices() to see the available devices");
+                     deviceConfig.enableDeviceByIndex(static_cast<size_t>(*device_index));
+                 } else {
+                     deviceConfig.enableBestDevice();
+                 }
+
                  rayx::Tracer tracer = rayx::Tracer(deviceConfig);
                  rayx::ObjectMask obj_mask = rayx::ObjectMask::all();
                  rayx::RayAttrMask attr_mask = rayx::RayAttrMask::All;
-                 rayx::Rays rays = tracer.trace(bl, rayx::Sequential::No, obj_mask, attr_mask, std::nullopt, std::nullopt);
+                 rayx::Sequential seq = sequential ? rayx::Sequential::Yes : rayx::Sequential::No;
+                 rayx::Rays rays = tracer.trace(bl, seq, obj_mask, attr_mask, max_events, std::nullopt);
                  return rays;
-             })
+             },
+             py::arg("sequential") = false, py::arg("seed") = std::optional<uint32_t>(), py::arg("max_events") = std::optional<int>(),
+             py::arg("device_index") = std::optional<int>(), py::arg("device_type") = rayx::DeviceConfig::DeviceType::All,
+             "Trace rays through the beamline.\n\n"
+             "sequential: if True, rays hit elements in beamline order (sequential tracing); "
+             "if False (default), tracing is non-sequential.\n"
+             "seed: if given, fixes the RNG seed for deterministic results; if None (default), a random seed is used.\n"
+             "max_events: optional maximum number of events recorded per ray (only used in non-sequential tracing).\n"
+             "device_index: optional index of the compute device to use (see list_devices()); if None (default), the best "
+             "available device is chosen automatically.\n"
+             "device_type: restrict device selection to a device class (DeviceType.Cpu, DeviceType.Gpu, or DeviceType.All; "
+             "default All).")
         .def("__getitem__", [](rayx::Beamline& bl, const std::string& name) {
             for (auto element : bl.getElements()) {
                 if (element->getName() == name) {
@@ -474,4 +587,39 @@ NB_MODULE(core, m) {
         });
 
     m.def("import_beamline", [](std::string path) { return rayx::importBeamline(path); }, "Import a beamline from an RML file", py::arg("path"));
+
+    m.def(
+        "list_devices",
+        [](rayx::DeviceConfig::DeviceType device_type) {
+            // Mirror the CLI's --list-devices: enumerate the available compute devices of the requested type.
+            auto typeName = [](rayx::DeviceConfig::DeviceType t) -> const char* {
+                switch (t) {
+                    case rayx::DeviceConfig::DeviceType::CpuSerial: return "CpuSerial";
+                    case rayx::DeviceConfig::DeviceType::CpuParallel: return "CpuParallel";
+                    case rayx::DeviceConfig::DeviceType::GpuCuda: return "GpuCuda";
+                    default: return "Unknown";
+                }
+            };
+            rayx::DeviceConfig config(device_type);
+            py::list result;
+            for (size_t i = 0; i < config.devices.size(); ++i) {
+                const auto& device = config.devices[i];
+                py::dict entry;
+                entry["index"] = static_cast<int>(i);
+                entry["name"] = device.name;
+                entry["type"] = typeName(device.type);
+                result.append(entry);
+            }
+            return result;
+        },
+        py::arg("device_type") = rayx::DeviceConfig::DeviceType::All,
+        "List the compute devices available for tracing as a list of dicts with 'index', 'name' and 'type'.\n"
+        "The 'index' is what you pass to Beamline.trace(device_index=...).\n"
+        "device_type: restrict the listing to a device class (DeviceType.Cpu, DeviceType.Gpu, or DeviceType.All; default All).\n"
+        "Note: GPU devices only appear in a CUDA-enabled build running on a machine with an NVIDIA GPU.");
+
+    m.def("fix_seed", &rayx::fixSeed, py::arg("seed") = rayx::FIXED_SEED,
+          "Fix the global RNG seed so that subsequent traces are deterministic. Defaults to the canonical fixed test seed.");
+    m.def("random_seed", &rayx::randomSeed, "Seed the global RNG randomly (based on system time).");
+    m.attr("FIXED_SEED") = rayx::FIXED_SEED;
 }

tests/test_orientation.py

@@ -0,0 +1,126 @@
+# tests/test_orientation.py
+"""Tests for the glm::dmat4x4 <-> numpy caster behind the `orientation` property.
+
+See the caster in main.cpp for the row/col convention and the rayx-core caveat:
+elements only persist the 3x3 rotation block, so these tests assert the 3x3 block
+for elements and the full 4x4 for sources.
+"""
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+sys.path.insert(0, str(Path(__file__).parent.parent / "python"))
+
+import rayx
+
+RML_FILE = Path(__file__).parent.parent / "examples" / "METRIX_U41_G1_H1_318eV_PS_MLearn_v114.rml"
+
+
+@pytest.fixture(scope="module")
+def beamline():
+    return rayx.import_beamline(str(RML_FILE))
+
+
+@pytest.fixture
+def element(beamline):
+    return beamline.elements[0]
+
+
+@pytest.fixture
+def source(beamline):
+    return beamline.sources[0]
+
+
+# --- read ---------------------------------------------------------------
+
+def test_element_orientation_reads_as_4x4(element):
+    a = np.asarray(element.orientation)
+    assert a.shape == (4, 4)
+    assert a.dtype == np.float64
+
+
+def test_source_orientation_reads_as_4x4(source):
+    a = np.asarray(source.orientation)
+    assert a.shape == (4, 4)
+    assert a.dtype == np.float64
+
+
+def test_source_orientation_homogeneous_row(source):
+    # DesignSource fills the homogeneous parts, so it is a clean rigid-transform matrix.
+    a = np.asarray(source.orientation)
+    assert np.allclose(a[3, :], [0, 0, 0, 1])
+    assert np.allclose(a[:3, 3], [0, 0, 0])
+
+
+# --- write + read-back --------------------------------------------------
+
+def test_element_write_rotation_block_persists(element):
+    # A proper rotation (90 deg about z) written into the 3x3 block must read back exactly.
+    m = np.eye(4)
+    m[:3, :3] = [[0, -1, 0], [1, 0, 0], [0, 0, 1]]
+    element.orientation = m
+    back = np.asarray(element.orientation)
+    assert np.allclose(back[:3, :3], m[:3, :3])
+
+
+def test_source_write_full_matrix_persists(source):
+    m = np.eye(4)
+    m[:3, :3] = [[0, 0, 1], [0, 1, 0], [-1, 0, 0]]
+    source.orientation = m
+    assert np.allclose(np.asarray(source.orientation), m)
+
+
+def test_write_accepts_nested_list(source):
+    m = [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
+    source.orientation = m
+    assert np.allclose(np.asarray(source.orientation), np.eye(4))
+
+
+def test_column_convention_not_transposed(source):
+    # An asymmetric (but homogeneous-bottom-row) matrix must come back un-transposed,
+    # proving python[row][col] maps to the same logical element (not its transpose).
+    m = np.array(
+        [[2.0, 3.0, 5.0, 0.0],
+         [7.0, 11.0, 13.0, 0.0],
+         [17.0, 19.0, 23.0, 0.0],
+         [0.0, 0.0, 0.0, 1.0]]
+    )
+    source.orientation = m
+    assert np.allclose(np.asarray(source.orientation), m)
+
+
+# --- bad input rejected -------------------------------------------------
+
+@pytest.mark.parametrize("bad", [np.eye(3), np.arange(16.0), np.zeros((4, 3)), 5.0])
+def test_wrong_shape_rejected(source, bad):
+    with pytest.raises(TypeError):
+        source.orientation = bad
+
+
+# --- round-trip identity through tracing --------------------------------
+
+def _stack(rays):
+    return np.column_stack(
+        [
+            np.asarray(rays.path_id),
+            np.asarray(rays.object_id),
+            np.asarray(rays.position_x),
+            np.asarray(rays.position_y),
+            np.asarray(rays.position_z),
+        ]
+    )
+
+
+def test_roundtrip_orientation_leaves_trace_unchanged(beamline):
+    # Capture every component's orientation, trace, reassign np.asarray(orientation)
+    # back verbatim, trace again with the same fixed seed, and assert the rays match.
+    SEED = rayx.FIXED_SEED
+    baseline = _stack(beamline.trace(seed=SEED))
+
+    for comp in list(beamline.elements) + list(beamline.sources):
+        comp.orientation = np.asarray(comp.orientation)
+
+    after = _stack(beamline.trace(seed=SEED))
+    assert np.array_equal(baseline, after)