tnc/builders/
tensorgeneration.rs

1use itertools::Itertools;
2use num_complex::Complex64;
3use rand::distr::Uniform;
4use rand::Rng;
5use tetra::Tensor as DataTensor;
6
7use crate::tensornetwork::tensordata::TensorData;
8
9/// Generates random sparse [`DataTensor`] object.
10/// Fills in sparse tensor based on `sparsity` value (defaults to `0.5`).
11///
12/// # Examples
13/// ```
14/// # use tnc::builders::tensorgeneration::random_sparse_tensor_data_with_rng;
15/// let shape = vec![5, 4, 3];
16/// random_sparse_tensor_data_with_rng(&shape, None, &mut rand::thread_rng());
17/// ```
18pub fn random_sparse_tensor_data_with_rng<R>(
19    dims: &[usize],
20    sparsity: Option<f32>,
21    rng: &mut R,
22) -> TensorData
23where
24    R: Rng,
25{
26    let sparsity = if let Some(sparsity) = sparsity {
27        assert!((0.0..=1.0).contains(&sparsity));
28        sparsity
29    } else {
30        0.5
31    };
32
33    let ranges = dims
34        .iter()
35        .map(|i| Uniform::new(0, *i).unwrap())
36        .collect_vec();
37    let size = dims.iter().product::<usize>();
38    let mut tensor = DataTensor::new(dims);
39
40    let mut nnz = 0;
41    let mut loc = Vec::new();
42    while (nnz as f32 / size as f32) < sparsity {
43        for r in &ranges {
44            loc.push(rng.sample(r));
45        }
46        let val = Complex64::new(rng.random(), rng.random());
47        tensor.set(&loc, val);
48        loc.clear();
49        nnz += 1;
50    }
51
52    TensorData::Matrix(tensor)
53}
54
55/// Generates random sparse [`DataTensor`] object.
56/// Fills in sparse tensor based on `sparsity` value (defaults to `0.5`). Uses the thread-local random number generator.
57///
58/// # Examples
59/// ```
60/// # use tnc::builders::tensorgeneration::random_sparse_tensor_data;
61/// let shape = vec![5,4,3];
62/// let r_tensor = random_sparse_tensor_data(&shape, None);
63/// ```
64#[must_use]
65pub fn random_sparse_tensor_data(shape: &[usize], sparsity: Option<f32>) -> TensorData {
66    random_sparse_tensor_data_with_rng(shape, sparsity, &mut rand::rng())
67}