11use ndarray:: { ArrayD , ArrayViewD , ArrayViewMutD } ;
2- use numpy:: { IntoPyArray , PyArrayDyn , PyReadonlyArrayDyn } ;
2+ use numpy:: { c64 , IntoPyArray , PyArrayDyn , PyReadonlyArrayDyn } ;
33use pyo3:: prelude:: { pymodule, PyModule , PyResult , Python } ;
44
55#[ pymodule]
@@ -14,13 +14,18 @@ fn rust_ext(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
1414 x *= a;
1515 }
1616
17+ // complex example
18+ fn conj ( x : ArrayViewD < ' _ , c64 > ) -> ArrayD < c64 > {
19+ x. map ( |c| c. conj ( ) )
20+ }
21+
1722 // wrapper of `axpy`
1823 #[ pyfn( m, "axpy" ) ]
1924 fn axpy_py < ' py > (
2025 py : Python < ' py > ,
2126 a : f64 ,
22- x : PyReadonlyArrayDyn < f64 > ,
23- y : PyReadonlyArrayDyn < f64 > ,
27+ x : PyReadonlyArrayDyn < ' _ , f64 > ,
28+ y : PyReadonlyArrayDyn < ' _ , f64 > ,
2429 ) -> & ' py PyArrayDyn < f64 > {
2530 let x = x. as_array ( ) ;
2631 let y = y. as_array ( ) ;
@@ -29,10 +34,15 @@ fn rust_ext(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
2934
3035 // wrapper of `mult`
3136 #[ pyfn( m, "mult" ) ]
32- fn mult_py ( _py : Python < ' _ > , a : f64 , x : & PyArrayDyn < f64 > ) -> PyResult < ( ) > {
37+ fn mult_py ( a : f64 , x : & PyArrayDyn < f64 > ) {
3338 let x = unsafe { x. as_array_mut ( ) } ;
3439 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)
3646 }
3747
3848 Ok ( ( ) )
0 commit comments