from typing import List, Tuple  
import numpy as np  
from dataclasses import dataclass  
from functools import cached_property  
Point = np.array  
@dataclass  
class Polygon:  
    x: np.array  
    y: np.array  
    @cached_property  
    def center(self) -> Point: ...  
    def area(self) -> float: ...  
def find_close_polygons(polygon_subset: List[Polygon], point: Point, max_dist: float) -> List[Polygon]:  
    ...  
def select_best_polygon(polygon_sets: List[Tuple[Point, List[Polygon]]]) -> List[Tuple[Point, Polygon]]:  
    ...  
def main(polygons: List[Polygon], points: np.ndarray) -> List[Tuple[Point, Polygon]]:  
    ...

pip install py-spy

# measure.py  
import time  
import poly_match  
import os  

# Reduce noise, actually improve perf in our case.  
os.environ["OPENBLAS_NUM_THREADS"] = "1"  

polygons, points = poly_match.generate_example() 

# We are going to increase this as the code gets faster and faster.  
NUM_ITER = 10  

t0 = time.perf_counter()  
for _ in range(NUM_ITER):  
    poly_match.main(polygons, points)  
t1 = time.perf_counter()  

took = (t1 - t0) / NUM_ITER  
print(f"Took and avg of {took * 1000:.2f}ms per iteration")

$ python measure.py  
Took an avg of 293.41ms per iteration

$ conda create -n pypyenv -c conda-forge pypy numpy && conda activate pypyenv  
$ pypy measure_with_warmup.py  
Took an avg of 1495.81ms per iteration

$ py-spy record --native -o profile.svg -- python measure.py  

py-spy> Sampling process 100 times a second. Press Control-C to exit.  

Took an avg of 365.43ms per iteration  

py-spy> Stopped sampling because process exited  
py-spy> Wrote flamegraph data to 'profile.svg'. Samples: 391 Errors: 0

$ python -m cProfile measure.py  
Took an avg of 546.47ms per iteration  
         7551778 function calls (7409483 primitive calls) in 7.806 seconds  
         ...

def find_close_polygons(  
    polygon_subset: List[Polygon], point: np.array, max_dist: float  
) -> List[Polygon]:  
    close_polygons = []  
    for poly in polygon_subset:  
        if np.linalg.norm(poly.center - point) < max_dist:  
            close_polygons.append(poly)  
    return close_polygons

mkdir poly_match_rs && cd "$_"  
pip install maturin  
maturin init --bindings pyo3  
maturin develop

use pyo3::prelude::*;  
#[pyfunction]  
fn find_close_polygons() -> PyResult<()> {  
    Ok(())  
}  
#[pymodule]  
fn poly_match_rs(_py: Python, m: &PyModule) -> PyResult<()> {  
    m.add_function(wrap_pyfunction!(find_close_polygons, m)?)?;  
    Ok(())  
}

>>> poly_match_rs.find_close_polygons(polygons, point, max_dist)  
E TypeError: poly_match_rs.poly_match_rs.find_close_polygons() takes no arguments (3 given)

#[pyfunction]  
fn find_close_polygons(polygons: Vec<PyObject>, point: PyObject, max_dist: f64) -> PyResult<Vec<PyObject>> {  
    Ok(vec![])  
}

#[pyfunction]  
fn find_close_polygons(polygons: Vec<PyObject>, point: PyObject, max_dist: f64) -> PyResult<Vec<PyObject>> {  
    let mut close_polygons = vec![];  


    for poly in polygons {  
        if norm(poly.center - point) < max_dist {  
            close_polygons.push(poly)  
        }  
    }  


    Ok(close_polygons)  
}

% maturin develop  
...  
error[E0609]: no field `center` on type `Py<PyAny>`  
 --> src/lib.rs:8:22  
  |  
8 |         if norm(poly.center - point) < max_dist {  
  |                      ^^^^^^ unknown field  
error[E0425]: cannot find function `norm` in this scope  
 --> src/lib.rs:8:12  
  |  
8 |         if norm(poly.center - point) < max_dist {  
  |            ^^^^ not found in this scope  
error: aborting due to 2 previous errors ] 58/59: poly_match_rs

# For Rust-native array operations.  
ndarray = "0.15"  
# For a `norm` function for arrays.  
ndarray-linalg = "0.16"    
# For accessing numpy-created objects, based on `ndarray`.  
numpy = "0.18"

use numpy::PyReadonlyArray1;  
#[pyfunction]  
fn find_close_polygons(  
    // An object which says "I have the GIL", so we can access Python-managed memory.  
    py: Python<'_>,  
    polygons: Vec<PyObject>,  
    // A reference to a numpy array we will be able to access.  
    point: PyReadonlyArray1<f64>,  
    max_dist: f64,  
) -> PyResult<Vec<PyObject>> {  
    // Convert to `ndarray::ArrayView1`, a fully operational native array.  
    let point = point.as_array();  
    ...  
}

/// Make the `norm` function available.  
use ndarray_linalg::Norm;  
assert_eq!((point.to_owned() - point).norm(), 0.); // Make the `norm` function available.

let center = poly  
  .getattr(py, "center")?                 // Python-style getattr, requires a GIL token (`py`).  
  .extract::<PyReadonlyArray1<f64>>(py)?  // Tell PyO3 what to convert the result to.  
  .as_array()                             // Like `point` before.  
  .to_owned();                            // We need one of the sides of the `-` to be "owned".

use pyo3::prelude::*;

use ndarray_linalg::Norm;
use numpy::PyReadonlyArray1;

#[pyfunction]
fn find_close_polygons(
    py: Python<'_>,
    polygons: Vec<PyObject>,
    point: PyReadonlyArray1<f64>,
    max_dist: f64,
) -> PyResult<Vec<PyObject>> {
    let mut close_polygons = vec![];
    let point = point.as_array();
    for poly in polygons {
        let center = poly
            .getattr(py, "center")?
            .extract::<PyReadonlyArray1<f64>>(py)?
            .as_array()
            .to_owned();

        if (center - point).norm() < max_dist {
            close_polygons.push(poly)
        }
    }

    Ok(close_polygons)
}

1.def find_close_polygons(  
2.    polygon_subset: List[Polygon], point: np.array, max_dist: float  
3.) -> List[Polygon]:  
4.    close_polygons = []  
5.    for poly in polygon_subset:  
6.        if np.linalg.norm(poly.center - point) < max_dist:  
7.            close_polygons.append(poly)  
8.  
9.    return close_polygons

1.$ (cd ./poly_match_rs/ && maturin develop)  
2.$ python measure.py  
3.Took an avg of 609.46ms per iteration

$ (cd ./poly_match_rs/ && maturin develop --release)  
$ python measure.py  
Took an avg of 23.44ms per iteration

1.# added to Cargo.toml  
2.[profile.release]  
3.debug = true       # Debug symbols for our profiler.  
4.lto = true         # Link-time optimization.  
5.codegen-units = 1  # Slower compilation but faster code.

$ py-spy record --native -o profile.svg -- python measure.py  
py-spy> Sampling process 100 times a second. Press Control-C to exit.

@dataclass  
class Polygon:  
    x: np.array  
    y: np.array  
    _area: float = None  
    @cached_property  
    def center(self) -> np.array:  
        centroid = np.array([self.x, self.y]).mean(axis=1)  
        return centroid  
    def area(self) -> float:  
        if self._area is None:  
            self._area = 0.5 * np.abs(  
                np.dot(self.x, np.roll(self.y, 1)) - np.dot(self.y, np.roll(self.x, 1))  
            )  
        return self._area

// `Array1` is a 1d array, and the `numpy` crate will play nicely with it.  
use ndarray::Array1;  


// `subclass` tells PyO3 to allow subclassing this in Python.  
#[pyclass(subclass)]  
struct Polygon {  
    x: Array1<f64>,  
    y: Array1<f64>,  
    center: Array1<f64>,  
}

use numpy::{PyArray1, PyReadonlyArray1, ToPyArray};  
#[pymethods]  
impl Polygon {  
    #[new]  
    fn new(x: PyReadonlyArray1<f64>, y: PyReadonlyArray1<f64>) -> Polygon {  
        let x = x.as_array();  
        let y = y.as_array();  
        let center = Array1::from_vec(vec![x.mean().unwrap(), y.mean().unwrap()]);  


        Polygon {  
            x: x.to_owned(),  
            y: y.to_owned(),  
            center,  
        }  
    }  
    // the `Py<..>` in the return type is a way of saying "an Object owned by Python".  
    #[getter]                 
    fn x(&self, py: Python<'_>) -> PyResult<Py<PyArray1<f64>>> {  
        Ok(self.x.to_pyarray(py).to_owned()) // Create a Python-owned, numpy version of `x`.  
    }  
    // Same for `y` and `center`.  
}

#[pymodule]
fn poly_match_rs(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<Polygon>()?; // new.
    m.add_function(wrap_pyfunction!(find_close_polygons, m)?)?;
    Ok(())
}

class Polygon(poly_match_rs.Polygon):
    _area: float = None

    def area(self) -> float:
        ...

#[pyfunction]
fn find_close_polygons(
    py: Python<'_>,
    polygons: Vec<Py<Polygon>>,             // References to Python-owned objects.
    point: PyReadonlyArray1<f64>,
    max_dist: f64,
) -> PyResult<Vec<Py<Polygon>>> {           // Return the same `Py` references, unmodified.
    let mut close_polygons = vec![];
    let point = point.as_array();
    for poly in polygons {
        let center = poly.borrow(py).center // Need to use the GIL (`py`) to borrow the underlying `Polygon`.
            .to_owned();

        if (center - point).norm() < max_dist {
            close_polygons.push(poly)
        }
    }

    Ok(close_polygons)
}

$ python measure.py
Took an avg of 6.29ms per iteration

let center = poly.borrow(py).center
    .to_owned();

if (center - point).norm() < max_dist { ... }

use ndarray_linalg::Scalar;  

let center = &poly.as_ref(py).borrow().center;  

if ((center[0] - point[0]).square() + (center[1] - point[1]).square()).sqrt() < max_dist {  
    close_polygons.push(poly)  
}

error[E0505]: cannot move out of `poly` because it is borrowed
  --> src/lib.rs:58:33
   |
55 |         let center = &poly.as_ref(py).borrow().center;
   |                       ------------------------
   |                       |
   |                       borrow of `poly` occurs here
   |                       a temporary with access to the borrow is created here ...
...
58 |             close_polygons.push(poly);
   |                                 ^^^^ move out of `poly` occurs here
59 |         }
60 |     }
   |     - ... and the borrow might be used here, when that temporary is dropped and runs the `Drop` code for type `PyRef`

let norm = {  
    let center = &poly.as_ref(py).borrow().center;  

    ((center[0] - point[0]).square() + (center[1] - point[1]).square()).sqrt()  
};  

if norm < max_dist {  
    close_polygons.push(poly)  
}

$ python measure.py  
Took an avg of 2.90ms per iteration

@dataclass  
class Polygon:  
    x: np.array  
    y: np.array  
    _area: float = None  

    @cached_property  
    def center(self) -> np.array:  
        centroid = np.array([self.x, self.y]).mean(axis=1)  
        return centroid  

    def area(self) -> float:  
        ...  

def find_close_polygons(  
    polygon_subset: List[Polygon], point: np.array, max_dist: float
) -> List[Polygon]:  
    close_polygons = []  
    for poly in polygon_subset:  
        if np.linalg.norm(poly.center - point) < max_dist:  
            close_polygons.append(poly)  

    return close_polygons  

# Rest of file (main, select_best_polygon).

import poly_match_rs  
from poly_match_rs import find_close_polygons  

class Polygon(poly_match_rs.Polygon):  
    _area: float = None  

    def area(self) -> float:  
        ...  

# Rest of file unchanged (main, select_best_polygon).

use pyo3::prelude::*;  

use ndarray::Array1;  
use ndarray_linalg::Scalar;  
use numpy::{PyArray1, PyReadonlyArray1, ToPyArray};  

#[pyclass(subclass)]  
struct Polygon {  
    x: Array1<f64>,  
    y: Array1<f64>,  
    center: Array1<f64>,  
}  

#[pymethods]  
impl Polygon {  
    #[new]  
    fn new(x: PyReadonlyArray1<f64>, y: PyReadonlyArray1<f64>) -> Polygon {  
        let x = x.as_array();  
        let y = y.as_array();  
        let center = Array1::from_vec(vec![x.mean().unwrap(), y.mean().unwrap()]);  

        Polygon {  
            x: x.to_owned(),  
            y: y.to_owned(),  
            center,  
        }  
    }  

    #[getter]  
    fn x(&self, py: Python<'_>) -> PyResult<Py<PyArray1<f64>>> {  
        Ok(self.x.to_pyarray(py).to_owned())  
    }  

    // Same for `y` and `center`.  
}  

#[pyfunction]  
fn find_close_polygons(  
    py: Python<'_>,  
    polygons: Vec<Py<Polygon>>,  
    point: PyReadonlyArray1<f64>,  
    max_dist: f64,  
) -> PyResult<Vec<Py<Polygon>>> {  
    let mut close_polygons = vec![];  
    let point = point.as_array();  
    for poly in polygons {  
        let norm = {  
            let center = &poly.as_ref(py).borrow().center;  

            ((center[0] - point[0]).square() + (center[1] - point[1]).square()).sqrt()  
        };  

        if norm < max_dist {  
            close_polygons.push(poly)  
        }  
    }  

    Ok(close_polygons)  
}  

#[pymodule]  
fn poly_match_rs(_py: Python, m: &PyModule) -> PyResult<()> {  
    m.add_class::<Polygon>()?;  
    m.add_function(wrap_pyfunction!(find_close_polygons, m)?)?;  
    Ok(())  
}