rusterize/
prelude.rs

1use crate::rasterization::pixel_functions::PixelFn;
2use num_traits::Num;
3use std::ops::AddAssign;
4
5pub use crate::{
6    encoding::arrays::{DenseArray, SparseArray},
7    error::{RusterizeError, RusterizeResult},
8    geo::raster::{RasterInfo, RasterInfoBuilder},
9    rasterization::pixel_functions::PixelFunction,
10    rasterize::{ArrayBuilder, FieldSource, Rasterize},
11};
12
13/// Trait to handle NaN check for dtypes that don't have it.
14pub trait NaNAware {
15    fn is_nan(&self) -> bool;
16}
17
18impl NaNAware for f32 {
19    fn is_nan(&self) -> bool {
20        f32::is_nan(*self)
21    }
22}
23
24impl NaNAware for f64 {
25    fn is_nan(&self) -> bool {
26        f64::is_nan(*self)
27    }
28}
29
30macro_rules! impl_maybe_nan_for_int {
31    ($($t:ty),*) => {
32        $(impl NaNAware for $t {
33            fn is_nan(&self) -> bool {
34                false
35            }
36        })*
37    };
38}
39
40impl_maybe_nan_for_int!(u8, u16, u32, u64, i8, i16, i32, i64);
41
42/// Handle polars dtypes and conversions.
43#[cfg(feature = "polars")]
44pub trait PolarsHandler: polars::prelude::Literal + Send + Sync {
45    type ChunkedArrayType: polars::prelude::PolarsNumericType<Native = Self> + 'static;
46    fn polars_dtype() -> polars::prelude::DataType;
47    fn from_named_vec(name: &str, vec: &[Self]) -> polars::prelude::Column
48    where
49        Self: Sized;
50}
51
52#[cfg(feature = "polars")]
53macro_rules! impl_polars_handler {
54    ($($t:ty => { dtype: $dtype:expr, catype: $catype:ty }),* $(,)?) => {
55        $(
56            impl PolarsHandler for $t {
57                type ChunkedArrayType = $catype;
58                fn polars_dtype() -> polars::prelude::DataType { $dtype }
59                fn from_named_vec(name: &str, vec: &[Self]) -> polars::prelude::Column {
60                    polars::prelude::Column::new(name.into(), vec)
61                }
62            }
63        )*
64    };
65}
66
67#[cfg(feature = "polars")]
68impl_polars_handler! {
69    f64 => { dtype: polars::prelude::DataType::Float64, catype: polars::prelude::Float64Type },
70    f32 => { dtype: polars::prelude::DataType::Float32, catype: polars::prelude::Float32Type },
71    u8  => { dtype: polars::prelude::DataType::UInt8,   catype: polars::prelude::UInt8Type },
72    i8  => { dtype: polars::prelude::DataType::Int8,    catype: polars::prelude::Int8Type },
73    u16 => { dtype: polars::prelude::DataType::UInt16,  catype: polars::prelude::UInt16Type },
74    i16 => { dtype: polars::prelude::DataType::Int16,   catype: polars::prelude::Int16Type },
75    u32 => { dtype: polars::prelude::DataType::UInt32,  catype: polars::prelude::UInt32Type },
76    i32 => { dtype: polars::prelude::DataType::Int32,   catype: polars::prelude::Int32Type },
77    u64 => { dtype: polars::prelude::DataType::UInt64,  catype: polars::prelude::UInt64Type },
78    i64 => { dtype: polars::prelude::DataType::Int64,   catype: polars::prelude::Int64Type },
79}
80
81/// Bound rasterization to a dtype.
82#[cfg(feature = "polars")]
83pub trait RasterDtype: Num + Copy + AddAssign + PartialOrd + NaNAware + PolarsHandler {}
84#[cfg(feature = "polars")]
85impl<N: Num + Copy + AddAssign + PartialOrd + NaNAware + PolarsHandler> RasterDtype for N {}
86#[cfg(not(feature = "polars"))]
87pub trait RasterDtype: Num + Copy + AddAssign + PartialOrd + NaNAware + Send + Sync {}
88#[cfg(not(feature = "polars"))]
89impl<N: Num + Copy + AddAssign + PartialOrd + NaNAware + Send + Sync> RasterDtype for N {}
90
91/// Spatial + value context handed to the rasterization engine.
92#[derive(Clone)]
93pub struct RasterizeContext<'a, N> {
94    /// The spatial information of the final raster.
95    pub raster_info: RasterInfo,
96    /// The values to burn.
97    pub field: FieldSource<'a, N>,
98    /// Specify the grouping of the geometries into multiple bands in the final raster. None is no grouping.
99    /// For this to work, `by` has to have the same length of the geometries.
100    pub by: Option<&'a [String]>,
101    /// Describes what happens to overlapping pixels.
102    pub pixel_fn: PixelFunction,
103    pub background: N,
104    /// Flags whether all pixels touching the geometry should be burned.
105    pub all_touched: bool,
106}
107
108impl<'a, N> RasterizeContext<'a, N> {
109    pub(crate) fn pixel_fn(&self) -> PixelFn<N>
110    where
111        N: Num + Copy + AddAssign + PartialOrd + NaNAware,
112    {
113        self.pixel_fn.to_function()
114    }
115
116    pub(crate) fn requires_dedup(&self) -> bool {
117        self.all_touched && matches!(self.pixel_fn, PixelFunction::Sum | PixelFunction::Count)
118    }
119}
rusterize/prelude.rs

rusterize/
prelude.rs
