pyo3

Module num_complex

Source
Expand description

Conversions to and from num-complexComplex<f32> and Complex<f64>.

num-complex’ Complex supports more operations than PyO3’s PyComplex and can be used with the rest of the Rust ecosystem.

§Setup

To use this feature, add this to your Cargo.toml:

[dependencies]
# change * to the latest versions
num-complex = "*"
pyo3 = { version = "0.23.2", features = ["num-complex"] }

Note that you must use compatible versions of num-complex and PyO3. The required num-complex version may vary based on the version of PyO3.

§Examples

Using num-complex and nalgebra to create a pyfunction that calculates the eigenvalues of a 2x2 matrix.

use nalgebra::base::{dimension::Const, Matrix};
use num_complex::Complex;
use pyo3::prelude::*;

type T = Complex<f64>;

#[pyfunction]
fn get_eigenvalues(m11: T, m12: T, m21: T, m22: T) -> Vec<T> {
    let mat = Matrix::<T, Const<2>, Const<2>, _>::new(m11, m12, m21, m22);

    match mat.eigenvalues() {
        Some(e) => e.data.as_slice().to_vec(),
        None => vec![],
    }
}

#[pymodule]
fn my_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(get_eigenvalues, m)?)?;
    Ok(())
}

Python code:

from my_module import get_eigenvalues

m11 = complex(0,-1)
m12 = complex(1,0)
m21 = complex(2,-1)
m22 = complex(-1,0)

result = get_eigenvalues(m11,m12,m21,m22)
assert result == [complex(1,-1), complex(-2,0)]