python/src/scene.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | #include <cstddef> | ||
| 3 | #include <cstdint> | ||
| 4 | #include <memory> | ||
| 5 | #include <optional> | ||
| 6 | #include <span> | ||
| 7 | #include <string> | ||
| 8 | #include <string_view> | ||
| 9 | #include <tuple> | ||
| 10 | #include <utility> | ||
| 11 | #include <vector> | ||
| 12 | |||
| 13 | #include <nanobind/ndarray.h> | ||
| 14 | #include <nanobind/stl/map.h> | ||
| 15 | #include <nanobind/stl/optional.h> | ||
| 16 | #include <nanobind/stl/string.h> | ||
| 17 | #include <nanobind/stl/string_view.h> | ||
| 18 | #include <nanobind/stl/tuple.h> | ||
| 19 | #include <nanobind/stl/vector.h> | ||
| 20 | |||
| 21 | #include "chrono_batch_render/device/queue.h" | ||
| 22 | #include "chrono_batch_render/render/shaders.h" | ||
| 23 | #include "chrono_batch_render/scene/scene_impl.h" | ||
| 24 | #include "chrono_batch_render/stage/views/poses.h" | ||
| 25 | #include "values.h" | ||
| 26 | |||
| 27 | namespace nb = nanobind; | ||
| 28 | using namespace nb::literals; | ||
| 29 | |||
| 30 | namespace chrono::batch_render::python { | ||
| 31 | |||
| 32 | 536 | SchemaId FindSchema(std::string_view name) { | |
| 33 | 13970 | for (uint32_t index = 0; index < SchemaCount(); ++index) { | |
| 34 | 27940 | if (Info(SchemaId(index)).name == name) { | |
| 35 | 536 | return SchemaId(index); | |
| 36 | } | ||
| 37 | } | ||
| 38 | ✗ | throw nb::value_error(("No schema is called " + std::string(name)).c_str()); | |
| 39 | ✗ | } | |
| 40 | |||
| 41 | namespace { | ||
| 42 | |||
| 43 | 31 | struct SceneSelection { | |
| 44 | Ref<Scene> scene; | ||
| 45 | View view; | ||
| 46 | |||
| 47 | 4 | std::vector<Prim> Prims() const { | |
| 48 | 4 | std::vector<Prim> prims; | |
| 49 | 111 | for (const uint32_t index : view.Prims()) { | |
| 50 | 214 | prims.push_back(scene->ByIndex(index)); | |
| 51 | } | ||
| 52 | 4 | return prims; | |
| 53 | ✗ | } | |
| 54 | }; | ||
| 55 | |||
| 56 | ✗ | struct SceneMapping { | |
| 57 | Ref<Scene> scene; | ||
| 58 | Mapping mapping; | ||
| 59 | std::optional<uintptr_t> stream; | ||
| 60 | }; | ||
| 61 | |||
| 62 | 41114 | SceneSelection Single(const ScenePrim& prim) { | |
| 63 | 41114 | return {prim.scene, prim.scene->Select(prim.prim)}; | |
| 64 | 41114 | } | |
| 65 | |||
| 66 | 7 | std::vector<ScenePrim> ScenePrims(const Ref<Scene>& scene, std::span<const Prim> prims) { | |
| 67 | 7 | std::vector<ScenePrim> wrapped; | |
| 68 | 129 | for (const Prim& prim : prims) { | |
| 69 | 244 | wrapped.push_back({scene, prim}); | |
| 70 | } | ||
| 71 | 7 | return wrapped; | |
| 72 | 122 | } | |
| 73 | |||
| 74 | 114 | bool IsText(ColumnType type) { | |
| 75 | 114 | return type == ColumnType::Enum || type == ColumnType::Token || type == ColumnType::Asset; | |
| 76 | } | ||
| 77 | |||
| 78 | 41063 | std::string NumpyType(DLDataType dtype) { | |
| 79 | 41063 | const char* kind = dtype.code == kDLFloat ? "float" : dtype.code == kDLInt ? "int" : "uint"; | |
| 80 | 82126 | return kind + std::to_string(dtype.bits); | |
| 81 | } | ||
| 82 | |||
| 83 | 41074 | std::vector<std::byte> ConvertCells(Scene& scene, const ColumnInfo& column, nb::handle value) { | |
| 84 | 41074 | switch (column.type) { | |
| 85 | 4 | case ColumnType::Enum: { | |
| 86 | 4 | std::vector<uint32_t> words; | |
| 87 | 4 | if (nb::isinstance<nb::str>(value) || nb::isinstance<nb::int_>(value)) { | |
| 88 | 4 | words.push_back(EnumIndex(value, column.tokens)); | |
| 89 | } else { | ||
| 90 | ✗ | for (nb::handle item : value) { | |
| 91 | ✗ | words.push_back(EnumIndex(item, column.tokens)); | |
| 92 | ✗ | } | |
| 93 | } | ||
| 94 | 3 | return Bytes(words); | |
| 95 | 4 | } | |
| 96 | 7 | case ColumnType::Token: | |
| 97 | 7 | case ColumnType::Asset: { | |
| 98 | 7 | std::vector<uint32_t> words; | |
| 99 | 14 | for (const std::string& text : Strings(value)) { | |
| 100 | 14 | words.push_back(scene.Token(text).id); | |
| 101 | 7 | } | |
| 102 | 7 | return Bytes(words); | |
| 103 | 7 | } | |
| 104 | ✗ | case ColumnType::Bool: | |
| 105 | ✗ | return NumpyBytes(nb::module_::import_("numpy").attr("asarray")(value, "dtype"_a = "bool"), | |
| 106 | ✗ | "uint32"); | |
| 107 | 41063 | default: | |
| 108 | 82125 | return NumpyBytes(value, NumpyType(Layout(column.type).dtype)); | |
| 109 | } | ||
| 110 | } | ||
| 111 | |||
| 112 | 41074 | bool OnDevice(nb::handle value) { | |
| 113 | 41074 | if (!nb::hasattr(value, "__dlpack_device__")) { | |
| 114 | return false; | ||
| 115 | } | ||
| 116 | 2 | const nb::tuple device = nb::cast<nb::tuple>(value.attr("__dlpack_device__")()); | |
| 117 | 2 | return nb::cast<int>(device[0]) != kDLCPU; | |
| 118 | 1 | } | |
| 119 | |||
| 120 | 1138 | nb::object MaskWords(nb::handle mask) { | |
| 121 | 1138 | if (mask.is_none() || nb::hasattr(mask, "__dlpack_device__")) { | |
| 122 | 1138 | const bool host = | |
| 123 | 1138 | mask.is_none() || | |
| 124 | 1143 | nb::cast<int>(nb::cast<nb::tuple>(mask.attr("__dlpack_device__")())[0]) == kDLCPU; | |
| 125 | 1138 | if (!host || mask.is_none()) { | |
| 126 | 1137 | return nb::borrow(mask); | |
| 127 | } | ||
| 128 | } | ||
| 129 | 1 | const nb::module_ numpy = nb::module_::import_("numpy"); | |
| 130 | 1 | const nb::object array = numpy.attr("asarray")(mask); | |
| 131 | 2 | if (nb::cast<std::string>(array.attr("dtype").attr("kind")) != "b") { | |
| 132 | 1 | return array; | |
| 133 | } | ||
| 134 | ✗ | nb::object bytes = numpy.attr("packbits")(array, "bitorder"_a = "little"); | |
| 135 | ✗ | const size_t padding = (8 - nb::len(bytes) % 8) % 8; | |
| 136 | ✗ | bytes = numpy.attr("pad")(bytes, nb::make_tuple(0, padding)); | |
| 137 | ✗ | return bytes.attr("view")(numpy.attr("uint64")); | |
| 138 | 1 | } | |
| 139 | |||
| 140 | ✗ | void WriteTensor(Scene& scene, UntypedAttribute& target, nb::handle value, | |
| 141 | std::optional<uintptr_t> stream, nb::handle mask = nb::none()) { | ||
| 142 | ✗ | HeldTensors held; | |
| 143 | ✗ | if (nb::isinstance<device::Tensor>(value)) { | |
| 144 | ✗ | const nb::object words = MaskWords(mask); | |
| 145 | ✗ | const device::Tensor bits = | |
| 146 | ✗ | words.is_none() ? device::Tensor() | |
| 147 | ✗ | : nb::isinstance<device::Tensor>(words) | |
| 148 | ✗ | ? nb::cast<device::Tensor>(words) | |
| 149 | ✗ | : held.Import(scene.GetDevice(), words, stream, "attribute mask"); | |
| 150 | ✗ | const device::Tensor values = nb::cast<device::Tensor>(value); | |
| 151 | ✗ | const Stream producer = | |
| 152 | ✗ | held.Producer().IsNone() && stream ? ToStream(stream) : held.Producer(); | |
| 153 | ✗ | nb::gil_scoped_release release; | |
| 154 | ✗ | held.Retire(scene, target.Write(values, producer, bits)); | |
| 155 | ✗ | return; | |
| 156 | ✗ | } | |
| 157 | ✗ | const DLTensor* values = held.Take(value, stream); | |
| 158 | ✗ | const DLTensor* bits = held.Take(MaskWords(mask), stream); | |
| 159 | ✗ | nb::gil_scoped_release release; | |
| 160 | ✗ | held.Retire(scene, target.Write(*values, held.Producer(), bits)); | |
| 161 | ✗ | } | |
| 162 | |||
| 163 | 41074 | void AssignCells(Scene& scene, UntypedAttribute& target, nb::handle value, | |
| 164 | std::optional<uintptr_t> stream) { | ||
| 165 | 41074 | const ColumnInfo& column = *target.Reference().column; | |
| 166 | 41074 | if (OnDevice(value)) { | |
| 167 | ✗ | return WriteTensor(scene, target, value, stream); | |
| 168 | } | ||
| 169 | |||
| 170 | 41074 | const std::vector<std::byte> cells = ConvertCells(scene, column, value); | |
| 171 | 41072 | if (cells.size() == ElementBytes(column.type)) { | |
| 172 | 41069 | target.Fill(cells); | |
| 173 | } else { | ||
| 174 | 4 | target.Write(cells); | |
| 175 | } | ||
| 176 | 41069 | } | |
| 177 | |||
| 178 | 26 | bool IsPrimvars(const ColumnRef& column) { | |
| 179 | 26 | return column.column->field == "primvars"; | |
| 180 | } | ||
| 181 | |||
| 182 | 13 | void SetArray(UntypedAttribute& target, uint32_t index, nb::handle value) { | |
| 183 | 13 | const ColumnRef column = target.Reference(); | |
| 184 | 13 | switch (column.column->type) { | |
| 185 | 3 | case ColumnType::PrimArray: | |
| 186 | 3 | return target.SetTargets(index, Strings(value)); | |
| 187 | ✗ | case ColumnType::TokenArray: | |
| 188 | ✗ | return target.SetTokens(index, Strings(value)); | |
| 189 | 8 | case ColumnType::FloatArray: | |
| 190 | 8 | case ColumnType::Float2Array: | |
| 191 | 8 | case ColumnType::Float3Array: | |
| 192 | 8 | case ColumnType::Float4Array: | |
| 193 | 24 | return target.SetArray(index, NumpyBytes(value, "float32")); | |
| 194 | ✗ | case ColumnType::Int64Array: | |
| 195 | ✗ | return target.SetArray(index, NumpyBytes(value, "int64")); | |
| 196 | 2 | default: | |
| 197 | 2 | if (IsPrimvars(column)) { | |
| 198 | ✗ | return target.SetPrimvars(index, nb::cast<std::vector<Primvar>>(value)); | |
| 199 | } | ||
| 200 | 6 | return target.SetArray(index, NumpyBytes(value, "int32")); | |
| 201 | } | ||
| 202 | } | ||
| 203 | |||
| 204 | 14 | size_t ArrayLanes(ColumnType type) { | |
| 205 | 14 | switch (type) { | |
| 206 | case ColumnType::Float2Array: | ||
| 207 | return 2; | ||
| 208 | 9 | case ColumnType::Float3Array: | |
| 209 | 9 | return 3; | |
| 210 | ✗ | case ColumnType::Float4Array: | |
| 211 | ✗ | return 4; | |
| 212 | 5 | default: | |
| 213 | 5 | return 1; | |
| 214 | } | ||
| 215 | } | ||
| 216 | |||
| 217 | 14 | bool IsNumericArray(const ColumnRef& column) { | |
| 218 | 14 | return column.column->type != ColumnType::PrimArray && | |
| 219 | 14 | column.column->type != ColumnType::TokenArray && !IsPrimvars(column); | |
| 220 | } | ||
| 221 | |||
| 222 | 3 | std::string ShapeText(const nb::object& array) { | |
| 223 | 6 | return nb::cast<std::string>(array.attr("shape").attr("__str__")()); | |
| 224 | } | ||
| 225 | |||
| 226 | 10 | void CheckRows(const UntypedAttribute& target, std::string_view field, nb::handle value) { | |
| 227 | 10 | const ColumnRef column = target.Reference(); | |
| 228 | 10 | const size_t lanes = ArrayLanes(column.column->type); | |
| 229 | 10 | if (!IsNumericArray(column) || lanes == 1) { | |
| 230 | 5 | return; | |
| 231 | } | ||
| 232 | 10 | const nb::object array = nb::module_::import_("numpy").attr("asarray")(value); | |
| 233 | 10 | const size_t rank = nb::cast<size_t>(array.attr("ndim")); | |
| 234 | 10 | const size_t size = nb::cast<size_t>(array.attr("size")); | |
| 235 | 24 | const bool rows = rank == 2 && nb::cast<size_t>(array.attr("shape")[nb::int_(1)]) == lanes; | |
| 236 | 5 | const bool flat = rank == 1 && size % lanes == 0; | |
| 237 | 5 | if (!rows && !flat) { | |
| 238 | 6 | throw nb::value_error((std::string(field) + " takes rows of " + std::to_string(lanes) + | |
| 239 | 10 | " values, shaped (rows, " + std::to_string(lanes) + "), got " + | |
| 240 | 2 | ShapeText(array)) | |
| 241 | 8 | .c_str()); | |
| 242 | } | ||
| 243 | 3 | } | |
| 244 | |||
| 245 | 4 | void AssignArrays(SceneSelection& source, UntypedAttribute& target, std::string_view field, | |
| 246 | nb::handle value) { | ||
| 247 | 4 | const uint32_t count = source.view.PrimCount(); | |
| 248 | 4 | const ColumnRef column = target.Reference(); | |
| 249 | 4 | if (!IsNumericArray(column)) { | |
| 250 | ✗ | for (uint32_t index = 0; index < count; ++index) { | |
| 251 | ✗ | SetArray(target, index, value); | |
| 252 | } | ||
| 253 | 1 | return; | |
| 254 | } | ||
| 255 | 8 | const nb::object array = nb::module_::import_("numpy").attr("asarray")(value); | |
| 256 | 4 | const size_t lanes = ArrayLanes(column.column->type); | |
| 257 | 8 | const size_t rank = nb::cast<size_t>(array.attr("ndim")); | |
| 258 | 4 | const size_t batched = lanes > 1 ? 3 : 2; | |
| 259 | 4 | if (rank == batched) { | |
| 260 | 4 | const nb::object shape = array.attr("shape"); | |
| 261 | 4 | const bool leading = nb::cast<size_t>(shape[nb::int_(0)]) == count; | |
| 262 | 8 | const bool rows = lanes == 1 || nb::cast<size_t>(shape[nb::int_(2)]) == lanes; | |
| 263 | 2 | if (!leading || !rows) { | |
| 264 | 1 | const std::string row = lanes > 1 ? ", rows, " + std::to_string(lanes) : ", rows"; | |
| 265 | 3 | throw nb::value_error((std::string(field) + " takes one array for each of the " + | |
| 266 | 4 | std::to_string(count) + " prims, shaped (" + | |
| 267 | 6 | std::to_string(count) + row + "), got " + ShapeText(array) + | |
| 268 | ". Use set_each for a different number of rows per prim") | ||
| 269 | 4 | .c_str()); | |
| 270 | 1 | } | |
| 271 | 4 | for (uint32_t index = 0; index < count; ++index) { | |
| 272 | 4 | const nb::object items = array[nb::int_(index)]; | |
| 273 | 3 | SetArray(target, index, items); | |
| 274 | 3 | } | |
| 275 | 1 | return; | |
| 276 | 2 | } | |
| 277 | 2 | if (rank > batched || (lanes == 1 && rank > 1)) { | |
| 278 | ✗ | throw nb::value_error((std::string(field) + | |
| 279 | ✗ | " takes one array for every prim or one per prim, got " + | |
| 280 | ✗ | ShapeText(array)) | |
| 281 | ✗ | .c_str()); | |
| 282 | } | ||
| 283 | 2 | CheckRows(target, field, array); | |
| 284 | 4 | for (uint32_t index = 0; index < count; ++index) { | |
| 285 | 3 | SetArray(target, index, array); | |
| 286 | } | ||
| 287 | 2 | } | |
| 288 | |||
| 289 | 41078 | void Assign(SceneSelection& source, std::string_view field, nb::handle value, | |
| 290 | std::optional<uintptr_t> stream = std::nullopt, nb::handle mask = nb::none()) { | ||
| 291 | 41078 | UntypedAttribute target = source.view.Get(field); | |
| 292 | 41078 | if (!mask.is_none()) { | |
| 293 | ✗ | if (IsArray(target.Reference().column->type)) { | |
| 294 | ✗ | throw nb::value_error("Masked writes take cell attributes, not arrays"); | |
| 295 | } | ||
| 296 | ✗ | const nb::object tensor = | |
| 297 | ✗ | nb::hasattr(value, "__dlpack__") | |
| 298 | ✗ | ? nb::borrow(value) | |
| 299 | 7 | : nb::module_::import_("numpy").attr("ascontiguousarray")(value); | |
| 300 | ✗ | return WriteTensor(*source.scene, target, tensor, stream, mask); | |
| 301 | ✗ | } | |
| 302 | 41078 | if (!IsArray(target.Reference().column->type)) { | |
| 303 | 41074 | return AssignCells(*source.scene, target, value, stream); | |
| 304 | } | ||
| 305 | 4 | AssignArrays(source, target, field, value); | |
| 306 | 41071 | } | |
| 307 | |||
| 308 | 3 | void AssignEach(SceneSelection& source, std::string_view field, nb::handle values) { | |
| 309 | 3 | UntypedAttribute target = source.view.Get(field); | |
| 310 | 3 | const uint32_t count = source.view.PrimCount(); | |
| 311 | 6 | if (nb::len(values) != count) { | |
| 312 | 2 | throw nb::value_error(("set_each(\"" + std::string(field) + | |
| 313 | 4 | "\") takes one value for each of the " + std::to_string(count) + | |
| 314 | 3 | " prims, got " + std::to_string(nb::len(values))) | |
| 315 | 4 | .c_str()); | |
| 316 | } | ||
| 317 | 2 | if (!IsArray(target.Reference().column->type)) { | |
| 318 | 1 | return AssignCells(*source.scene, target, values, std::nullopt); | |
| 319 | } | ||
| 320 | 2 | uint32_t index = 0; | |
| 321 | 6 | for (nb::handle value : values) { | |
| 322 | 4 | CheckRows(target, field, value); | |
| 323 | 4 | SetArray(target, index++, value); | |
| 324 | 4 | } | |
| 325 | 2 | } | |
| 326 | |||
| 327 | 13 | nb::object ReadArray(const UntypedAttribute& origin, uint32_t index) { | |
| 328 | 13 | const ColumnRef column = origin.Reference(); | |
| 329 | 13 | switch (column.column->type) { | |
| 330 | ✗ | case ColumnType::PrimArray: | |
| 331 | ✗ | return nb::cast(origin.GetTargets(index)); | |
| 332 | ✗ | case ColumnType::TokenArray: | |
| 333 | ✗ | return nb::cast(origin.GetTokens(index)); | |
| 334 | 13 | default: | |
| 335 | 13 | if (IsPrimvars(column)) { | |
| 336 | ✗ | return nb::cast(origin.GetPrimvars(index)); | |
| 337 | } | ||
| 338 | } | ||
| 339 | |||
| 340 | 13 | std::vector<std::byte> bytes; | |
| 341 | 13 | { | |
| 342 | 13 | nb::gil_scoped_release release; | |
| 343 | 26 | bytes = origin.GetArray(index); | |
| 344 | ✗ | } | |
| 345 | 13 | if (column.column->type == ColumnType::FloatArray) { | |
| 346 | ✗ | const size_t count = bytes.size() / 4; | |
| 347 | ✗ | return Numpy(std::move(bytes), {count}, {kDLFloat, 32, 1}); | |
| 348 | } | ||
| 349 | 13 | if (column.column->type == ColumnType::Float2Array || | |
| 350 | 2 | column.column->type == ColumnType::Float3Array || | |
| 351 | column.column->type == ColumnType::Float4Array) { | ||
| 352 | 11 | const size_t lanes = column.column->type == ColumnType::Float2Array ? 2 | |
| 353 | : column.column->type == ColumnType::Float3Array ? 3 | ||
| 354 | : 4; | ||
| 355 | 11 | const size_t rows = bytes.size() / (lanes * 4); | |
| 356 | 22 | return Numpy(std::move(bytes), {rows, lanes}, {kDLFloat, 32, 1}); | |
| 357 | } | ||
| 358 | 2 | if (column.column->type == ColumnType::Int64Array) { | |
| 359 | ✗ | const size_t count = bytes.size() / 8; | |
| 360 | ✗ | return Numpy(std::move(bytes), {count}, {kDLInt, 64, 1}); | |
| 361 | } | ||
| 362 | 2 | const size_t count = bytes.size() / 4; | |
| 363 | 4 | return Numpy(std::move(bytes), {count}, {kDLInt, 32, 1}); | |
| 364 | } | ||
| 365 | |||
| 366 | 63 | nb::object ReadCells(Scene& scene, const UntypedAttribute& origin) { | |
| 367 | 63 | std::vector<std::byte> bytes; | |
| 368 | 63 | { | |
| 369 | 63 | nb::gil_scoped_release release; | |
| 370 | 126 | bytes = origin.ReadHost(); | |
| 371 | ✗ | } | |
| 372 | |||
| 373 | 63 | const ColumnInfo& column = *origin.Reference().column; | |
| 374 | 63 | const ColumnType type = column.type; | |
| 375 | 63 | if (IsText(type)) { | |
| 376 | 38 | nb::list texts; | |
| 377 | 76 | for (const uint32_t word : Items<uint32_t>(bytes)) { | |
| 378 | 38 | const std::string_view text = | |
| 379 | 38 | type == ColumnType::Enum ? column.tokens[word] : scene.Text(Token{word}); | |
| 380 | 76 | texts.append(nb::str(text.data(), text.size())); | |
| 381 | 38 | } | |
| 382 | 38 | return texts; | |
| 383 | 38 | } | |
| 384 | 25 | if (type == ColumnType::Bool) { | |
| 385 | ✗ | return Numpy(std::move(bytes), {origin.Size()}, {kDLUInt, 32, 1}).attr("astype")("bool"); | |
| 386 | } | ||
| 387 | |||
| 388 | 25 | const ColumnLayout layout = Layout(type); | |
| 389 | 25 | std::vector<size_t> shape = {origin.Size()}; | |
| 390 | 25 | if (layout.lanes > 1) { | |
| 391 | 12 | shape.push_back(layout.lanes); | |
| 392 | } | ||
| 393 | 50 | return Numpy(std::move(bytes), shape, layout.dtype); | |
| 394 | 88 | } | |
| 395 | |||
| 396 | 68 | nb::object Read(SceneSelection& source, std::string_view field) { | |
| 397 | 68 | const UntypedAttribute origin = source.view.Get(field); | |
| 398 | 68 | if (!IsArray(origin.Reference().column->type)) { | |
| 399 | 63 | return ReadCells(*source.scene, origin); | |
| 400 | } | ||
| 401 | 5 | nb::list values; | |
| 402 | 18 | for (uint32_t index = 0; index < origin.Size(); ++index) { | |
| 403 | 26 | values.append(ReadArray(origin, index)); | |
| 404 | } | ||
| 405 | 5 | return values; | |
| 406 | 73 | } | |
| 407 | |||
| 408 | 41063 | void AssignOne(ScenePrim& prim, std::string_view field, nb::handle value, | |
| 409 | std::optional<uintptr_t> stream) { | ||
| 410 | 41063 | SceneSelection single = Single(prim); | |
| 411 | 41063 | UntypedAttribute target = single.view.Get(field); | |
| 412 | 41063 | if (IsArray(target.Reference().column->type)) { | |
| 413 | 4 | CheckRows(target, field, value); | |
| 414 | 3 | return SetArray(target, 0, value); | |
| 415 | } | ||
| 416 | 82119 | Assign(single, field, value, stream); | |
| 417 | 41063 | } | |
| 418 | |||
| 419 | 51 | nb::object ReadFirst(SceneSelection& single, std::string_view field) { | |
| 420 | 51 | const nb::object values = Read(single, field); | |
| 421 | 51 | const ColumnType type = single.view.Column(field).column->type; | |
| 422 | 51 | if (IsArray(type) || IsText(type)) { | |
| 423 | 74 | return values[0]; | |
| 424 | } | ||
| 425 | |||
| 426 | 14 | const nb::object first = values[0]; | |
| 427 | 42 | return Layout(type).lanes > 1 ? nb::tuple(first.attr("tolist")()) : first.attr("item")(); | |
| 428 | 65 | } | |
| 429 | |||
| 430 | ✗ | AttachFrom FindAttachFrom(std::string_view from) { | |
| 431 | ✗ | if (from == "tensor") { | |
| 432 | return AttachFrom::Tensor; | ||
| 433 | } | ||
| 434 | ✗ | if (from == "scene") { | |
| 435 | return AttachFrom::Scene; | ||
| 436 | } | ||
| 437 | ✗ | throw nb::value_error("Values are attached from the tensor or the scene"); | |
| 438 | } | ||
| 439 | |||
| 440 | ✗ | void Attach(const SceneSelection& source, std::string_view field, nb::handle tensor, | |
| 441 | std::string_view initial, std::optional<uintptr_t> stream) { | ||
| 442 | ✗ | UntypedAttribute target = source.view.Get(field); | |
| 443 | ✗ | const AttachFrom values = FindAttachFrom(initial); | |
| 444 | ✗ | const TakenTensor taken = TakeTensor(tensor, stream); | |
| 445 | ✗ | try { | |
| 446 | ✗ | nb::gil_scoped_release release; | |
| 447 | ✗ | target.Attach(taken.managed, taken.producer, values); | |
| 448 | ✗ | } catch (...) { | |
| 449 | ✗ | if (taken.managed->deleter) { | |
| 450 | ✗ | taken.managed->deleter(taken.managed); | |
| 451 | } | ||
| 452 | ✗ | throw; | |
| 453 | ✗ | } | |
| 454 | ✗ | } | |
| 455 | |||
| 456 | ✗ | void MarkChanged(const SceneSelection& source, std::string_view field, | |
| 457 | std::optional<uintptr_t> stream, nb::handle mask) { | ||
| 458 | ✗ | UntypedAttribute target = source.view.Get(field); | |
| 459 | ✗ | HeldTensors held; | |
| 460 | ✗ | const DLTensor* bits = held.Take(MaskWords(mask), stream); | |
| 461 | ✗ | nb::gil_scoped_release release; | |
| 462 | ✗ | held.Retire(*source.scene, target.MarkChanged(ToStream(stream), bits)); | |
| 463 | ✗ | } | |
| 464 | |||
| 465 | ✗ | void WaitForReads(const SceneSelection& source, std::string_view field, | |
| 466 | std::optional<uintptr_t> stream) { | ||
| 467 | ✗ | const UntypedAttribute target = source.view.Get(field); | |
| 468 | ✗ | nb::gil_scoped_release release; | |
| 469 | ✗ | target.WaitForReads(stream ? ToStream(stream) : Stream::Host()); | |
| 470 | ✗ | } | |
| 471 | |||
| 472 | ✗ | void Detach(const SceneSelection& source, std::string_view field, std::optional<uintptr_t> stream) { | |
| 473 | ✗ | UntypedAttribute target = source.view.Get(field); | |
| 474 | ✗ | nb::gil_scoped_release release; | |
| 475 | ✗ | target.Detach(ToStream(stream)); | |
| 476 | ✗ | } | |
| 477 | |||
| 478 | 1149 | QuaternionOrder FindOrder(std::string_view order) { | |
| 479 | 1149 | if (order == "xyzw") { | |
| 480 | return QuaternionOrder::Xyzw; | ||
| 481 | } | ||
| 482 | 2 | if (order == "wxyz") { | |
| 483 | return QuaternionOrder::Wxyz; | ||
| 484 | } | ||
| 485 | ✗ | throw nb::value_error("Quaternions are ordered xyzw or wxyz"); | |
| 486 | } | ||
| 487 | |||
| 488 | 3 | MapAccess Access(std::string_view access) { | |
| 489 | 3 | if (access == "read") { | |
| 490 | return MapAccess::Read; | ||
| 491 | } | ||
| 492 | 2 | if (access == "write") { | |
| 493 | return MapAccess::Write; | ||
| 494 | } | ||
| 495 | ✗ | if (access == "read_write") { | |
| 496 | return MapAccess::ReadWrite; | ||
| 497 | } | ||
| 498 | ✗ | throw nb::value_error("Access is read, write or read_write"); | |
| 499 | } | ||
| 500 | |||
| 501 | 274 | SchemaId DefinedSchema(Scene& scene, std::string_view name, const std::optional<nb::dict>& values) { | |
| 502 | 274 | const SchemaId schema = FindSchema(name); | |
| 503 | 274 | if (schema != SchemaId::Shader) { | |
| 504 | return schema; | ||
| 505 | } | ||
| 506 | 48 | const auto text = [&](const char* key) { | |
| 507 | 26 | return values && values->contains(key) ? nb::cast<std::string>((*values)[key]) | |
| 508 | 32 | : std::string(); | |
| 509 | 22 | }; | |
| 510 | 22 | if (text("implementationSource") == "sourceAsset") { | |
| 511 | 2 | const std::string node = text("slangSourceAssetSubIdentifier"); | |
| 512 | 2 | const std::string file = text("slangSourceAsset"); | |
| 513 | 2 | if (node.empty() || file.empty()) { | |
| 514 | ✗ | throw nb::value_error("A Shader loaded from a source asset needs slangSourceAsset and " | |
| 515 | ✗ | "slangSourceAssetSubIdentifier"); | |
| 516 | } | ||
| 517 | 2 | if (ShaderSchema(node) == SchemaId::None) { | |
| 518 | 2 | LoadShaders(scene, file); | |
| 519 | } | ||
| 520 | 2 | const SchemaId loaded = ShaderSchema(node); | |
| 521 | 2 | if (loaded == SchemaId::None) { | |
| 522 | ✗ | throw nb::value_error((file + " has no material node called " + node).c_str()); | |
| 523 | } | ||
| 524 | 2 | return loaded; | |
| 525 | 2 | } | |
| 526 | 20 | if (!values || !values->contains("id")) { | |
| 527 | 1 | throw nb::value_error("A Shader needs an id naming its node, such as UsdPreviewSurface"); | |
| 528 | } | ||
| 529 | 38 | const std::string id = nb::cast<std::string>((*values)["id"]); | |
| 530 | 19 | const SchemaId shader = ShaderSchema(id); | |
| 531 | 19 | if (shader == SchemaId::None) { | |
| 532 | 3 | throw nb::value_error(("No shader has info:id " + id).c_str()); | |
| 533 | } | ||
| 534 | 18 | return shader; | |
| 535 | 19 | } | |
| 536 | |||
| 537 | 1 | std::vector<Prim> Handles(std::span<const ScenePrim> prims) { | |
| 538 | 1 | std::vector<Prim> handles; | |
| 539 | 1 | handles.reserve(prims.size()); | |
| 540 | 2 | for (const ScenePrim& prim : prims) { | |
| 541 | 1 | handles.push_back(prim.prim); | |
| 542 | } | ||
| 543 | 1 | return handles; | |
| 544 | ✗ | } | |
| 545 | |||
| 546 | 4 | SceneSelection DefineBatch(Ref<Scene> scene, const std::vector<std::string>& paths, | |
| 547 | std::string_view name, std::optional<nb::dict> values, | ||
| 548 | std::optional<nb::dict> apis, bool attachable) { | ||
| 549 | 4 | const SchemaId schema = DefinedSchema(*scene, name, values); | |
| 550 | 4 | const PrimRange range = scene->Define(paths, schema, {.attachable = attachable}); | |
| 551 | 4 | SceneSelection defined{scene, scene->Select(range, schema)}; | |
| 552 | 4 | if (values) { | |
| 553 | 12 | for (auto [key, value] : *values) { | |
| 554 | 16 | Assign(defined, nb::cast<std::string>(key), value); | |
| 555 | } | ||
| 556 | } | ||
| 557 | 4 | if (apis) { | |
| 558 | ✗ | const std::vector<Prim> prims = scene->Handles(range); | |
| 559 | ✗ | for (auto [key, value] : *apis) { | |
| 560 | ✗ | ApplySchema(*scene, prims, FindSchema(nb::cast<std::string>(key)), "", | |
| 561 | ✗ | nb::cast<nb::dict>(value)); | |
| 562 | } | ||
| 563 | ✗ | } | |
| 564 | 4 | return defined; | |
| 565 | ✗ | } | |
| 566 | |||
| 567 | 270 | ScenePrim Define(Ref<Scene> scene, std::string_view path, std::string_view name, | |
| 568 | std::optional<nb::dict> values, std::optional<nb::dict> apis) { | ||
| 569 | 270 | const SchemaId schema = DefinedSchema(*scene, name, values); | |
| 570 | 268 | if (IsRuntime(schema)) { | |
| 571 | 24 | DefineBatch(scene, {std::string(path)}, name, values, apis, false); | |
| 572 | 4 | const std::optional<Prim> prim = scene->Find(path); | |
| 573 | 4 | return {std::move(scene), *prim}; | |
| 574 | } | ||
| 575 | 1047 | const Prim prim = DefineSchema(*scene, path, schema, values ? *values : nb::dict(), | |
| 576 | 264 | apis ? *apis : nb::dict()); | |
| 577 | 262 | return {std::move(scene), prim}; | |
| 578 | 274 | } | |
| 579 | |||
| 580 | 1 | Interpolation FindInterpolation(std::string_view name) { | |
| 581 | 4 | for (uint32_t mode = 0; mode < std::size(kInterpolationTokens); ++mode) { | |
| 582 | 4 | if (kInterpolationTokens[mode] == name) { | |
| 583 | 1 | return Interpolation(mode); | |
| 584 | } | ||
| 585 | } | ||
| 586 | ✗ | throw nb::value_error((std::string(name) + " isn't an interpolation").c_str()); | |
| 587 | } | ||
| 588 | |||
| 589 | 1 | Primvar CreatePrimvar(std::string name, nb::handle values, std::string_view interpolation, | |
| 590 | std::optional<nb::handle> indices, uint32_t elementSize) { | ||
| 591 | 1 | const nb::object numpy = nb::module_::import_("numpy"); | |
| 592 | 1 | nb::object array = numpy.attr("asarray")(values); | |
| 593 | 1 | const bool integral = | |
| 594 | 2 | nb::cast<bool>(numpy.attr("issubdtype")(array.attr("dtype"), numpy.attr("integer"))); | |
| 595 | 2 | const size_t ndim = nb::cast<size_t>(array.attr("ndim")); | |
| 596 | 1 | const uint32_t components = | |
| 597 | 1 | ndim > 1 ? nb::cast<uint32_t>(nb::cast<nb::tuple>(array.attr("shape"))[ndim - 1]) : 1; | |
| 598 | 1 | if (ndim > 2 || components < 1 || components > 4) { | |
| 599 | ✗ | throw nb::value_error("A primvar's values are n values or n rows of 2, 3 or 4"); | |
| 600 | } | ||
| 601 | |||
| 602 | 1 | Primvar primvar; | |
| 603 | 1 | primvar.name = std::move(name); | |
| 604 | 2 | primvar.type = | |
| 605 | 1 | PrimvarType(uint32_t(integral ? PrimvarType::Int : PrimvarType::Float) + components - 1); | |
| 606 | 1 | primvar.interpolation = FindInterpolation(interpolation); | |
| 607 | 1 | primvar.elementSize = elementSize; | |
| 608 | 3 | primvar.words = Items<uint32_t>(NumpyBytes(array, integral ? "int32" : "float32")); | |
| 609 | 1 | if (indices) { | |
| 610 | ✗ | primvar.indices = Items<int32_t>(NumpyBytes(*indices, "int32")); | |
| 611 | } | ||
| 612 | 1 | return primvar; | |
| 613 | 1 | } | |
| 614 | |||
| 615 | 2 | void RegisterPrimvar(nb::module_& module) { | |
| 616 | 2 | nb::class_<Primvar>(module, "Primvar") | |
| 617 | 6 | .def(nb::new_(&CreatePrimvar), nb::arg("name"), nb::arg("values"), | |
| 618 | 6 | nb::arg("interpolation") = "vertex", nb::arg("indices") = nb::none(), | |
| 619 | 6 | nb::arg("element_size") = 1) | |
| 620 | 2 | .def_ro("name", &Primvar::name) | |
| 621 | 2 | .def_prop_ro("interpolation", | |
| 622 | 2 | [](const Primvar& primvar) { | |
| 623 | ✗ | return std::string(kInterpolationTokens[uint32_t(primvar.interpolation)]); | |
| 624 | }) | ||
| 625 | 2 | .def_ro("element_size", &Primvar::elementSize) | |
| 626 | 2 | .def_prop_ro("values", | |
| 627 | ✗ | [](const Primvar& primvar) { | |
| 628 | ✗ | const uint32_t components = PrimvarComponents(primvar.type); | |
| 629 | ✗ | std::vector<size_t> shape = {primvar.words.size() / components}; | |
| 630 | ✗ | if (components > 1) { | |
| 631 | ✗ | shape.push_back(components); | |
| 632 | } | ||
| 633 | ✗ | const bool integral = PrimvarIsInt(primvar.type); | |
| 634 | ✗ | return Numpy(Bytes(primvar.words), shape, | |
| 635 | ✗ | {uint8_t(integral ? kDLInt : kDLFloat), 32, 1}); | |
| 636 | ✗ | }) | |
| 637 | 2 | .def_prop_ro("indices", | |
| 638 | ✗ | [](const Primvar& primvar) { | |
| 639 | ✗ | return Numpy(Bytes(primvar.indices), {primvar.indices.size()}, | |
| 640 | ✗ | {kDLInt, 32, 1}); | |
| 641 | }) | ||
| 642 | 2 | .def("__repr__", [](const Primvar& primvar) { | |
| 643 | ✗ | return "Primvar(" + primvar.name + ", " + | |
| 644 | ✗ | std::string(kInterpolationTokens[uint32_t(primvar.interpolation)]) + ")"; | |
| 645 | }); | ||
| 646 | 2 | } | |
| 647 | |||
| 648 | 2 | void RegisterScene(nb::module_& module) { | |
| 649 | 2 | nb::class_<Scene>(module, "Scene") | |
| 650 | 6 | .def(nb::new_([](std::optional<Ref<device::Device>> device, device::Traversal traversal, | |
| 651 | uint32_t maxRefits, uint32_t pointsPerBox, float curveTolerance, | ||
| 652 | uint32_t maxCurveSteps, bool exactCurves, bool waitForKernels) { | ||
| 653 | 42 | nb::gil_scoped_release release; | |
| 654 | 42 | const SceneDesc desc{.traversal = traversal, | |
| 655 | .maxRefits = maxRefits, | ||
| 656 | .pointsPerBox = pointsPerBox, | ||
| 657 | .curveTolerance = curveTolerance, | ||
| 658 | .maxCurveSteps = maxCurveSteps, | ||
| 659 | .exactCurves = exactCurves, | ||
| 660 | 42 | .waitForKernels = waitForKernels}; | |
| 661 | 42 | return device ? Scene::Create(**device, desc) : Scene::Create(desc); | |
| 662 | 42 | }), | |
| 663 | 4 | nb::arg("device") = nb::none(), nb::kw_only(), | |
| 664 | 4 | nb::arg("traversal") = device::Traversal::Auto, | |
| 665 | 6 | nb::arg("max_refits") = SceneDesc{}.maxRefits, | |
| 666 | 6 | nb::arg("points_per_box") = SceneDesc{}.pointsPerBox, | |
| 667 | 6 | nb::arg("curve_tolerance") = SceneDesc{}.curveTolerance, | |
| 668 | 6 | nb::arg("max_curve_steps") = SceneDesc{}.maxCurveSteps, | |
| 669 | 8 | nb::arg("exact_curves") = SceneDesc{}.exactCurves, nb::arg("wait_for_kernels") = true) | |
| 670 | 4 | .def( | |
| 671 | "load_shaders", | ||
| 672 | 6 | [](Ref<Scene> scene, const std::string& path) { return LoadShaders(*scene, path); }, | |
| 673 | 2 | nb::arg("path")) | |
| 674 | 5 | .def("reload_shaders", [](Ref<Scene> scene) { return ReloadShaders(*scene); }) | |
| 675 | 4 | .def("define", &Define, nb::arg("path"), nb::arg("schema"), nb::arg("values") = nb::none(), | |
| 676 | 4 | nb::kw_only(), nb::arg("apis") = nb::none()) | |
| 677 | 4 | .def("define", &DefineBatch, nb::arg("paths"), nb::arg("schema"), | |
| 678 | 4 | nb::arg("values") = nb::none(), nb::kw_only(), nb::arg("apis") = nb::none(), | |
| 679 | 6 | nb::arg("attachable") = false) | |
| 680 | 2 | .def( | |
| 681 | "apply", | ||
| 682 | 1 | [](Scene& scene, const std::vector<ScenePrim>& prims, std::string_view api, | |
| 683 | std::optional<nb::dict> values, std::optional<std::string_view> instance) { | ||
| 684 | 1 | const SchemaId schema = FindSchema(api); | |
| 685 | 1 | if (IsMultipleApi(schema) && !instance) { | |
| 686 | ✗ | throw nb::value_error((std::string(api) + " needs an instance name").c_str()); | |
| 687 | } | ||
| 688 | 1 | if (!IsMultipleApi(schema) && instance) { | |
| 689 | ✗ | throw nb::value_error((std::string(api) + " takes no instance name").c_str()); | |
| 690 | } | ||
| 691 | 2 | ApplySchema(scene, Handles(prims), schema, instance.value_or(""), | |
| 692 | 1 | values ? *values : nb::dict()); | |
| 693 | 1 | }, | |
| 694 | 4 | nb::arg("prims"), nb::arg("api"), nb::arg("values") = nb::none(), nb::kw_only(), | |
| 695 | 4 | nb::arg("instance") = nb::none()) | |
| 696 | 4 | .def( | |
| 697 | "unapply", | ||
| 698 | ✗ | [](Scene& scene, const std::vector<ScenePrim>& prims, std::string_view api, | |
| 699 | std::optional<std::string_view> instance) { | ||
| 700 | ✗ | if (instance) { | |
| 701 | ✗ | scene.Unapply(Handles(prims), FindSchema(api), *instance); | |
| 702 | } else { | ||
| 703 | ✗ | scene.Unapply(Handles(prims), FindSchema(api)); | |
| 704 | } | ||
| 705 | ✗ | }, | |
| 706 | 4 | nb::arg("prims"), nb::arg("api"), nb::kw_only(), nb::arg("instance") = nb::none()) | |
| 707 | 4 | .def( | |
| 708 | 3 | "remove", [](Scene& scene, const ScenePrim& prim) { scene.Remove(prim.prim); }, | |
| 709 | 2 | nb::arg("prim")) | |
| 710 | 4 | .def( | |
| 711 | "move", | ||
| 712 | 1 | [](Ref<Scene> scene, const ScenePrim& prim, std::string_view path) { | |
| 713 | 1 | return ScenePrim{scene, scene->Move(prim.prim, path)}; | |
| 714 | }, | ||
| 715 | 2 | nb::arg("prim"), nb::arg("path")) | |
| 716 | 4 | .def( | |
| 717 | "rename", | ||
| 718 | 1 | [](Ref<Scene> scene, const ScenePrim& prim, std::string_view name) { | |
| 719 | 1 | return ScenePrim{scene, scene->Rename(prim.prim, name)}; | |
| 720 | }, | ||
| 721 | 2 | nb::arg("prim"), nb::arg("name")) | |
| 722 | 4 | .def( | |
| 723 | 1 | "copy", [](Scene& scene, const ScenePrim& prim) { return scene.Copy(prim.prim); }, | |
| 724 | 2 | nb::arg("prim")) | |
| 725 | 4 | .def( | |
| 726 | "paste", | ||
| 727 | 1 | [](Ref<Scene> scene, const PrimCopy& copy, std::string_view path) { | |
| 728 | 1 | return ScenePrim{scene, scene->Paste(copy, path)}; | |
| 729 | }, | ||
| 730 | 2 | nb::arg("copy"), nb::arg("path")) | |
| 731 | 4 | .def( | |
| 732 | "find", | ||
| 733 | 41262 | [](Ref<Scene> scene, std::string_view path) -> std::optional<ScenePrim> { | |
| 734 | 41262 | if (const std::optional<Prim> prim = scene->Find(path)) { | |
| 735 | 41259 | return ScenePrim{scene, *prim}; | |
| 736 | } | ||
| 737 | 3 | return std::nullopt; | |
| 738 | }, | ||
| 739 | 2 | nb::arg("path")) | |
| 740 | 2 | .def( | |
| 741 | "select", | ||
| 742 | 23 | [](Ref<Scene> scene, std::string_view pattern, std::optional<std::string_view> schema, | |
| 743 | std::optional<std::string_view> instance) { | ||
| 744 | 23 | View view = instance | |
| 745 | 23 | ? scene->Select(pattern, FindSchema(schema.value_or("")), *instance) | |
| 746 | 23 | : schema ? scene->Select(pattern, FindSchema(*schema)) | |
| 747 | 23 | : scene->Select(pattern); | |
| 748 | 23 | return SceneSelection{std::move(scene), std::move(view)}; | |
| 749 | 23 | }, | |
| 750 | 4 | nb::arg("pattern"), nb::arg("schema") = nb::none(), nb::kw_only(), | |
| 751 | 4 | nb::arg("instance") = nb::none()) | |
| 752 | 2 | .def( | |
| 753 | "select", | ||
| 754 | 6 | [](Ref<Scene> scene, const std::vector<std::string>& paths) { | |
| 755 | 6 | View view = scene->Select(paths); | |
| 756 | 6 | return SceneSelection{std::move(scene), std::move(view)}; | |
| 757 | 6 | }, | |
| 758 | ✗ | nb::arg("paths")) | |
| 759 | 2 | .def( | |
| 760 | "clone", | ||
| 761 | ✗ | [](Ref<Scene> scene, const ScenePrim& source, const std::vector<std::string>& targets, | |
| 762 | bool moveSource, bool attachable) { | ||
| 763 | ✗ | scene->Clone(source.prim, targets, | |
| 764 | {.moveSource = moveSource, .attachable = attachable}); | ||
| 765 | ✗ | std::vector<ScenePrim> roots; | |
| 766 | ✗ | for (const std::string& target : targets) { | |
| 767 | ✗ | roots.push_back({scene, *scene->Find(target)}); | |
| 768 | } | ||
| 769 | ✗ | return roots; | |
| 770 | ✗ | }, | |
| 771 | 4 | nb::arg("source"), nb::arg("targets"), nb::kw_only(), nb::arg("move_source") = false, | |
| 772 | 6 | nb::arg("attachable") = false) | |
| 773 | 2 | .def( | |
| 774 | "clone", | ||
| 775 | ✗ | [](Ref<Scene> scene, const ScenePrim& source, uint32_t count, bool moveSource, | |
| 776 | bool attachable) { | ||
| 777 | ✗ | const PrimRange copies = scene->Clone( | |
| 778 | source.prim, count, {.moveSource = moveSource, .attachable = attachable}); | ||
| 779 | ✗ | std::vector<ScenePrim> roots; | |
| 780 | ✗ | for (uint32_t copy = 0; copy < copies.copies; ++copy) { | |
| 781 | ✗ | roots.push_back({scene, scene->ByIndex(copies.Index(copy, 0))}); | |
| 782 | } | ||
| 783 | ✗ | return roots; | |
| 784 | ✗ | }, | |
| 785 | 4 | nb::arg("source"), nb::arg("count"), nb::kw_only(), nb::arg("move_source") = false, | |
| 786 | 6 | nb::arg("attachable") = false) | |
| 787 | 2 | .def( | |
| 788 | "publish", | ||
| 789 | 3 | [](Scene& scene, std::optional<Ordinal> ordinal, std::optional<uintptr_t> stream, | |
| 790 | 1 | bool wait) { scene.Publish({.ordinal = ordinal, .wait = wait}, ToStream(stream)); }, | |
| 791 | 4 | nb::arg("ordinal") = nb::none(), nb::arg("stream") = nb::none(), nb::kw_only(), | |
| 792 | 6 | nb::arg("wait") = false, nb::call_guard<nb::gil_scoped_release>()) | |
| 793 | 2 | .def("advance", &Scene::Advance, nb::arg("dt")) | |
| 794 | 2 | .def_prop_ro("published", &Scene::Published) | |
| 795 | 2 | .def_prop_rw("time", &Scene::Time, &Scene::SetTime) | |
| 796 | 4 | .def_prop_ro("structure_version", &Scene::StructureVersion) | |
| 797 | 2 | .def( | |
| 798 | "acquire", | ||
| 799 | 2 | [](Scene& scene, std::optional<uintptr_t> stream) { | |
| 800 | ✗ | return scene.Acquire(stream ? ToStream(stream) : Stream::Host()); | |
| 801 | }, | ||
| 802 | 4 | nb::arg("stream") = nb::none(), nb::call_guard<nb::gil_scoped_release>()) | |
| 803 | 4 | .def( | |
| 804 | "release", | ||
| 805 | ✗ | [](Scene& scene, const Snapshot& snapshot, std::optional<uintptr_t> stream) { | |
| 806 | ✗ | scene.Release(snapshot, ToStream(stream)); | |
| 807 | ✗ | }, | |
| 808 | 4 | nb::arg("snapshot"), nb::arg("stream") = nb::none(), | |
| 809 | nb::call_guard<nb::gil_scoped_release>()) | ||
| 810 | 2 | .def( | |
| 811 | "bounds", | ||
| 812 | 2 | [](const Scene& scene, std::optional<uint32_t> world) -> nb::object { | |
| 813 | ✗ | const std::optional<Box> box = scene.Bounds(world.value_or(kGlobalWorld)); | |
| 814 | ✗ | if (!box) { | |
| 815 | ✗ | return nb::none(); | |
| 816 | } | ||
| 817 | ✗ | return nb::make_tuple(nb::make_tuple(box->lower.x, box->lower.y, box->lower.z), | |
| 818 | ✗ | nb::make_tuple(box->upper.x, box->upper.y, box->upper.z)); | |
| 819 | }, | ||
| 820 | 4 | nb::arg("world") = nb::none()) | |
| 821 | 2 | .def( | |
| 822 | "prim_bounds", | ||
| 823 | 1 | [](const Scene& scene, const std::vector<std::string>& paths) { | |
| 824 | 1 | std::vector<Box> boxes; | |
| 825 | 1 | { | |
| 826 | 1 | const nb::gil_scoped_release release; | |
| 827 | 2 | boxes = scene.PrimBounds(Span<const std::string>(paths)); | |
| 828 | ✗ | } | |
| 829 | 1 | nb::list rows; | |
| 830 | 3 | for (const Box& box : boxes) { | |
| 831 | 2 | rows.append( | |
| 832 | 4 | nb::make_tuple(nb::make_tuple(box.lower.x, box.lower.y, box.lower.z), | |
| 833 | 4 | nb::make_tuple(box.upper.x, box.upper.y, box.upper.z))); | |
| 834 | } | ||
| 835 | 4 | return nb::module_::import_("numpy").attr("asarray")(rows, "dtype"_a = "float32"); | |
| 836 | 1 | }, | |
| 837 | 2 | nb::arg("paths"), | |
| 838 | "World-space boxes of the prims at paths, relative to each prim's world origin, as an " | ||
| 839 | "(N, 2, 3) array of lower and upper corners as of the latest publish. Prims without " | ||
| 840 | "bounds read as zeros.") | ||
| 841 | 2 | .def_prop_ro("pending", &Scene::Pending) | |
| 842 | 2 | .def_prop_ro("world_count", &Scene::WorldCount) | |
| 843 | 4 | .def_prop_ro("prim_count", &Scene::PrimCount) | |
| 844 | 2 | .def_prop_ro("device", | |
| 845 | ✗ | [](Scene& scene) { return Ref<device::Device>(&scene.GetDevice()); }); | |
| 846 | 2 | } | |
| 847 | |||
| 848 | 2 | void RegisterPrimCopy(nb::module_& module) { | |
| 849 | 2 | nb::class_<PrimCopy>(module, "PrimCopy") | |
| 850 | 4 | .def_prop_ro("empty", &PrimCopy::Empty) | |
| 851 | ✗ | .def_prop_ro("name", [](const PrimCopy& copy) { return std::string(copy.Name()); }) | |
| 852 | 4 | .def_prop_ro("prim_count", &PrimCopy::PrimCount); | |
| 853 | 2 | } | |
| 854 | |||
| 855 | 2 | void RegisterPrim(nb::module_& module) { | |
| 856 | 2 | nb::class_<ScenePrim>(module, "Prim") | |
| 857 | 38 | .def_prop_ro("path", [](const ScenePrim& prim) { return prim.prim.Path(); }) | |
| 858 | 3 | .def_prop_ro("name", [](const ScenePrim& prim) { return std::string(prim.prim.Name()); }) | |
| 859 | 2 | .def_prop_ro("index", [](const ScenePrim& prim) { return prim.prim.Index(); }) | |
| 860 | 2 | .def_prop_ro("world", | |
| 861 | ✗ | [](const ScenePrim& prim) -> std::optional<uint32_t> { | |
| 862 | ✗ | const uint32_t world = prim.prim.World(); | |
| 863 | ✗ | if (world == kGlobalWorld) { | |
| 864 | ✗ | return std::nullopt; | |
| 865 | } | ||
| 866 | ✗ | return world; | |
| 867 | }) | ||
| 868 | 2 | .def_prop_ro("generation", [](const ScenePrim& prim) { return prim.prim.Generation(); }) | |
| 869 | 2 | .def_prop_ro("valid", [](const ScenePrim& prim) { return prim.prim.Valid(); }) | |
| 870 | 2 | .def_prop_ro( | |
| 871 | "schema", | ||
| 872 | ✗ | [](const ScenePrim& prim) { return std::string(Info(prim.prim.Schema()).name); }) | |
| 873 | 19 | .def_prop_ro("scene", [](const ScenePrim& prim) { return prim.scene; }) | |
| 874 | 2 | .def_prop_ro("parent", | |
| 875 | ✗ | [](const ScenePrim& prim) -> std::optional<ScenePrim> { | |
| 876 | ✗ | if (const std::optional<Prim> parent = prim.prim.Parent()) { | |
| 877 | ✗ | return ScenePrim{prim.scene, *parent}; | |
| 878 | } | ||
| 879 | ✗ | return std::nullopt; | |
| 880 | }) | ||
| 881 | 2 | .def_prop_ro( | |
| 882 | "children", | ||
| 883 | 5 | [](const ScenePrim& prim) { return ScenePrims(prim.scene, prim.prim.Children()); }) | |
| 884 | 2 | .def_prop_ro("connections", | |
| 885 | 35 | [](const ScenePrim& prim) { | |
| 886 | 35 | std::vector<std::tuple<std::string, std::string, std::string>> found; | |
| 887 | 64 | for (const Connection& connection : prim.prim.Connections()) { | |
| 888 | 29 | found.emplace_back(connection.input, connection.source, | |
| 889 | 29 | connection.output); | |
| 890 | 35 | } | |
| 891 | 35 | return found; | |
| 892 | ✗ | }) | |
| 893 | 4 | .def( | |
| 894 | "connect", | ||
| 895 | 21 | [](const ScenePrim& prim, std::string_view input, const ScenePrim& source, | |
| 896 | 21 | std::string_view output) { prim.prim.Connect(input, source.prim, output); }, | |
| 897 | 2 | nb::arg("input"), nb::arg("source"), nb::arg("output")) | |
| 898 | 4 | .def( | |
| 899 | "disconnect", | ||
| 900 | 1 | [](const ScenePrim& prim, std::string_view input) { prim.prim.Disconnect(input); }, | |
| 901 | 2 | nb::arg("input")) | |
| 902 | 4 | .def( | |
| 903 | "is_a", | ||
| 904 | 229 | [](const ScenePrim& prim, std::string_view schema) { | |
| 905 | 229 | return prim.prim.IsA(FindSchema(schema)); | |
| 906 | }, | ||
| 907 | 2 | nb::arg("schema")) | |
| 908 | 4 | .def( | |
| 909 | "has_api", | ||
| 910 | ✗ | [](const ScenePrim& prim, std::string_view api) { | |
| 911 | ✗ | return prim.prim.HasApi(FindSchema(api)); | |
| 912 | }, | ||
| 913 | 2 | nb::arg("api")) | |
| 914 | 4 | .def( | |
| 915 | "attach", | ||
| 916 | ✗ | [](const ScenePrim& prim, std::string_view field, nb::handle tensor, | |
| 917 | std::string_view initial, std::optional<uintptr_t> stream) { | ||
| 918 | ✗ | Attach(Single(prim), field, tensor, initial, stream); | |
| 919 | ✗ | }, | |
| 920 | 6 | nb::arg("field"), nb::arg("tensor"), nb::kw_only(), nb::arg("initial") = "tensor", | |
| 921 | 4 | nb::arg("stream") = nb::none()) | |
| 922 | 4 | .def( | |
| 923 | "mark_changed", | ||
| 924 | ✗ | [](const ScenePrim& prim, std::string_view field, std::optional<uintptr_t> stream) { | |
| 925 | ✗ | MarkChanged(Single(prim), field, stream, nb::none()); | |
| 926 | ✗ | }, | |
| 927 | 4 | nb::arg("field"), nb::arg("stream") = nb::none()) | |
| 928 | 4 | .def( | |
| 929 | "detach", | ||
| 930 | ✗ | [](const ScenePrim& prim, std::string_view field, std::optional<uintptr_t> stream) { | |
| 931 | ✗ | Detach(Single(prim), field, stream); | |
| 932 | ✗ | }, | |
| 933 | 4 | nb::arg("field"), nb::arg("stream") = nb::none()) | |
| 934 | 4 | .def( | |
| 935 | "wait_for_reads", | ||
| 936 | ✗ | [](const ScenePrim& prim, std::string_view field, std::optional<uintptr_t> stream) { | |
| 937 | ✗ | WaitForReads(Single(prim), field, stream); | |
| 938 | ✗ | }, | |
| 939 | 4 | nb::arg("field"), nb::arg("stream") = nb::none()) | |
| 940 | 4 | .def( | |
| 941 | "attached", | ||
| 942 | ✗ | [](const ScenePrim& prim, std::string_view field) { | |
| 943 | ✗ | return Single(prim).view.Get(field).Attached(); | |
| 944 | }, | ||
| 945 | 2 | nb::arg("field")) | |
| 946 | 2 | .def("__getitem__", | |
| 947 | 51 | [](ScenePrim& prim, std::string_view field) { | |
| 948 | 51 | SceneSelection single = Single(prim); | |
| 949 | 102 | return ReadFirst(single, field); | |
| 950 | 51 | }) | |
| 951 | 2 | .def("__setitem__", [](ScenePrim& prim, std::string_view field, | |
| 952 | 41063 | nb::handle value) { AssignOne(prim, field, value, std::nullopt); }) | |
| 953 | 4 | .def( | |
| 954 | "instances", | ||
| 955 | ✗ | [](const ScenePrim& prim, std::string_view api) { | |
| 956 | ✗ | return prim.prim.Instances(FindSchema(api)); | |
| 957 | }, | ||
| 958 | 2 | nb::arg("api")) | |
| 959 | 4 | .def( | |
| 960 | "get_instance", | ||
| 961 | ✗ | [](ScenePrim& prim, std::string_view api, std::string_view instance, | |
| 962 | std::string_view field) { | ||
| 963 | ✗ | SceneSelection single{prim.scene, | |
| 964 | ✗ | prim.scene->Select(prim.prim, FindSchema(api), instance)}; | |
| 965 | ✗ | return ReadFirst(single, field); | |
| 966 | ✗ | }, | |
| 967 | 2 | nb::arg("api"), nb::arg("instance"), nb::arg("field")) | |
| 968 | 4 | .def( | |
| 969 | "set_instance", | ||
| 970 | ✗ | [](ScenePrim& prim, std::string_view api, std::string_view instance, | |
| 971 | std::string_view field, nb::handle value) { | ||
| 972 | ✗ | SceneSelection single{prim.scene, | |
| 973 | ✗ | prim.scene->Select(prim.prim, FindSchema(api), instance)}; | |
| 974 | ✗ | Assign(single, field, value); | |
| 975 | ✗ | }, | |
| 976 | 2 | nb::arg("api"), nb::arg("instance"), nb::arg("field"), nb::arg("value")) | |
| 977 | 4 | .def( | |
| 978 | "instance_values", | ||
| 979 | ✗ | [](ScenePrim& prim, std::string_view api, std::string_view field) { | |
| 980 | ✗ | const SchemaId schema = FindSchema(api); | |
| 981 | ✗ | nb::dict values; | |
| 982 | ✗ | for (const std::string& name : prim.prim.Instances(schema)) { | |
| 983 | ✗ | SceneSelection single{prim.scene, prim.scene->Select(prim.prim, schema, name)}; | |
| 984 | ✗ | values[nb::str(name.c_str())] = ReadFirst(single, field); | |
| 985 | ✗ | } | |
| 986 | ✗ | return values; | |
| 987 | ✗ | }, | |
| 988 | 2 | nb::arg("api"), nb::arg("field")) | |
| 989 | 4 | .def( | |
| 990 | "set_instances", | ||
| 991 | ✗ | [](ScenePrim& prim, std::string_view api, std::string_view field, | |
| 992 | const nb::dict& values) { | ||
| 993 | ✗ | const SchemaId schema = FindSchema(api); | |
| 994 | ✗ | for (const std::string& name : prim.prim.Instances(schema)) { | |
| 995 | ✗ | if (!values.contains(name.c_str())) { | |
| 996 | ✗ | prim.scene->Unapply(Span<const Prim>(&prim.prim, 1), schema, name); | |
| 997 | } | ||
| 998 | ✗ | } | |
| 999 | ✗ | for (auto [key, value] : values) { | |
| 1000 | ✗ | const std::string name = nb::cast<std::string>(key); | |
| 1001 | ✗ | prim.scene->Apply(Span<const Prim>(&prim.prim, 1), schema, name); | |
| 1002 | ✗ | SceneSelection single{prim.scene, prim.scene->Select(prim.prim, schema, name)}; | |
| 1003 | ✗ | Assign(single, field, value); | |
| 1004 | ✗ | } | |
| 1005 | ✗ | }, | |
| 1006 | 2 | nb::arg("api"), nb::arg("field"), nb::arg("values")) | |
| 1007 | 2 | .def("__eq__", | |
| 1008 | ✗ | [](const ScenePrim& prim, const ScenePrim& other) { | |
| 1009 | ✗ | return prim.scene == other.scene && prim.prim == other.prim; | |
| 1010 | }) | ||
| 1011 | 2 | .def("__hash__", | |
| 1012 | ✗ | [](const ScenePrim& prim) { | |
| 1013 | ✗ | return (size_t(prim.prim.Index()) << 32) ^ prim.prim.Generation() ^ | |
| 1014 | ✗ | size_t(prim.scene.Get()); | |
| 1015 | }) | ||
| 1016 | 2 | .def("__repr__", [](const ScenePrim& prim) { | |
| 1017 | ✗ | return "Prim(" + prim.prim.Path() + ", " + std::string(Info(prim.prim.Schema()).name) + | |
| 1018 | ✗ | ")"; | |
| 1019 | }); | ||
| 1020 | 2 | } | |
| 1021 | |||
| 1022 | 2 | void RegisterMapping(nb::module_& module) { | |
| 1023 | 2 | nb::class_<SceneMapping>(module, "Mapping") | |
| 1024 | 2 | .def_prop_ro("open", [](const SceneMapping& held) { return held.mapping.Open(); }) | |
| 1025 | 2 | .def_prop_ro("aliased", [](const SceneMapping& held) { return held.mapping.Aliased(); }) | |
| 1026 | 6 | .def_prop_ro("values", [](const SceneMapping& held) { return held.mapping.Values(); }) | |
| 1027 | 2 | .def_prop_ro("rows", | |
| 1028 | ✗ | [](const SceneMapping& held) -> std::optional<device::Tensor> { | |
| 1029 | ✗ | if (held.mapping.Rows().Empty()) { | |
| 1030 | ✗ | return std::nullopt; | |
| 1031 | } | ||
| 1032 | ✗ | return held.mapping.Rows(); | |
| 1033 | }) | ||
| 1034 | 2 | .def_prop_ro("offsets", | |
| 1035 | ✗ | [](const SceneMapping& held) -> std::optional<device::Tensor> { | |
| 1036 | ✗ | if (held.mapping.Offsets().Empty()) { | |
| 1037 | ✗ | return std::nullopt; | |
| 1038 | } | ||
| 1039 | ✗ | return held.mapping.Offsets(); | |
| 1040 | }) | ||
| 1041 | 2 | .def( | |
| 1042 | "close", | ||
| 1043 | ✗ | [](SceneMapping& held, std::optional<uintptr_t> stream, nb::handle mask) { | |
| 1044 | ✗ | if (!held.mapping.Open()) { | |
| 1045 | ✗ | return; | |
| 1046 | } | ||
| 1047 | ✗ | HeldTensors kept; | |
| 1048 | ✗ | const DLTensor* bits = kept.Take(MaskWords(mask), stream); | |
| 1049 | ✗ | nb::gil_scoped_release release; | |
| 1050 | ✗ | kept.Retire(*held.scene, | |
| 1051 | ✗ | held.mapping.Close(ToStream(stream ? stream : held.stream), bits)); | |
| 1052 | ✗ | }, | |
| 1053 | 4 | nb::kw_only(), nb::arg("stream") = nb::none(), nb::arg("mask") = nb::none()) | |
| 1054 | 2 | .def( | |
| 1055 | "__enter__", [](SceneMapping& held) -> SceneMapping& { return held; }, | ||
| 1056 | nb::rv_policy::reference) | ||
| 1057 | 2 | .def( | |
| 1058 | "__exit__", | ||
| 1059 | 2 | [](SceneMapping& held, nb::handle, nb::handle, nb::handle) { | |
| 1060 | 2 | if (held.mapping.Open()) { | |
| 1061 | 2 | nb::gil_scoped_release release; | |
| 1062 | 2 | held.mapping.Close(ToStream(held.stream)); | |
| 1063 | 2 | } | |
| 1064 | 2 | }, | |
| 1065 | 4 | nb::arg("type").none(), nb::arg("value").none(), nb::arg("traceback").none()); | |
| 1066 | 2 | } | |
| 1067 | |||
| 1068 | 2 | void RegisterSnapshot(nb::module_& module) { | |
| 1069 | 2 | nb::class_<Snapshot>(module, "Snapshot") | |
| 1070 | 2 | .def_ro("ordinal", &Snapshot::ordinal) | |
| 1071 | 2 | .def_ro("tables", &Snapshot::tables) | |
| 1072 | 2 | .def_prop_ro("ready", [](const Snapshot& snapshot) { return snapshot.ready.Ready(); }) | |
| 1073 | 2 | .def_prop_ro("version", [](const Snapshot& snapshot) { return snapshot.Version(); }) | |
| 1074 | 2 | .def( | |
| 1075 | "table_version", | ||
| 1076 | ✗ | [](const Snapshot& snapshot, std::string_view table) { | |
| 1077 | ✗ | return snapshot.Version(table); | |
| 1078 | }, | ||
| 1079 | 2 | nb::arg("table")); | |
| 1080 | 2 | } | |
| 1081 | |||
| 1082 | 1138 | void WritePosesFrom(const SceneSelection& target, nb::handle tensor, nb::handle translation, | |
| 1083 | nb::handle rotation, std::string_view order, nb::handle rows, nb::handle mask, | ||
| 1084 | std::optional<uintptr_t> stream, bool world) { | ||
| 1085 | 1138 | const bool split = !translation.is_none() || !rotation.is_none(); | |
| 1086 | 1138 | if (split == !tensor.is_none() || (split && (translation.is_none() || rotation.is_none()))) { | |
| 1087 | ✗ | throw nb::value_error("Give poses as one [..., 7] tensor, or as a translation " | |
| 1088 | ✗ | "and a rotation"); | |
| 1089 | } | ||
| 1090 | 1138 | const QuaternionOrder chosen = FindOrder(order); | |
| 1091 | 1138 | const nb::object maskWords = MaskWords(mask); | |
| 1092 | 1138 | HeldTensors held; | |
| 1093 | 1138 | const bool native = | |
| 1094 | 1138 | nb::isinstance<device::Tensor>(tensor) || nb::isinstance<device::Tensor>(translation) || | |
| 1095 | 3414 | nb::isinstance<device::Tensor>(rotation) || nb::isinstance<device::Tensor>(rows); | |
| 1096 | 1138 | if (!native) { | |
| 1097 | 1138 | const DLTensor* packed = held.Take(tensor, stream); | |
| 1098 | 1138 | const DLTensor* translations = held.Take(translation, stream); | |
| 1099 | 1138 | const DLTensor* rotations = held.Take(rotation, stream); | |
| 1100 | 1138 | const DLTensor* sources = held.Take(rows, stream); | |
| 1101 | 1138 | const DLTensor* bits = held.Take(maskWords, stream); | |
| 1102 | 1138 | const Stream producer = | |
| 1103 | 1138 | held.Producer().IsNone() && stream ? ToStream(stream) : held.Producer(); | |
| 1104 | 1138 | nb::gil_scoped_release release; | |
| 1105 | 1138 | if (world) { | |
| 1106 | 1 | held.Retire( | |
| 1107 | *target.scene, | ||
| 1108 | 1 | split ? target.view.WriteWorldPoses(*translations, *rotations, producer, chosen, | |
| 1109 | sources, bits) | ||
| 1110 | 1 | : target.view.WriteWorldPoses(*packed, producer, chosen, sources, bits)); | |
| 1111 | 1 | return; | |
| 1112 | } | ||
| 1113 | 1137 | held.Retire(*target.scene, | |
| 1114 | 1137 | split ? target.view.WritePoses(*translations, *rotations, producer, chosen, | |
| 1115 | sources, bits) | ||
| 1116 | 1136 | : target.view.WritePoses(*packed, producer, chosen, sources, bits)); | |
| 1117 | 1137 | return; | |
| 1118 | 1138 | } | |
| 1119 | ✗ | device::Device& device = target.scene->GetDevice(); | |
| 1120 | ✗ | const auto input = [&](nb::handle value, const char* label) { | |
| 1121 | ✗ | if (value.is_none()) { | |
| 1122 | ✗ | return device::Tensor(); | |
| 1123 | } | ||
| 1124 | ✗ | if (nb::isinstance<device::Tensor>(value)) { | |
| 1125 | ✗ | return nb::cast<device::Tensor>(value); | |
| 1126 | } | ||
| 1127 | ✗ | return held.Import(device, value, stream, label); | |
| 1128 | ✗ | }; | |
| 1129 | ✗ | nb::gil_scoped_release release; | |
| 1130 | ✗ | const auto lock = SceneImpl::Get(*target.scene).Lock(); | |
| 1131 | ✗ | const device::Tensor packed = input(tensor, "pose staging"); | |
| 1132 | ✗ | const device::Tensor translations = input(translation, "translation staging"); | |
| 1133 | ✗ | const device::Tensor rotations = input(rotation, "rotation staging"); | |
| 1134 | ✗ | const device::Tensor sources = input(rows, "pose rows"); | |
| 1135 | ✗ | const device::Tensor bits = input(maskWords, "pose mask"); | |
| 1136 | ✗ | const Stream producer = held.Producer().IsNone() && stream ? ToStream(stream) : held.Producer(); | |
| 1137 | ✗ | if (world) { | |
| 1138 | ✗ | held.Retire(*target.scene, | |
| 1139 | ✗ | split ? target.view.WriteWorldPoses(translations, rotations, producer, chosen, | |
| 1140 | sources, bits) | ||
| 1141 | ✗ | : target.view.WriteWorldPoses(packed, producer, chosen, sources, bits)); | |
| 1142 | ✗ | return; | |
| 1143 | } | ||
| 1144 | ✗ | held.Retire( | |
| 1145 | *target.scene, | ||
| 1146 | ✗ | split ? target.view.WritePoses(translations, rotations, producer, chosen, sources, bits) | |
| 1147 | ✗ | : target.view.WritePoses(packed, producer, chosen, sources, bits)); | |
| 1148 | 1138 | } | |
| 1149 | |||
| 1150 | 2 | void RegisterSelection(nb::module_& module) { | |
| 1151 | 2 | nb::class_<SceneSelection>(module, "Selection") | |
| 1152 | 5 | .def("__len__", [](const SceneSelection& source) { return source.view.PrimCount(); }) | |
| 1153 | 2 | .def_prop_ro( | |
| 1154 | "prims", | ||
| 1155 | 2 | [](const SceneSelection& source) { return ScenePrims(source.scene, source.Prims()); }) | |
| 1156 | 2 | .def_prop_ro("paths", | |
| 1157 | 2 | [](const SceneSelection& source) { | |
| 1158 | 2 | std::vector<std::string> paths; | |
| 1159 | 6 | for (const Prim& prim : source.Prims()) { | |
| 1160 | 8 | paths.push_back(prim.Path()); | |
| 1161 | 2 | } | |
| 1162 | 2 | return paths; | |
| 1163 | ✗ | }) | |
| 1164 | 2 | .def_prop_ro("dense", | |
| 1165 | ✗ | [](const SceneSelection& source) { | |
| 1166 | ✗ | return source.view.Schema() != SchemaId::None && | |
| 1167 | ✗ | source.view.Dense(source.view.Schema()); | |
| 1168 | }) | ||
| 1169 | 4 | .def( | |
| 1170 | "world", | ||
| 1171 | 2 | [](const SceneSelection& source, uint32_t world) { | |
| 1172 | 2 | return SceneSelection{source.scene, source.view.World(world)}; | |
| 1173 | }, | ||
| 1174 | 2 | nb::arg("world")) | |
| 1175 | 4 | .def("attach", &Attach, nb::arg("field"), nb::arg("tensor"), nb::kw_only(), | |
| 1176 | 6 | nb::arg("initial") = "tensor", nb::arg("stream") = nb::none()) | |
| 1177 | 6 | .def("mark_changed", &MarkChanged, nb::arg("field"), nb::arg("stream") = nb::none(), | |
| 1178 | 4 | nb::kw_only(), nb::arg("mask") = nb::none()) | |
| 1179 | 6 | .def("detach", &Detach, nb::arg("field"), nb::arg("stream") = nb::none()) | |
| 1180 | 6 | .def("wait_for_reads", &WaitForReads, nb::arg("field"), nb::arg("stream") = nb::none()) | |
| 1181 | 4 | .def( | |
| 1182 | "attached", | ||
| 1183 | ✗ | [](const SceneSelection& source, std::string_view field) { | |
| 1184 | ✗ | return source.view.Get(field).Attached(); | |
| 1185 | }, | ||
| 1186 | 2 | nb::arg("field")) | |
| 1187 | 4 | .def( | |
| 1188 | "read", | ||
| 1189 | ✗ | [](const SceneSelection& source, std::string_view field, | |
| 1190 | std::optional<uintptr_t> stream) { | ||
| 1191 | ✗ | const UntypedAttribute origin = source.view.Get(field); | |
| 1192 | ✗ | nb::gil_scoped_release release; | |
| 1193 | ✗ | return origin.Read(stream ? ToStream(stream) : Stream::Host()); | |
| 1194 | ✗ | }, | |
| 1195 | 4 | nb::arg("field"), nb::arg("stream") = nb::none()) | |
| 1196 | 2 | .def( | |
| 1197 | "write", | ||
| 1198 | ✗ | [](SceneSelection& source, std::string_view field, nb::handle value, | |
| 1199 | std::optional<uintptr_t> stream, | ||
| 1200 | ✗ | nb::handle mask) { Assign(source, field, value, stream, mask); }, | |
| 1201 | 4 | nb::arg("field"), nb::arg("value"), nb::kw_only(), nb::arg("stream") = nb::none(), | |
| 1202 | 4 | nb::arg("mask") = nb::none()) | |
| 1203 | 4 | .def( | |
| 1204 | "map", | ||
| 1205 | ✗ | [](const SceneSelection& source, std::string_view field, std::string_view access, | |
| 1206 | std::optional<uintptr_t> stream) { | ||
| 1207 | ✗ | UntypedAttribute target = source.view.Get(field); | |
| 1208 | ✗ | nb::gil_scoped_release release; | |
| 1209 | ✗ | return new SceneMapping{source.scene, target.Map(Access(access), ToStream(stream)), | |
| 1210 | ✗ | stream}; | |
| 1211 | ✗ | }, | |
| 1212 | 6 | nb::arg("field"), nb::arg("access") = "write", nb::arg("stream") = nb::none(), | |
| 1213 | nb::keep_alive<0, 1>()) | ||
| 1214 | 4 | .def( | |
| 1215 | "map_primvar", | ||
| 1216 | 3 | [](const SceneSelection& source, std::string_view name, std::string_view access, | |
| 1217 | std::optional<uintptr_t> stream) { | ||
| 1218 | 3 | UntypedAttribute target = source.view.Get("primvars"); | |
| 1219 | 3 | nb::gil_scoped_release release; | |
| 1220 | 6 | return new SceneMapping{source.scene, | |
| 1221 | target.MapPrimvar(name, Access(access), ToStream(stream)), | ||
| 1222 | 4 | stream}; | |
| 1223 | 5 | }, | |
| 1224 | 6 | nb::arg("name"), nb::arg("access") = "write", nb::arg("stream") = nb::none(), | |
| 1225 | nb::keep_alive<0, 1>()) | ||
| 1226 | 2 | .def( | |
| 1227 | "read_poses", | ||
| 1228 | 7 | [](const SceneSelection& source, device::DType dtype, std::string_view order, | |
| 1229 | std::optional<uintptr_t> stream) { | ||
| 1230 | 5 | const QuaternionOrder chosen = FindOrder(order); | |
| 1231 | 5 | nb::gil_scoped_release release; | |
| 1232 | ✗ | return source.view.ReadPoses(dtype, stream ? ToStream(stream) : Stream::Host(), | |
| 1233 | 5 | chosen); | |
| 1234 | 5 | }, | |
| 1235 | 8 | nb::arg("dtype") = device::DType::Float32, nb::kw_only(), nb::arg("order") = "xyzw", | |
| 1236 | 4 | nb::arg("stream") = nb::none()) | |
| 1237 | 2 | .def( | |
| 1238 | "write_poses", | ||
| 1239 | 1139 | [](const SceneSelection& target, nb::handle tensor, nb::handle translation, | |
| 1240 | nb::handle rotation, std::string_view order, nb::handle rows, nb::handle mask, | ||
| 1241 | std::optional<uintptr_t> stream) { | ||
| 1242 | 1137 | WritePosesFrom(target, tensor, translation, rotation, order, rows, mask, stream, | |
| 1243 | false); | ||
| 1244 | }, | ||
| 1245 | 4 | nb::arg("tensor") = nb::none(), nb::kw_only(), nb::arg("translation") = nb::none(), | |
| 1246 | 4 | nb::arg("rotation") = nb::none(), nb::arg("order") = "xyzw", | |
| 1247 | 4 | nb::arg("rows") = nb::none(), nb::arg("mask") = nb::none(), | |
| 1248 | 4 | nb::arg("stream") = nb::none()) | |
| 1249 | 2 | .def( | |
| 1250 | "write_world_poses", | ||
| 1251 | 3 | [](const SceneSelection& target, nb::handle tensor, nb::handle translation, | |
| 1252 | nb::handle rotation, std::string_view order, nb::handle rows, nb::handle mask, | ||
| 1253 | std::optional<uintptr_t> stream) { | ||
| 1254 | 1 | WritePosesFrom(target, tensor, translation, rotation, order, rows, mask, stream, | |
| 1255 | true); | ||
| 1256 | }, | ||
| 1257 | 4 | nb::arg("tensor") = nb::none(), nb::kw_only(), nb::arg("translation") = nb::none(), | |
| 1258 | 4 | nb::arg("rotation") = nb::none(), nb::arg("order") = "xyzw", | |
| 1259 | 4 | nb::arg("rows") = nb::none(), nb::arg("mask") = nb::none(), | |
| 1260 | 4 | nb::arg("stream") = nb::none(), | |
| 1261 | "Places the prims at world poses, setting the local translations and rotations that " | ||
| 1262 | "put them there under their parents as of every write so far.") | ||
| 1263 | 2 | .def( | |
| 1264 | "read_world_poses", | ||
| 1265 | 8 | [](const SceneSelection& source, device::DType dtype, std::string_view order, | |
| 1266 | std::optional<uintptr_t> stream) { | ||
| 1267 | 6 | const QuaternionOrder chosen = FindOrder(order); | |
| 1268 | 6 | nb::gil_scoped_release release; | |
| 1269 | ✗ | return source.view.ReadWorldPoses(dtype, stream ? ToStream(stream) : Stream::Host(), | |
| 1270 | 6 | chosen); | |
| 1271 | 6 | }, | |
| 1272 | 8 | nb::arg("dtype") = device::DType::Float32, nb::kw_only(), nb::arg("order") = "xyzw", | |
| 1273 | 4 | nb::arg("stream") = nb::none(), | |
| 1274 | "The prims' world poses as an [N, 7] tensor of translations and rotations, as of " | ||
| 1275 | "every write so far. Prims without a transform read as zeros.") | ||
| 1276 | 2 | .def( | |
| 1277 | "read_world_scales", | ||
| 1278 | 3 | [](const SceneSelection& source, device::DType dtype, std::optional<uintptr_t> stream) { | |
| 1279 | 1 | nb::gil_scoped_release release; | |
| 1280 | 1 | return source.view.ReadWorldScales(dtype, | |
| 1281 | 1 | stream ? ToStream(stream) : Stream::Host()); | |
| 1282 | 1 | }, | |
| 1283 | 6 | nb::arg("dtype") = device::DType::Float32, nb::kw_only(), | |
| 1284 | 4 | nb::arg("stream") = nb::none(), | |
| 1285 | "The prims' world scales as an [N, 3] tensor, the lengths of their world axes.") | ||
| 1286 | 2 | .def("__getitem__", &Read) | |
| 1287 | 2 | .def("__setitem__", [](SceneSelection& source, std::string_view field, | |
| 1288 | 11 | nb::handle value) { Assign(source, field, value); }) | |
| 1289 | 4 | .def("set_each", &AssignEach, nb::arg("field"), nb::arg("values")) | |
| 1290 | 2 | .def("__repr__", [](const SceneSelection& source) { | |
| 1291 | ✗ | return "Selection(" + std::to_string(source.view.PrimCount()) + " prims)"; | |
| 1292 | }); | ||
| 1293 | 2 | } | |
| 1294 | |||
| 1295 | } | ||
| 1296 | |||
| 1297 | 2 | void BindScene(nb::module_& module) { | |
| 1298 | 2 | RegisterPrimvar(module); | |
| 1299 | 2 | RegisterPrimCopy(module); | |
| 1300 | 2 | RegisterScene(module); | |
| 1301 | 2 | RegisterPrim(module); | |
| 1302 | 2 | RegisterMapping(module); | |
| 1303 | 2 | RegisterSnapshot(module); | |
| 1304 | 2 | RegisterSelection(module); | |
| 1305 | 2 | } | |
| 1306 | |||
| 1307 | } | ||
| 1308 |