thrust_py/
intervals.rs

1use thrust::intervals::{Interval, IntervalCollection};
2
3use numpy::{PyArray1, PyReadonlyArray1};
4use pyo3::types::IntoPyDict;
5use pyo3::{prelude::*, types::PyDict};
6
7fn get_ic(start: PyReadonlyArray1<i64>, stop: PyReadonlyArray1<i64>) -> IntervalCollection<i64> {
8    let size1 = start.len().unwrap();
9    let size2 = stop.len().unwrap();
10    let size = std::cmp::min(size1, size2);
11
12    let mut elts = Vec::<Interval<i64>>::with_capacity(size);
13    for i in 0..size {
14        elts.push(Interval {
15            start: *start.get(i).unwrap(),
16            stop: *stop.get(i).unwrap(),
17        })
18    }
19    IntervalCollection { elts }
20}
21
22#[pyfunction]
23fn interval_and(py: Python, start1: i64, stop1: i64, start2: i64, stop2: i64) -> PyResult<Bound<PyDict>> {
24    let left = Interval {
25        start: start1,
26        stop: stop1,
27    };
28    let right = Interval {
29        start: start2,
30        stop: stop2,
31    };
32    let res = match &left & &right {
33        None => [("empty", true)].into_py_dict(py),
34        Some(Interval { start, stop }) => [("start", start), ("stop", stop)].into_py_dict(py),
35    };
36    res
37}
38
39#[pyfunction]
40fn collection_and<'a>(
41    py: Python<'a>,
42    start1: PyReadonlyArray1<i64>,
43    stop1: PyReadonlyArray1<i64>,
44    start2: PyReadonlyArray1<i64>,
45    stop2: PyReadonlyArray1<i64>,
46) -> PyResult<Bound<'a, PyDict>> {
47    let left = get_ic(start1, stop1);
48    let right = get_ic(start2, stop2);
49    let res = &left & &right;
50    let start: Vec<i64> = res.elts.iter().map(|elt| elt.start).collect();
51    let stop: Vec<i64> = res.elts.iter().map(|elt| elt.stop).collect();
52
53    let wrapped_res = PyDict::new(py);
54    wrapped_res.set_item("start", PyArray1::from_vec(py, start))?;
55    wrapped_res.set_item("stop", PyArray1::from_vec(py, stop))?;
56    Ok(wrapped_res)
57}
58
59#[pyfunction]
60fn collection_andi<'a>(
61    py: Python<'a>,
62    start1: PyReadonlyArray1<i64>,
63    stop1: PyReadonlyArray1<i64>,
64    start2: i64,
65    stop2: i64,
66) -> PyResult<Bound<'a, PyDict>> {
67    let left = get_ic(start1, stop1);
68    let right = Interval {
69        start: start2,
70        stop: stop2,
71    };
72
73    let res = &left & &right;
74    let start: Vec<i64> = res.elts.iter().map(|elt| elt.start).collect();
75    let stop: Vec<i64> = res.elts.iter().map(|elt| elt.stop).collect();
76
77    let wrapped_res = PyDict::new(py);
78    wrapped_res.set_item("start", PyArray1::from_vec(py, start))?;
79    wrapped_res.set_item("stop", PyArray1::from_vec(py, stop))?;
80    Ok(wrapped_res)
81}
82
83#[pyfunction]
84fn interval_add(py: Python, start1: i64, stop1: i64, start2: i64, stop2: i64) -> PyResult<Bound<PyDict>> {
85    let left = Interval {
86        start: start1,
87        stop: stop1,
88    };
89    let right = Interval {
90        start: start2,
91        stop: stop2,
92    };
93
94    let res = left + right;
95    let start: Vec<i64> = res.elts.iter().map(|elt| elt.start).collect();
96    let stop: Vec<i64> = res.elts.iter().map(|elt| elt.stop).collect();
97
98    let wrapped_res = PyDict::new(py);
99    wrapped_res.set_item("start", PyArray1::from_vec(py, start))?;
100    wrapped_res.set_item("stop", PyArray1::from_vec(py, stop))?;
101    Ok(wrapped_res)
102}
103
104#[pyfunction]
105fn collection_add<'a>(
106    py: Python<'a>,
107    start1: PyReadonlyArray1<i64>,
108    stop1: PyReadonlyArray1<i64>,
109    start2: PyReadonlyArray1<i64>,
110    stop2: PyReadonlyArray1<i64>,
111) -> PyResult<Bound<'a, PyDict>> {
112    let left = get_ic(start1, stop1);
113    let right = get_ic(start2, stop2);
114    let res = left + right;
115
116    let start: Vec<i64> = res.elts.iter().map(|elt| elt.start).collect();
117    let stop: Vec<i64> = res.elts.iter().map(|elt| elt.stop).collect();
118
119    let wrapped_res = PyDict::new(py);
120    wrapped_res.set_item("start", PyArray1::from_vec(py, start))?;
121    wrapped_res.set_item("stop", PyArray1::from_vec(py, stop))?;
122    Ok(wrapped_res)
123}
124
125#[pyfunction]
126fn collection_addi<'a>(
127    py: Python<'a>,
128    start1: PyReadonlyArray1<i64>,
129    stop1: PyReadonlyArray1<i64>,
130    start2: i64,
131    stop2: i64,
132) -> PyResult<Bound<'a, PyDict>> {
133    let left = get_ic(start1, stop1);
134    let right = Interval {
135        start: start2,
136        stop: stop2,
137    };
138    let res = &left + &right;
139
140    let start: Vec<i64> = res.elts.iter().map(|elt| elt.start).collect();
141    let stop: Vec<i64> = res.elts.iter().map(|elt| elt.stop).collect();
142
143    let wrapped_res = PyDict::new(py);
144    wrapped_res.set_item("start", PyArray1::from_vec(py, start))?;
145    wrapped_res.set_item("stop", PyArray1::from_vec(py, stop))?;
146    Ok(wrapped_res)
147}
148
149#[pyfunction]
150fn interval_sub(py: Python, start1: i64, stop1: i64, start2: i64, stop2: i64) -> PyResult<Bound<PyDict>> {
151    let left = Interval {
152        start: start1,
153        stop: stop1,
154    };
155    let right = Interval {
156        start: start2,
157        stop: stop2,
158    };
159
160    let res = left - right;
161    let start: Vec<i64> = res.elts.iter().map(|elt| elt.start).collect();
162    let stop: Vec<i64> = res.elts.iter().map(|elt| elt.stop).collect();
163
164    let wrapped_res = PyDict::new(py);
165    wrapped_res.set_item("start", PyArray1::from_vec(py, start))?;
166    wrapped_res.set_item("stop", PyArray1::from_vec(py, stop))?;
167    Ok(wrapped_res)
168}
169
170#[pyfunction]
171fn collection_sub<'a>(
172    py: Python<'a>,
173    start1: PyReadonlyArray1<i64>,
174    stop1: PyReadonlyArray1<i64>,
175    start2: PyReadonlyArray1<i64>,
176    stop2: PyReadonlyArray1<i64>,
177) -> PyResult<Bound<'a, PyDict>> {
178    let left = get_ic(start1, stop1);
179    let right = get_ic(start2, stop2);
180    let res = left - right;
181
182    let start: Vec<i64> = res.elts.iter().map(|elt| elt.start).collect();
183    let stop: Vec<i64> = res.elts.iter().map(|elt| elt.stop).collect();
184
185    let wrapped_res = PyDict::new(py);
186    wrapped_res.set_item("start", PyArray1::from_vec(py, start))?;
187    wrapped_res.set_item("stop", PyArray1::from_vec(py, stop))?;
188    Ok(wrapped_res)
189}
190
191#[pyfunction]
192fn collection_subi<'a>(
193    py: Python<'a>,
194    start1: PyReadonlyArray1<i64>,
195    stop1: PyReadonlyArray1<i64>,
196    start2: i64,
197    stop2: i64,
198) -> PyResult<Bound<'a, PyDict>> {
199    let left = get_ic(start1, stop1);
200    let right = Interval {
201        start: start2,
202        stop: stop2,
203    };
204    let res = left - right;
205
206    let start: Vec<i64> = res.elts.iter().map(|elt| elt.start).collect();
207    let stop: Vec<i64> = res.elts.iter().map(|elt| elt.stop).collect();
208
209    let wrapped_res = PyDict::new(py);
210    wrapped_res.set_item("start", PyArray1::from_vec(py, start))?;
211    wrapped_res.set_item("stop", PyArray1::from_vec(py, stop))?;
212    Ok(wrapped_res)
213}
214
215pub fn init(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
216    let m = PyModule::new(py, "intervals")?;
217
218    m.add_function(wrap_pyfunction!(interval_and, &m)?)?;
219    m.add_function(wrap_pyfunction!(collection_and, &m)?)?;
220    m.add_function(wrap_pyfunction!(collection_andi, &m)?)?;
221
222    m.add_function(wrap_pyfunction!(interval_add, &m)?)?;
223    m.add_function(wrap_pyfunction!(collection_add, &m)?)?;
224    m.add_function(wrap_pyfunction!(collection_addi, &m)?)?;
225
226    m.add_function(wrap_pyfunction!(interval_sub, &m)?)?;
227    m.add_function(wrap_pyfunction!(collection_sub, &m)?)?;
228    m.add_function(wrap_pyfunction!(collection_subi, &m)?)?;
229
230    Ok(m)
231}
thrust_py/intervals.rs

thrust_py/
intervals.rs
