burn_tensor/tensor/grid/
meshgrid.rs

1use crate::kind::Basic;
2use crate::tensor::Tensor;
3use crate::tensor::grid::{GridIndexing, GridOptions, GridSparsity, IndexPos};
4use alloc::vec::Vec;
5
6/// Return a collection of coordinate matrices for coordinate vectors.
7///
8/// Takes N 1D tensors and returns N tensors where each tensor represents the coordinates
9/// in one dimension across an N-dimensional grid.
10///
11/// Based upon `options.sparse`, the generated coordinate tensors can either be `Sparse` or `Dense`:
12/// * In `Sparse` mode, output tensors will have shape 1 everywhere except their cardinal dimension.
13/// * In `Dense` mode, output tensors will be expanded to the full grid shape.
14///
15/// Based upon `options.indexing`, the generated coordinate tensors will use either:
16/// * `Matrix` indexing, where dimensions are in the same order as their cardinality.
17/// * `Cartesian` indexing; where the first two dimensions are swapped.
18///
19/// See:
20///  - [numpy.meshgrid](https://numpy.org/doc/stable/reference/generated/numpy.meshgrid.html)
21///  - [torch.meshgrid](https://pytorch.org/docs/stable/generated/torch.meshgrid.html)
22///
23/// # Arguments
24///
25/// * `tensors` - A slice of 1D tensors
26/// * `options` - the options.
27///
28/// # Returns
29///
30/// A vector of N N-dimensional tensors representing the grid coordinates.
31pub fn meshgrid<const N: usize, K, O>(tensors: &[Tensor<1, K>; N], options: O) -> [Tensor<N, K>; N]
32where
33    K: Basic,
34    O: Into<GridOptions>,
35{
36    let options = options.into();
37    let swap_dims = options.indexing == GridIndexing::Cartesian && N > 1;
38    let dense = options.sparsity == GridSparsity::Dense;
39
40    let grid_shape: [usize; N] = tensors
41        .iter()
42        .map(|t| t.dims()[0])
43        .collect::<Vec<_>>()
44        .try_into()
45        .unwrap();
46
47    tensors
48        .iter()
49        .enumerate()
50        .map(|(i, tensor)| {
51            let mut coord_tensor_shape = [1; N];
52            coord_tensor_shape[i] = grid_shape[i];
53
54            // Reshape the tensor to have singleton dimensions in all but the i-th dimension
55            let mut tensor = tensor.clone().reshape(coord_tensor_shape);
56
57            if dense {
58                tensor = tensor.expand(grid_shape);
59            }
60            if swap_dims {
61                tensor = tensor.swap_dims(0, 1);
62            }
63
64            tensor
65        })
66        .collect::<Vec<_>>()
67        .try_into()
68        .unwrap()
69}
70
71/// Return a coordinate matrix for a given set of 1D coordinate tensors.
72///
73/// Equivalent to stacking a dense matrix `meshgrid`,
74/// where the stack is along the first or last dimension.
75///
76/// # Arguments
77///
78/// * `tensors`: A slice of 1D tensors.
79/// * `index_pos`: The position of the index in the output tensor.
80///
81/// # Returns
82///
83/// A tensor of either ``(N, ..., |T[i]|, ...)`` or ``(..., |T[i]|, ..., N)``,
84/// of coordinates, indexed on the first or last dimension.
85pub fn meshgrid_stack<const D: usize, const D2: usize, K>(
86    tensors: &[Tensor<1, K>; D],
87    index_pos: IndexPos,
88) -> Tensor<D2, K>
89where
90    K: Basic,
91{
92    assert_eq!(D2, D + 1, "D2 ({D2}) != D ({D}) + 1");
93
94    let xs: Vec<Tensor<D, K>> = meshgrid(tensors, GridOptions::default())
95        .into_iter()
96        .collect();
97
98    let dim = match index_pos {
99        IndexPos::First => 0,
100        IndexPos::Last => D,
101    };
102
103    Tensor::stack(xs, dim)
104}
burn_tensor/tensor/grid/meshgrid.rs

burn_tensor/tensor/grid/
meshgrid.rs
