diff --git a/setup.py b/setup.py index b041648..8fe5916 100644 --- a/setup.py +++ b/setup.py @@ -5,8 +5,11 @@ import numpy astar_module = Extension( - 'pyastar2d.astar', sources=['src/cpp/astar.cpp'], - include_dirs=[numpy.get_include()], # for numpy/arrayobject.h + 'pyastar2d.astar', sources=['src/cpp/astar.cpp', 'src/cpp/experimental_heuristics.cpp'], + include_dirs=[ + numpy.get_include(), # for numpy/arrayobject.h + 'src/cpp' # for experimental_heuristics.h + ], extra_compile_args=["-O3", "-Wall", "-shared", "-fpic"], ) diff --git a/src/cpp/astar.cpp b/src/cpp/astar.cpp index 11131eb..94ab3a1 100644 --- a/src/cpp/astar.cpp +++ b/src/cpp/astar.cpp @@ -4,6 +4,7 @@ #include #include #include +#include const float INF = std::numeric_limits::infinity(); @@ -49,13 +50,15 @@ static PyObject *astar(PyObject *self, PyObject *args) { int start; int goal; int diag_ok; + int heuristic_override; if (!PyArg_ParseTuple( - args, "Oiiiii", // i = int, O = object + args, "Oiiiiii", // i = int, O = object &weights_object, &h, &w, &start, &goal, - &diag_ok)) + &diag_ok, &heuristic_override + )) return NULL; float* weights = (float*) weights_object->data; @@ -73,6 +76,13 @@ static PyObject *astar(PyObject *self, PyObject *args) { nodes_to_visit.push(start_node); int* nbrs = new int[8]; + + int goal_i = goal / w; + int goal_j = goal % w; + int start_i = start / w; + int start_j = start % w; + + heuristic_ptr heuristic_func = select_heuristic(heuristic_override); while (!nodes_to_visit.empty()) { // .top() doesn't actually remove the node @@ -104,13 +114,16 @@ static PyObject *astar(PyObject *self, PyObject *args) { float new_cost = costs[cur.idx] + weights[nbrs[i]]; if (new_cost < costs[nbrs[i]]) { // estimate the cost to the goal based on legal moves - if (diag_ok) { - heuristic_cost = linf_norm(nbrs[i] / w, nbrs[i] % w, - goal / w, goal % w); - } - else { - heuristic_cost = l1_norm(nbrs[i] / w, nbrs[i] % w, - goal / w, goal % w); + // Get the heuristic method to use + if (heuristic_override == DEFAULT) { + if (diag_ok) { + heuristic_cost = linf_norm(nbrs[i] / w, nbrs[i] % w, goal_i, goal_j); + } else { + heuristic_cost = l1_norm(nbrs[i] / w, nbrs[i] % w, goal_i, goal_j); + } + } else { + heuristic_cost = heuristic_func( + nbrs[i] / w, nbrs[i] % w, goal_i, goal_j, start_i, start_j); } // paths with lower expected cost are explored first diff --git a/src/cpp/experimental_heuristics.cpp b/src/cpp/experimental_heuristics.cpp new file mode 100644 index 0000000..c6a5a17 --- /dev/null +++ b/src/cpp/experimental_heuristics.cpp @@ -0,0 +1,42 @@ +// Please note below heuristics are experimental and only for pretty lines. +// They may not take the shortest path and require additional cpu cycles. + +#include +#include +#include + + +heuristic_ptr select_heuristic(int h) { + switch (h) { + case ORTHOGONAL_X: + return orthogonal_x; + case ORTHOGONAL_Y: + return orthogonal_y; + default: + return NULL; + } +} + +// Orthogonal x (moves by x first, then half way by y) +float orthogonal_x(int i0, int j0, int i1, int j1, int i2, int j2) { + int di = std::abs(i0 - i1); + int dim = std::abs(i1 - i2); + int djm = std::abs(j1 - j2); + if (di > (dim * 0.5)) { + return di + djm; + } else { + return std::abs(j0 - j1); + } +} + +// Orthogonal y (moves by y first, then half way by x) +float orthogonal_y(int i0, int j0, int i1, int j1, int i2, int j2) { + int dj = std::abs(j0 - j1); + int djm = std::abs(j1 - j2); + int dim = std::abs(i1 - i2); + if (dj > (djm * 0.5)) { + return dj + dim; + } else { + return std::abs(i0 - i1); + } +} diff --git a/src/cpp/experimental_heuristics.h b/src/cpp/experimental_heuristics.h new file mode 100644 index 0000000..98c0562 --- /dev/null +++ b/src/cpp/experimental_heuristics.h @@ -0,0 +1,20 @@ +// Please note below heuristics are experimental and only for pretty lines. +// They may not take the shortest path and require additional cpu cycles. + +#ifndef EXPERIMENTAL_HEURISTICS_H_ +#define EXPERIMENTAL_HEURISTICS_H_ + + +enum Heuristic { DEFAULT, ORTHOGONAL_X, ORTHOGONAL_Y }; + +typedef float (*heuristic_ptr)(int, int, int, int, int, int); + +heuristic_ptr select_heuristic(int); + +// Orthogonal x (moves by x first, then half way by y) +float orthogonal_x(int, int, int, int, int, int); + +// Orthogonal y (moves by y first, then half way by x) +float orthogonal_y(int, int, int, int, int, int); + +#endif diff --git a/src/pyastar2d/__init__.py b/src/pyastar2d/__init__.py index 3579838..644dc48 100644 --- a/src/pyastar2d/__init__.py +++ b/src/pyastar2d/__init__.py @@ -1,2 +1,2 @@ -from pyastar2d.astar_wrapper import astar_path -__all__ = ["astar_path"] +from pyastar2d.astar_wrapper import astar_path, Heuristic +__all__ = ["astar_path", "Heuristic"] diff --git a/src/pyastar2d/astar_wrapper.py b/src/pyastar2d/astar_wrapper.py index 95c5827..0d36f95 100644 --- a/src/pyastar2d/astar_wrapper.py +++ b/src/pyastar2d/astar_wrapper.py @@ -1,6 +1,7 @@ import ctypes import numpy as np import pyastar2d.astar +from enum import IntEnum from typing import Optional, Tuple @@ -19,14 +20,32 @@ ctypes.c_int, # start index in flattened grid ctypes.c_int, # goal index in flattened grid ctypes.c_bool, # allow diagonal + ctypes.c_int, # heuristic_override ] +class Heuristic(IntEnum): + """The supported heuristics.""" + + DEFAULT = 0 + ORTHOGONAL_X = 1 + ORTHOGONAL_Y = 2 def astar_path( weights: np.ndarray, start: Tuple[int, int], goal: Tuple[int, int], - allow_diagonal: bool = False) -> Optional[np.ndarray]: + allow_diagonal: bool = False, + heuristic_override: Heuristic = Heuristic.DEFAULT) -> Optional[np.ndarray]: + """ + Run astar algorithm on 2d weights. + + param np.ndarray weights: A grid of weights e.g. np.ones((10, 10), dtype=np.float32) + param Tuple[int, int] start: (i, j) + param Tuple[int, int] goal: (i, j) + param bool allow_diagonal: Whether to allow diagonal moves + param Heuristic heuristic_override: Override heuristic, see Heuristic(IntEnum) + + """ assert weights.dtype == np.float32, ( f"weights must have np.float32 data type, but has {weights.dtype}" ) @@ -49,5 +68,6 @@ def astar_path( path = pyastar2d.astar.astar( weights.flatten(), height, width, start_idx, goal_idx, allow_diagonal, + int(heuristic_override) ) return path diff --git a/tests/test_astar.py b/tests/test_astar.py index 59904a3..deb8d4f 100644 --- a/tests/test_astar.py +++ b/tests/test_astar.py @@ -5,6 +5,8 @@ import sys import pyastar2d +from pyastar2d import Heuristic + def test_small(): weights = np.array([[1, 3, 3, 3, 3], @@ -125,3 +127,19 @@ def test_bad_weights_dtype(): with pytest.raises(AssertionError) as exc: pyastar2d.astar_path(weights, (0, 0), (2, 2)) assert "float64" in exc.value.args[0] + + +def test_orthogonal_x(): + weights = np.ones((5, 5), dtype=np.float32) + path = pyastar2d.astar_path(weights, (0, 0), (4, 4), allow_diagonal=False, heuristic_override=Heuristic.ORTHOGONAL_X) + expected = np.array([[0, 0], [1, 0], [2, 0], [2, 1], [2, 2], [2, 3], [2, 4], [3, 4], [4, 4]]) + + assert np.all(path == expected) + + +def test_orthogonal_y(): + weights = np.ones((5, 5), dtype=np.float32) + path = pyastar2d.astar_path(weights, (0, 0), (4, 4), allow_diagonal=False, heuristic_override=Heuristic.ORTHOGONAL_Y) + expected = np.array([[0, 0], [0, 1], [0, 2], [1, 2], [2, 2], [3, 2], [4, 2], [4, 3], [4, 4]]) + + assert np.all(path == expected)