1
1
use ndarray:: { ArrayD , ArrayViewD , ArrayViewMutD } ;
2
- use numpy:: { IntoPyArray , PyArrayDyn , PyReadonlyArrayDyn } ;
2
+ use numpy:: { c64 , IntoPyArray , PyArrayDyn , PyReadonlyArrayDyn } ;
3
3
use pyo3:: prelude:: { pymodule, PyModule , PyResult , Python } ;
4
4
5
5
#[ pymodule]
@@ -14,13 +14,18 @@ fn rust_ext(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
14
14
x *= a;
15
15
}
16
16
17
+ // complex example
18
+ fn conj ( x : ArrayViewD < ' _ , c64 > ) -> ArrayD < c64 > {
19
+ x. map ( |c| c. conj ( ) )
20
+ }
21
+
17
22
// wrapper of `axpy`
18
23
#[ pyfn( m, "axpy" ) ]
19
24
fn axpy_py < ' py > (
20
25
py : Python < ' py > ,
21
26
a : f64 ,
22
- x : PyReadonlyArrayDyn < f64 > ,
23
- y : PyReadonlyArrayDyn < f64 > ,
27
+ x : PyReadonlyArrayDyn < ' _ , f64 > ,
28
+ y : PyReadonlyArrayDyn < ' _ , f64 > ,
24
29
) -> & ' py PyArrayDyn < f64 > {
25
30
let x = x. as_array ( ) ;
26
31
let y = y. as_array ( ) ;
@@ -29,10 +34,15 @@ fn rust_ext(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
29
34
30
35
// wrapper of `mult`
31
36
#[ pyfn( m, "mult" ) ]
32
- fn mult_py ( _py : Python < ' _ > , a : f64 , x : & PyArrayDyn < f64 > ) -> PyResult < ( ) > {
37
+ fn mult_py ( a : f64 , x : & PyArrayDyn < f64 > ) {
33
38
let x = unsafe { x. as_array_mut ( ) } ;
34
39
mult ( a, x) ;
35
- Ok ( ( ) )
40
+ }
41
+
42
+ // wrapper of `conj`
43
+ #[ pyfn( m, "conj" ) ]
44
+ fn conj_py < ' py > ( py : Python < ' py > , x : PyReadonlyArrayDyn < ' _ , c64 > ) -> & ' py PyArrayDyn < c64 > {
45
+ conj ( x. as_array ( ) ) . into_pyarray ( py)
36
46
}
37
47
38
48
Ok ( ( ) )
0 commit comments