def separate_units_by_type(all_units):
    """Separate all_units to their respective unit type group."""
    immune = {}
    infection = {}
    unit_type_to_group = {
        'immune': immune,
        'infection': infection,
    }
    for unit_no, unit in all_units.items():
        group = unit_type_to_group[unit['type']]
        group[unit_no] = unit
    return immune, infection