Skip to content

utils

get_locations_depth_first(hierarchy: pd.DataFrame) -> list[int]

Return location ids sorted by a depth first search of the hierarchy.

Locations at the same level are sorted alphabetically by name.

Source code in src/climate_data/diagnostics/utils.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def get_locations_depth_first(hierarchy: pd.DataFrame) -> list[int]:
    """Return location ids sorted by a depth first search of the hierarchy.

    Locations at the same level are sorted alphabetically by name.
    """

    def _get_locations(location: pd.Series):
        locs = [location.location_id]

        children = hierarchy[
            (hierarchy.parent_id == location.location_id)
            & (hierarchy.location_id != location.location_id)
        ]
        for child in children.sort_values("location_ascii_name").itertuples():
            locs.extend(_get_locations(child))
        return locs

    top_locs = hierarchy[hierarchy.location_id == hierarchy.parent_id]
    locations = []
    for top_loc in top_locs.sort_values("location_ascii_name").itertuples():
        locations.extend(_get_locations(top_loc))

    return locations