hand.rs
18.0 kB · rust · 556 lines
1use mrlyrs::core::cell::Cell;2use mrlyrs::core::colors::Color;3use mrlyrs::core::error::{Error, Result};4use mrlyrs::core::rng::Rng;5use mrlyrs::core::tensor::{Dtype, Tensor};6use mrlyrs::math::bang::Code;7use mrlyrs::math::cell::models::CellNd;8use mrlyrs::math::six::{Cell6d, Orientation, Projection};9use numpy::ndarray::{ArrayD, IxDyn};10use numpy::{11 Element, IntoPyArray, PyReadonlyArrayDyn, PyReadwriteArrayDyn, PyUntypedArrayMethods,12};13use pyo3::exceptions::{PyOverflowError, PyValueError};14use pyo3::prelude::*;15use pyo3::types::{PyDict, PyList, PyTuple};16use pyo3::{Borrowed, IntoPyObjectExt};17use pythonize::{depythonize, pythonize};18use serde::de::DeserializeOwned;19use serde::Serialize;2021// ERRORS2223/// Maps one crate error onto its Python class, overflow apart and everything else a ValueError.24pub fn py_error(error: Error) -> PyErr {25 match error {26 Error::Overflow(message) => PyOverflowError::new_err(message),27 other => PyValueError::new_err(other.to_string()),28 }29}3031/// Unwraps a crate result into a Python result.32pub fn ok<T>(result: Result<T>) -> PyResult<T> {33 result.map_err(py_error)34}3536fn bad(message: impl Into<String>) -> PyErr {37 PyValueError::new_err(message.into())38}3940// TENSOR4142/// A tensor crossing as a numpy array of its own dtype.43pub struct PyTensor(pub Tensor);4445fn owned<'py, T: Element>(46 py: Python<'py>,47 shape: &[usize],48 data: Vec<T>,49) -> PyResult<Bound<'py, PyAny>> {50 let array =51 ArrayD::from_shape_vec(IxDyn(shape), data).map_err(|error| bad(error.to_string()))?;52 Ok(array.into_pyarray(py).into_any())53}5455fn read<T: Element + Copy>(obj: &Bound<'_, PyAny>) -> Option<PyResult<(Vec<T>, Vec<usize>)>> {56 let view = obj.extract::<PyReadonlyArrayDyn<T>>().ok()?;57 let shape = view.shape().to_vec();58 Some(match view.as_slice() {59 Ok(data) => Ok((data.to_vec(), shape)),60 Err(_) => Err(bad("an array crossing into Rust must be C-contiguous.")),61 })62}6364/// Builds a numpy array of the tensor's shape and dtype, the buffer moved into numpy.65pub fn tensor_into_py<'py>(py: Python<'py>, tensor: &Tensor) -> PyResult<Bound<'py, PyAny>> {66 let shape = &tensor.shape;67 match tensor.dtype() {68 Dtype::U8 => owned(py, shape, ok(tensor.bytes())?.to_vec()),69 Dtype::U16 => owned(py, shape, ok(tensor.u16s())?.to_vec()),70 Dtype::U32 => owned(py, shape, ok(tensor.u32s())?.to_vec()),71 Dtype::I32 => owned(py, shape, ok(tensor.i32s())?.to_vec()),72 }73}7475fn copy_back<T: Element + Copy>(obj: &Bound<'_, PyAny>, data: &[T]) -> PyResult<()> {76 let mut view = obj.extract::<PyReadwriteArrayDyn<T>>().map_err(|_| {77 bad("a mutated tensor writes back into a writable array of its own dtype.")78 })?;79 let slot = view80 .as_slice_mut()81 .map_err(|_| bad("a mutated tensor writes back into a C-contiguous array."))?;82 if slot.len() != data.len() {83 return Err(bad("a mutated tensor keeps its shape."));84 }85 slot.copy_from_slice(data);86 Ok(())87}8889/// Copies a mutated tensor back into the numpy array it was read from.90pub fn tensor_write_back(obj: &Bound<'_, PyAny>, tensor: &Tensor) -> PyResult<()> {91 match tensor.dtype() {92 Dtype::U8 => copy_back(obj, ok(tensor.bytes())?),93 Dtype::U16 => copy_back(obj, ok(tensor.u16s())?),94 Dtype::U32 => copy_back(obj, ok(tensor.u32s())?),95 Dtype::I32 => copy_back(obj, ok(tensor.i32s())?),96 }97}9899/// Reads a C-contiguous uint8, uint16, uint32 or int32 array into a tensor.100pub fn tensor_from_py(obj: &Bound<'_, PyAny>) -> PyResult<Tensor> {101 if let Some(result) = read::<u8>(obj) {102 let (data, shape) = result?;103 return ok(Tensor::u8(data, shape));104 }105 if let Some(result) = read::<u16>(obj) {106 let (data, shape) = result?;107 return ok(Tensor::u16(data, shape));108 }109 if let Some(result) = read::<u32>(obj) {110 let (data, shape) = result?;111 return ok(Tensor::u32(data, shape));112 }113 if let Some(result) = read::<i32>(obj) {114 let (data, shape) = result?;115 return ok(Tensor::i32(data, shape));116 }117 Err(bad(118 "a tensor wants a uint8, uint16, uint32 or int32 numpy array.",119 ))120}121122impl<'py> IntoPyObject<'py> for PyTensor {123 type Target = PyAny;124 type Output = Bound<'py, PyAny>;125 type Error = PyErr;126 fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {127 tensor_into_py(py, &self.0)128 }129}130131impl<'py> FromPyObject<'_, 'py> for PyTensor {132 type Error = PyErr;133 fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<PyTensor> {134 Ok(PyTensor(tensor_from_py(&obj)?))135 }136}137138// CELL139140/// A cell crossing as the dict of its types, colors and tags.141pub struct PyCell(pub Cell);142143fn field<'py>(obj: &Bound<'py, PyAny>, name: &str) -> Option<Bound<'py, PyAny>> {144 match obj.get_item(name) {145 Ok(value) if !value.is_none() => Some(value),146 _ => None,147 }148}149150/// Builds the `{types, colors, tags}` dict of a cell, each array moved into numpy.151pub fn cell_into_py<'py>(py: Python<'py>, cell: Cell) -> PyResult<Bound<'py, PyAny>> {152 let Cell {153 types,154 colors,155 tags,156 } = cell;157 let dict = PyDict::new(py);158 dict.set_item("types", tensor_into_py(py, &types)?)?;159 match colors {160 Some(colors) => {161 let mut shape = types.shape.clone();162 shape.push(4);163 let flat: Vec<u8> = colors.into_iter().flatten().collect();164 dict.set_item("colors", owned(py, &shape, flat)?)?;165 }166 None => dict.set_item("colors", py.None())?,167 }168 match tags {169 Some(tags) => dict.set_item("tags", tensor_into_py(py, &tags)?)?,170 None => dict.set_item("tags", py.None())?,171 }172 Ok(dict.into_any())173}174175/// Reads a `{types, colors, tags}` dict into a cell.176pub fn cell_from_py(obj: &Bound<'_, PyAny>) -> PyResult<Cell> {177 let types = match field(obj, "types") {178 Some(value) => tensor_from_py(&value)?,179 None => return Err(bad("a cell wants a \"types\" array.")),180 };181 let colors = match field(obj, "colors") {182 Some(value) => {183 let view = value184 .extract::<PyReadonlyArrayDyn<u8>>()185 .map_err(|_| bad("cell colors want a uint8 array."))?;186 if view.shape().last() != Some(&4) {187 return Err(bad("cell colors want a trailing axis of four channels."));188 }189 let data = view190 .as_slice()191 .map_err(|_| bad("cell colors want a C-contiguous array."))?;192 if data.len() != types.size() * 4 {193 return Err(bad("cell colors want one rgba per cell."));194 }195 Some(196 data.chunks_exact(4)197 .map(|rgba| [rgba[0], rgba[1], rgba[2], rgba[3]])198 .collect(),199 )200 }201 None => None,202 };203 let tags = match field(obj, "tags") {204 Some(value) => Some(tensor_from_py(&value)?),205 None => None,206 };207 Ok(Cell {208 types,209 colors,210 tags,211 })212}213214/// Writes a mutated cell's types, colors and tags back into the dict it was read from.215pub fn cell_write_back(obj: &Bound<'_, PyAny>, cell: Cell) -> PyResult<()> {216 let fresh = cell_into_py(obj.py(), cell)?;217 for key in ["types", "colors", "tags"] {218 obj.set_item(key, fresh.get_item(key)?)?;219 }220 Ok(())221}222223/// Reads the rank of the array, the cell dict or the first of a list of them.224pub fn ndim(obj: &Bound<'_, PyAny>) -> PyResult<usize> {225 if let Ok(dict) = obj.cast::<PyDict>() {226 return match dict.get_item("types")? {227 Some(types) => ndim(&types),228 None => Err(bad("a cell wants a \"types\" array.")),229 };230 }231 if let Ok(list) = obj.cast::<PyList>() {232 return match list.get_item(0) {233 Ok(first) => ndim(&first),234 Err(_) => Err(bad("a dispatch wants at least one cell.")),235 };236 }237 if let Ok(tuple) = obj.cast::<PyTuple>() {238 return match tuple.get_item(0) {239 Ok(first) => ndim(&first),240 Err(_) => Err(bad("a dispatch wants at least one cell.")),241 };242 }243 obj.getattr("ndim")244 .and_then(|rank| rank.extract::<usize>())245 .map_err(|_| bad("a dispatch wants a numpy array, a cell dict or a list of them."))246}247248impl<'py> IntoPyObject<'py> for PyCell {249 type Target = PyAny;250 type Output = Bound<'py, PyAny>;251 type Error = PyErr;252 fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {253 cell_into_py(py, self.0)254 }255}256257impl<'py> FromPyObject<'_, 'py> for PyCell {258 type Error = PyErr;259 fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<PyCell> {260 Ok(PyCell(cell_from_py(&obj)?))261 }262}263264// CELLND265266/// An N-dimensional cell crossing as the same dict, N read off the types array.267pub struct PyCellNd<const N: usize>(pub CellNd<N>);268269/// The two-dimensional crossing.270pub type PyCell2d = PyCellNd<2>;271272/// The three-dimensional crossing.273pub type PyCell3d = PyCellNd<3>;274275/// Reads the dict into an N-dimensional cell, erring when the types array is not N-dimensional.276pub fn cell_nd_from_py<const N: usize>(obj: &Bound<'_, PyAny>) -> PyResult<CellNd<N>> {277 let cell = cell_from_py(obj)?;278 let rank = cell.types.shape.len();279 if rank != N {280 return Err(bad(format!(281 "a {N}d cell wants a {N}d types array, got {rank}d."282 )));283 }284 Ok(CellNd { cell })285}286287impl<'py, const N: usize> IntoPyObject<'py> for PyCellNd<N> {288 type Target = PyAny;289 type Output = Bound<'py, PyAny>;290 type Error = PyErr;291 fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {292 cell_into_py(py, self.0.cell)293 }294}295296impl<'py, const N: usize> FromPyObject<'_, 'py> for PyCellNd<N> {297 type Error = PyErr;298 fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<PyCellNd<N>> {299 Ok(PyCellNd(cell_nd_from_py(&obj)?))300 }301}302303// CELL6D304305/// A hexagonal cell crossing as the dict of its four fields.306pub struct PyCell6d(pub Cell6d);307308/// Builds the `{cell, projection, orientation, start}` dict of a hexagonal cell.309pub fn cell_6d_into_py<'py>(py: Python<'py>, cell: Cell6d) -> PyResult<Bound<'py, PyAny>> {310 let dict = PyDict::new(py);311 dict.set_item("cell", cell_into_py(py, cell.cell.cell)?)?;312 dict.set_item("projection", serde_into_py(py, &cell.projection)?)?;313 dict.set_item("orientation", serde_into_py(py, &cell.orientation)?)?;314 dict.set_item("start", cell.start)?;315 Ok(dict.into_any())316}317318/// Reads a `{cell, projection, orientation, start}` dict into a hexagonal cell.319pub fn cell_6d_from_py(obj: &Bound<'_, PyAny>) -> PyResult<Cell6d> {320 let inner = match field(obj, "cell") {321 Some(value) => cell_nd_from_py::<2>(&value)?,322 None => return Err(bad("a 6d cell wants a \"cell\" dict.")),323 };324 let projection: Projection = match field(obj, "projection") {325 Some(value) => serde_from_py(&value)?,326 None => return Err(bad("a 6d cell wants a \"projection\".")),327 };328 let orientation: Orientation = match field(obj, "orientation") {329 Some(value) => serde_from_py(&value)?,330 None => return Err(bad("a 6d cell wants an \"orientation\".")),331 };332 let start: u8 = match field(obj, "start") {333 Some(value) => value.extract()?,334 None => return Err(bad("a 6d cell wants a \"start\".")),335 };336 Ok(Cell6d::new(inner, projection, orientation, start))337}338339impl<'py> IntoPyObject<'py> for PyCell6d {340 type Target = PyAny;341 type Output = Bound<'py, PyAny>;342 type Error = PyErr;343 fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {344 cell_6d_into_py(py, self.0)345 }346}347348impl<'py> FromPyObject<'_, 'py> for PyCell6d {349 type Error = PyErr;350 fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<PyCell6d> {351 Ok(PyCell6d(cell_6d_from_py(&obj)?))352 }353}354355// COLOR356357/// A color crossing as the four-tuple of its channels.358pub struct PyColor(pub Color);359360impl<'py> IntoPyObject<'py> for PyColor {361 type Target = PyAny;362 type Output = Bound<'py, PyAny>;363 type Error = PyErr;364 fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {365 let color = self.0;366 (color.r, color.g, color.b, color.a).into_bound_py_any(py)367 }368}369370impl<'py> FromPyObject<'_, 'py> for PyColor {371 type Error = PyErr;372 fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<PyColor> {373 let (r, g, b, a) = obj374 .extract::<(u8, u8, u8, u8)>()375 .map_err(|_| bad("a color wants four channel bytes."))?;376 Ok(PyColor(Color::rgba(r, g, b, a)))377 }378}379380// PIXELS381382/// One rgba pixel crossing as a four-tuple of channel bytes.383pub struct PyRgba(pub [u8; 4]);384385impl<'py> IntoPyObject<'py> for PyRgba {386 type Target = PyAny;387 type Output = Bound<'py, PyAny>;388 type Error = PyErr;389 fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {390 let [r, g, b, a] = self.0;391 (r, g, b, a).into_bound_py_any(py)392 }393}394395impl<'py> FromPyObject<'_, 'py> for PyRgba {396 type Error = PyErr;397 fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<PyRgba> {398 let rgba = obj399 .extract::<[u8; 4]>()400 .map_err(|_| bad("a pixel wants four channel bytes."))?;401 Ok(PyRgba(rgba))402 }403}404405/// A run of rgba pixels crossing as an (n, 4) uint8 array, a list of four-tuples read too.406pub struct PyPixels(pub Vec<[u8; 4]>);407408impl<'py> IntoPyObject<'py> for PyPixels {409 type Target = PyAny;410 type Output = Bound<'py, PyAny>;411 type Error = PyErr;412 fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {413 let count = self.0.len();414 let flat: Vec<u8> = self.0.into_iter().flatten().collect();415 owned(py, &[count, 4], flat)416 }417}418419impl<'py> FromPyObject<'_, 'py> for PyPixels {420 type Error = PyErr;421 fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<PyPixels> {422 if let Ok(view) = obj.extract::<PyReadonlyArrayDyn<u8>>() {423 if view.shape().last() != Some(&4) {424 return Err(bad("pixels want a trailing axis of four channels."));425 }426 let data = view427 .as_slice()428 .map_err(|_| bad("pixels want a C-contiguous array."))?;429 return Ok(PyPixels(430 data.chunks_exact(4)431 .map(|rgba| [rgba[0], rgba[1], rgba[2], rgba[3]])432 .collect(),433 ));434 }435 let listed = obj436 .extract::<Vec<PyRgba>>()437 .map_err(|_| bad("pixels want an (n, 4) uint8 array or a list of four-tuples."))?;438 Ok(PyPixels(listed.into_iter().map(|p| p.0).collect()))439 }440}441442// CODE443444/// A design code crossing as a Python int.445pub struct PyCode(pub Code);446447impl<'py> IntoPyObject<'py> for PyCode {448 type Target = PyAny;449 type Output = Bound<'py, PyAny>;450 type Error = PyErr;451 fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {452 self.0.get().into_bound_py_any(py)453 }454}455456impl<'py> FromPyObject<'_, 'py> for PyCode {457 type Error = PyErr;458 fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<PyCode> {459 let bits = obj460 .extract::<u128>()461 .map_err(|_| bad("a code wants a non-negative integer below 2^128."))?;462 Ok(PyCode(Code::from(bits)))463 }464}465466// SERDE467468/// Every other serde type crossing as plain Python data.469pub struct PySerde<T>(pub T);470471/// Turns any serde value into plain Python data.472pub fn serde_into_py<'py, T: Serialize + ?Sized>(473 py: Python<'py>,474 value: &T,475) -> PyResult<Bound<'py, PyAny>> {476 Ok(pythonize(py, value)?)477}478479/// Reads plain Python data back into any serde type.480pub fn serde_from_py<T: DeserializeOwned>(obj: &Bound<'_, PyAny>) -> PyResult<T> {481 Ok(depythonize(obj)?)482}483484impl<'py, T: Serialize> IntoPyObject<'py> for PySerde<T> {485 type Target = PyAny;486 type Output = Bound<'py, PyAny>;487 type Error = PyErr;488 fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {489 serde_into_py(py, &self.0)490 }491}492493impl<'py, T: DeserializeOwned> FromPyObject<'_, 'py> for PySerde<T> {494 type Error = PyErr;495 fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult<PySerde<T>> {496 Ok(PySerde(serde_from_py(&obj)?))497 }498}499500// RNG501502/// The seeded random stream, one class, passed wherever Rust takes a mutable stream.503#[pyclass(name = "Rng", module = "mrlypy.core")]504pub struct PyRng(pub Rng);505506#[pymethods]507impl PyRng {508 /// Opens the stream of the seed.509 #[new]510 pub fn new(seed: u64) -> PyRng {511 PyRng(Rng::new(seed))512 }513 /// Draws a float at or above zero and below one.514 pub fn unit(&mut self) -> f64 {515 self.0.unit()516 }517 /// Draws an integer below n, or zero when n is zero.518 pub fn below(&mut self, n: usize) -> usize {519 self.0.below(n)520 }521 /// Draws an integer between lo and hi inclusive.522 pub fn range(&mut self, lo: i64, hi: i64) -> i64 {523 self.0.range(lo, hi)524 }525 /// Draws a fair coin flip.526 pub fn boolean(&mut self) -> bool {527 self.0.boolean()528 }529 /// Returns true with probability p.530 pub fn chance(&mut self, p: f64) -> bool {531 self.0.chance(p)532 }533 /// Draws amount distinct indices below length.534 pub fn sample_indices(&mut self, length: usize, amount: usize) -> Vec<usize> {535 self.0.sample_indices(length, amount)536 }537 /// Draws one item of the sequence, the same draw as Rust's choice.538 pub fn choice<'py>(&mut self, seq: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {539 let indices: Vec<usize> = (0..seq.len()?).collect();540 let index = *ok(self.0.choice(&indices))?;541 seq.get_item(index)542 }543 /// Shuffles the list in place, the same permutation as Rust's shuffle.544 pub fn shuffle(&mut self, seq: &Bound<'_, PyList>) -> PyResult<()> {545 let mut order: Vec<usize> = (0..seq.len()).collect();546 self.0.shuffle(&mut order);547 let items: Vec<Bound<'_, PyAny>> = order548 .into_iter()549 .map(|index| seq.get_item(index))550 .collect::<PyResult<_>>()?;551 for (at, item) in items.into_iter().enumerate() {552 seq.set_item(at, item)?;553 }554 Ok(())555 }556}