refactoring of rail connection functions

flatland/envs/grid4_generators_utils.py

@@ -14,7 +14,7 @@ from flatland.core.grid.grid_utils import Vec2dOperations as Vec2d
 from flatland.core.transition_map import GridTransitionMap, RailEnvTransitions
 
 
-def connect_basic_operation(
+def connect_rail(
     rail_trans: RailEnvTransitions,
     grid_map: GridTransitionMap,
     start: IntVector2D,
@@ -26,9 +26,20 @@ def connect_basic_operation(
     forbidden_cells=None
 ) -> IntVector2DArray:
     """
-    Creates a new path [start,end] in `grid_map.grid`, based on rail_trans, and
+        Creates a new path [start,end] in `grid_map.grid`, based on rail_trans, and
     returns the path created as a list of positions.
+    :param rail_trans:
+    :param grid_map:
+    :param start:
+    :param end:
+    :param flip_start_node_trans:
+    :param flip_end_node_trans:
+    :param nice:
+    :param a_star_distance_function:
+    :param forbidden_cells:
+    :return:
     """
+
     # in the worst case we will need to do a A* search, so we might as well set that up
     path: IntVector2DArray = a_star(grid_map, start, end, a_star_distance_function, nice, forbidden_cells)
     if len(path) < 2:
@@ -79,7 +90,7 @@ def connect_basic_operation(
     return path
 
 
-def connect_line(rail_trans, grid_map, start, end, openend=False):
+def connect_straigt_line(rail_trans, grid_map, start, end, openend=False):
     """
     Generates a straight rail line from start cell to end cell.
     Diagonal lines are not allowed
@@ -126,47 +137,3 @@ def connect_line(rail_trans, grid_map, start, end, openend=False):
         path.append(current_cell)
     return path
 
-def connect_rail(rail_trans: RailEnvTransitions, grid_map: GridTransitionMap,
-                 start: IntVector2D, end: IntVector2D,
-                 a_star_distance_function: IntVector2DDistance = Vec2d.get_manhattan_distance) -> IntVector2DArray:
-    """
-
-    :param rail_trans:
-    :param grid_map:
-    :param start:
-    :param end:
-    :param a_star_distance_function:
-    :return:
-    """
-    return connect_basic_operation(rail_trans, grid_map, start, end, True, True, True, a_star_distance_function)
-
-
-def connect_cities(rail_trans: RailEnvTransitions, grid_map: GridTransitionMap,
-                   start: IntVector2D, end: IntVector2D,
-                   a_star_distance_function: IntVector2DDistance = Vec2d.get_manhattan_distance,
-                   forbidden_cells=None) -> IntVector2DArray:
-    """
-
-    :param rail_trans:
-    :param grid_map:
-    :param start:
-    :param end:
-    :param forbidden_cells:
-    :param a_star_distance_function:
-    :return:
-    """
-    return connect_basic_operation(rail_trans, grid_map, start, end, False, False, False, a_star_distance_function,
-                                   forbidden_cells)
-
-def connect_straigt_line(rail_trans: RailEnvTransitions, grid_map: GridTransitionMap, start: IntVector2D,
-                         end: IntVector2D, openend=False) -> IntVector2DArray:
-    """
-
-    :param rail_trans:
-    :param grid_map:
-    :param start:
-    :param end:
-    :param openend:
-    :return:
-    """
-    return connect_line(rail_trans, grid_map, start, end, openend)

flatland/envs/rail_generators.py

@@ -11,7 +11,7 @@ from flatland.core.grid.grid_utils import distance_on_rail, direction_to_city, I
     Vec2dOperations
 from flatland.core.grid.rail_env_grid import RailEnvTransitions
 from flatland.core.transition_map import GridTransitionMap
-from flatland.envs.grid4_generators_utils import connect_rail, connect_cities, connect_straigt_line
+from flatland.envs.grid4_generators_utils import connect_rail, connect_straigt_line
 
 RailGeneratorProduct = Tuple[GridTransitionMap, Optional[Dict]]
 RailGenerator = Callable[[int, int, int, int], RailGeneratorProduct]
@@ -125,7 +125,9 @@ def complex_rail_generator(nr_start_goal=1,
                 # we might as well give up at this point
                 break
 
-            new_path = connect_rail(rail_trans, grid_map, start, goal)
+            new_path = connect_rail(rail_trans, grid_map, start, goal, flip_start_node_trans=True,
+                                    flip_end_node_trans=True, nice=True,
+                                    forbidden_cells=None)
             if len(new_path) >= 2:
                 nr_created += 1
                 start_goal.append([start, goal])
@@ -150,7 +152,10 @@ def complex_rail_generator(nr_start_goal=1,
                     break
             if not all_ok:
                 break
-            new_path = connect_rail(rail_trans, grid_map, start, goal)
+            new_path = connect_rail(rail_trans, grid_map, start, goal, flip_start_node_trans=True,
+                                    flip_end_node_trans=True, nice=True,
+                                    forbidden_cells=None)
+
             if len(new_path) >= 2:
                 nr_created += 1
 
@@ -735,9 +740,10 @@ def sparse_rail_generator(max_num_cities: int = 5, grid_mode: bool = False, max_
                                 min_connection_dist = tmp_dist
                                 neighb_connection_point = tmp_in_connection_point
                                 neighbour_direction = dir
-                    new_line = connect_cities(rail_trans, grid_map, tmp_out_connection_point,
-                                              neighb_connection_point,
-                                              forbidden_cells=city_cells)
+                    new_line = connect_rail(rail_trans, grid_map, tmp_out_connection_point,
+                                            neighb_connection_point, flip_start_node_trans=False,
+                                            flip_end_node_trans=False, nice=True,
+                                            forbidden_cells=city_cells)
                     all_paths.extend(new_line)
 
         return all_paths

tests/test_flatland_core_grid4_generators_util.py

@@ -3,7 +3,7 @@ import numpy as np
 from flatland.core.grid.grid_utils import Vec2dOperations as Vec2d
 from flatland.core.grid.rail_env_grid import RailEnvTransitions
 from flatland.core.transition_map import GridTransitionMap
-from flatland.envs.grid4_generators_utils import connect_rail, connect_cities
+from flatland.envs.grid4_generators_utils import connect_rail
 
 
 def test_build_railway_infrastructure():
@@ -20,17 +20,17 @@ def test_build_railway_infrastructure():
 
     start_point = (1, 3)
     end_point = (1, 7)
-    connection_002 = connect_cities(rail_trans, grid_map, start_point, end_point, Vec2d.get_manhattan_distance)
+    connection_002 = connect_rail(rail_trans, grid_map, start_point, end_point, Vec2d.get_manhattan_distance)
     connection_002_expected = [(1, 3), (1, 4), (1, 5), (1, 6), (1, 7)]
 
     start_point = (6, 2)
     end_point = (6, 5)
-    connection_003 = connect_cities(rail_trans, grid_map, start_point, end_point, Vec2d.get_manhattan_distance)
+    connection_003 = connect_rail(rail_trans, grid_map, start_point, end_point, Vec2d.get_manhattan_distance)
     connection_003_expected = [(6, 2), (6, 3), (6, 4), (6, 5)]
 
     start_point = (7, 5)
     end_point = (8, 9)
-    connection_004 = connect_cities(rail_trans, grid_map, start_point, end_point, Vec2d.get_manhattan_distance)
+    connection_004 = connect_rail(rail_trans, grid_map, start_point, end_point, Vec2d.get_manhattan_distance)
     connection_004_expected = [(7, 5), (7, 6), (7, 7), (7, 8), (7, 9), (8, 9)]
 
     assert connection_001 == connection_001_expected, \