unit_test.py 1.99 KB
Newer Older
Larkin Heintzman's avatar
Larkin Heintzman committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74
#!/usr/bin/env python

"""
boostrtrees.pyx and Rtree test code

designed to be run-able with py.test
"""
import pytest
import numpy as np
from boostrtrees import RTree
import math


def test_size():
    rec_ptr = RTree()
    assert(rec_ptr.size() == 0)
    rec_ptr.insert_point(1, 2, 10)
    assert(rec_ptr.size() == 1)
    rec_ptr.insert_point(3, 4, 10)
    assert (rec_ptr.size() == 2)
    rec_ptr.insert_point(3, 4, 10)
    assert (rec_ptr.size() == 3)


def test_knn():
    rec_ptr = RTree()
    rec_ptr.insert_point(1, 2, 1)
    rec_ptr.insert_point(3, 4, 2)
    rec_ptr.insert_point(5, 6, 3)
    rec_ptr.insert_point(8, 9, 4)
    assert (rec_ptr.knn(5, 6, 2) == [3, 2])


def test_bounds():
    rec_ptr = RTree()
    rec_ptr.insert_point(3, 5, 10)
    rec_ptr.insert_point(6, 7, 10)
    assert rec_ptr.bounds() == {'min_x': 3.0, 'max_x': 6.0, 'min_y': 5.0, 'max_y': 7.0}


def test_insert_numpy_points_size():
    rec_ptr = RTree()
    pts = np.array([(1.0, 2.0, 10), (1.0, 4.0, 11.0),
                    (3.0, 5.0, 12), (5.0, 7.0, 13)])
    rec_ptr.insert_points(pts)
    assert (rec_ptr.size() == 4)


def test_insert_numpy_points_bounds():
    rec_ptr = RTree()
    pts = np.array([(1.0, 2.0, 10), (1.0, 4.0, 11.0),
                    (3.0, 5.0, 12), (5.0, 7.0, 13)])
    rec_ptr.insert_points(pts)
    assert rec_ptr.bounds() == {'min_x': 1.0, 'max_x': 5.0, 'min_y': 2.0, 'max_y': 7.0}


def test_insert_numpy_points_withquery():
    rec_ptr = RTree()
    pts = np.array([(1.0, 2.0, 10), (1.0, 4.0, 11.0),
                    (3.0, 5.0, 12), (5.0, 7.0, 13)])
    rec_ptr.insert_points(pts)
    nn = rec_ptr.knn(0.0, 0.0, 1)
    assert nn == [10]


def test_minimum_distance():
    rec_ptr = RTree()
    pts = np.array([(1.0, 2.0, 10), (1.0, 4.0, 11.0),
                    (3.0, 5.0, 12), (5.0, 7.0, 13)])
    rec_ptr.insert_points(pts)
    dist = rec_ptr.min_distance(0.0, 0.0)
    assert math.fabs(dist - math.sqrt(5.0)) < 10e-5
    dist = rec_ptr.min_distance(7.0, 9.0)
    assert math.fabs(dist - math.sqrt(8.0)) < 10e-5