test_flatland_core_transitions.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Tests for `flatland` package."""
import numpy as np

from flatland.core.transitions import RailEnvTransitions, Grid8Transitions
from flatland.envs.env_utils import validate_new_transition


def test_rotate_railenv_transition():
    rail_env_transitions = RailEnvTransitions()

    # remove whitespace in string; keep whitespace below for easier reading
    def rw(s):
        return s.replace(" ", "")

    # TODO test all cases
    transition_cycles = [
        # empty cell - Case 0
        [int('0000000000000000', 2), int('0000000000000000', 2), int('0000000000000000', 2),
         int('0000000000000000', 2)],
        # Case 1 - straight
        #     |
        #     |
        #     |
        [int(rw('1000 0000 0010 0000'), 2), int(rw('0000 0100 0000 0001'), 2)],
        # Case 1b (8)  - simple turn right
        #      _
        #     |
        #     |
        [
            int(rw('0100 0000 0000 0010'), 2),
            int(rw('0001 0010 0000 0000'), 2),
            int(rw('0000 1000 0001 0000'), 2),
            int(rw('0000 0000 0100 1000'), 2),
        ],
        # Case 1c (9)  - simple turn left
        #    _
        #     |
        #     |

        #                int('0001001000000000', 2),

        # Case 2 - simple left switch
        #  _ _|
        #     |
        #     |
        [
            int(rw('1001 0010 0010 0000'), 2),
            int(rw('0000 1100 0001 0001'), 2),
            int(rw('1000 0000 0110 1000'), 2),
            int(rw('0100 0100 0000 0011'), 2),
        ],
        # Case 2b (10) - simple right switch
        #     |
        #     |
        #     |
        #                int('1100000000100010', 2)]
        #                int('1000010000100001', 2),  # Case 3 - diamond drossing
        #                int('1001011000100001', 2),  # Case 4 - single slip
        #                int('1100110000110011', 2),  # Case 5 - double slip
        #                int('0101001000000010', 2),  # Case 6 - symmetrical
        #                int('0010000000000000', 2),  # Case 7 - dead end

    ]

    for index, cycle in enumerate(transition_cycles):
        for i in range(4):
            actual_transition = rail_env_transitions.rotate_transition(cycle[0], i * 90)
            expected_transition = cycle[i % len(cycle)]
            try:
                assert actual_transition == expected_transition, \
                    "Case {}: rotate_transition({}, {}) should equal {} but was {}." \
                        .format(i, cycle[0], i, expected_transition, actual_transition)
            except Exception as e:
                print("expected:")
                rail_env_transitions.print(expected_transition)
                print("actual:")
                rail_env_transitions.print(actual_transition)

                raise e


def test_is_valid_railenv_transitions():
    rail_env_trans = RailEnvTransitions()
    transition_list = rail_env_trans.transitions

    for t in transition_list:
        assert (rail_env_trans.is_valid(t) is True)
        for i in range(3):
            rot_trans = rail_env_trans.rotate_transition(t, 90 * i)
            assert (rail_env_trans.is_valid(rot_trans) is True)

    assert (rail_env_trans.is_valid(int('1111111111110010', 2)) is False)
    assert (rail_env_trans.is_valid(int('1001111111110010', 2)) is False)
    assert (rail_env_trans.is_valid(int('1001111001110110', 2)) is False)


def test_adding_new_valid_transition():
    rail_trans = RailEnvTransitions()
    rail_array = np.zeros(shape=(15, 15), dtype=np.uint16)

    # adding straight
    assert (validate_new_transition(rail_trans, rail_array, (4, 5), (5, 5), (6, 5), (10, 10)) is True)

    # adding valid right turn
    assert (validate_new_transition(rail_trans, rail_array, (5, 4), (5, 5), (5, 6), (10, 10)) is True)
    # adding valid left turn
    assert (validate_new_transition(rail_trans, rail_array, (5, 6), (5, 5), (5, 6), (10, 10)) is True)

    # adding invalid turn
    rail_array[(5, 5)] = rail_trans.transitions[2]
    assert (validate_new_transition(rail_trans, rail_array, (4, 5), (5, 5), (5, 6), (10, 10)) is False)

    # should create #4 -> valid
    rail_array[(5, 5)] = rail_trans.transitions[3]
    assert (validate_new_transition(rail_trans, rail_array, (4, 5), (5, 5), (5, 6), (10, 10)) is True)

    # adding invalid turn
    rail_array[(5, 5)] = rail_trans.transitions[7]
    assert (validate_new_transition(rail_trans, rail_array, (4, 5), (5, 5), (5, 6), (10, 10)) is False)

    # test path start condition
    rail_array[(5, 5)] = rail_trans.transitions[0]
    assert (validate_new_transition(rail_trans, rail_array, None, (5, 5), (5, 6), (10, 10)) is True)

    # test path end condition
    rail_array[(5, 5)] = rail_trans.transitions[0]
    assert (validate_new_transition(rail_trans, rail_array, (5, 4), (5, 5), (6, 5), (6, 5)) is True)


def test_valid_railenv_transitions():
    rail_env_trans = RailEnvTransitions()

    # directions:
    #            'N': 0
    #            'E': 1
    #            'S': 2
    #            'W': 3

    for i in range(2):
        assert (rail_env_trans.get_transitions(
            int('1100110000110011', 2), i) == (1, 1, 0, 0))
        assert (rail_env_trans.get_transitions(
            int('1100110000110011', 2), 2 + i) == (0, 0, 1, 1))

    no_transition_cell = int('0000000000000000', 2)

    for i in range(4):
        assert (rail_env_trans.get_transitions(
            no_transition_cell, i) == (0, 0, 0, 0))

    # Facing south, going south
    north_south_transition = rail_env_trans.set_transitions(no_transition_cell, 2, (0, 0, 1, 0))
    assert (rail_env_trans.set_transition(
        north_south_transition, 2, 2, 0) == no_transition_cell)
    assert (rail_env_trans.get_transition(
        north_south_transition, 2, 2))

    # Facing north, going east
    south_east_transition = \
        rail_env_trans.set_transition(no_transition_cell, 0, 1, 1)
    assert (rail_env_trans.get_transition(
        south_east_transition, 0, 1))

    # The opposite transitions are not feasible
    assert (not rail_env_trans.get_transition(
        north_south_transition, 2, 0))
    assert (not rail_env_trans.get_transition(
        south_east_transition, 2, 1))

    east_west_transition = rail_env_trans.rotate_transition(north_south_transition, 90)
    north_west_transition = rail_env_trans.rotate_transition(south_east_transition, 180)

    # Facing west, going west
    assert (rail_env_trans.get_transition(
        east_west_transition, 3, 3))
    # Facing south, going west
    assert (rail_env_trans.get_transition(
        north_west_transition, 2, 3))

    assert (south_east_transition == rail_env_trans.rotate_transition(
        south_east_transition, 360))


def test_diagonal_transitions():
    diagonal_trans_env = Grid8Transitions([])

    # Facing north, going north-east
    south_northeast_transition = int('01000000' + '0' * 8 * 7, 2)
    assert (diagonal_trans_env.get_transitions(
        south_northeast_transition, 0) == (0, 1, 0, 0, 0, 0, 0, 0))

    # Allowing transition from north to southwest: Facing south, going SW
    north_southwest_transition = \
        diagonal_trans_env.set_transitions(int('0' * 64, 2), 4, (0, 0, 0, 0, 0, 1, 0, 0))

    assert (diagonal_trans_env.rotate_transition(
        south_northeast_transition, 180) == north_southwest_transition)