diff --git a/pyproject.toml b/pyproject.toml index bd5d490..d86eea4 100644 --- a/pyproject.toml +++ b/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" } diff --git a/rayx/main.cpp b/rayx/main.cpp index fd2369f..ba94542 100644 --- a/rayx/main.cpp +++ b/rayx/main.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include #include @@ -7,6 +8,7 @@ #include #include #include +#include #include #include @@ -22,6 +24,75 @@ std::filesystem::path getModulePath(py::module_ m) { return std::filesystem::path(py::cast(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 { + 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(p); }); + ndarray> arr(data, {4, 4}, owner); + return type_caster::from_cpp(arr, rv_policy::reference, cleanup); + } +}; + +} // namespace nanobind::detail + namespace reflect { template <> @@ -269,7 +340,6 @@ static_assert(Structure); } // 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_(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_(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 seed, std::optional max_events, + std::optional 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(*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(*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(), py::arg("max_events") = std::optional(), + py::arg("device_index") = std::optional(), 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(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; } diff --git a/tests/test_orientation.py b/tests/test_orientation.py new file mode 100644 index 0000000..94ce5fd --- /dev/null +++ b/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)