Merge pull request #35 from peterchenadded/master

added changes for #29

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"],
 )
 

src/cpp/astar.cpp

@@ -4,6 +4,7 @@
 #include <Python.h>
 #include <numpy/arrayobject.h>
 #include <iostream>
+#include <experimental_heuristics.h>
 
 
 const float INF = std::numeric_limits<float>::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

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 <cmath>
+#include <cstddef>
+#include <experimental_heuristics.h>
+
+
+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);
+  }
+}

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

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"]

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

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)