Strahler stream order in Python

1   Introduction

Implement the Strahler stream order.

2   Python code

from osgeo import ogr
import matplotlib.pyplot as plt

def find_head_lines(lines):
    head_idx = []

    num_lines = len(lines)
    for i in range(num_lines):
        line = lines[i]
        first_point = line[0]

        has_upstream = False

        for j in range(num_lines):
            if j == i:
                continue
            line = lines[j]
            last_point = line[len(line)-1]

            if first_point == last_point:
                has_upstream = True

        if not has_upstream:
            head_idx.append(i)

    return head_idx


def find_prev_lines(curr_idx, lines):
    prev_idx = []

    num_lines = len(lines)

    line = lines[curr_idx]
    first_point = line[0]

    for i in range(num_lines):
        if i == curr_idx:
            continue
        line = lines[i]
        last_point = line[len(line)-1]

        if first_point == last_point:
            prev_idx.append(i)

    return prev_idx


def find_next_line(curr_idx, lines):
    num_lines = len(lines)

    line = lines[curr_idx]
    last_point = line[len(line)-1]

    next_idx = None

    for i in range(num_lines):
        if i == curr_idx:
            continue
        line = lines[i]
        first_point = line[0]

        if last_point == first_point:
            next_idx = i
            break

    return next_idx


def find_sibling_line(curr_idx, lines):
    num_lines = len(lines)

    line = lines[curr_idx]
    last_point = line[len(line)-1]

    sibling_idx = None

    for i in range(num_lines):
        if i == curr_idx:
            continue
        line = lines[i]
        last_point2 = line[len(line)-1]

        if last_point == last_point2:
            sibling_idx = i
            break

    return sibling_idx


shapefile_path = "p:/courses/python/streamnetwork.shp"

file = ogr.Open(shapefile_path, 1)
lyr = file.GetLayer(0)

num_features = lyr.GetFeatureCount()

lines = []

field_name = "Strahler"

for i in range(num_features):
    feat = lyr.GetFeature(i)
    feat.SetField(field_name, 0)
    lyr.SetFeature(feat)
    geom = feat.geometry()
    line = geom.GetPoints()
    lines.append(line)

head_idx = find_head_lines(lines)

for idx in head_idx:
    curr_idx = idx
    curr_feat = lyr.GetFeature(curr_idx)
    curr_ord = 1
    curr_feat.SetField(field_name, curr_ord) # head line always 1
    lyr.SetFeature(curr_feat)

    while True:
        next_idx = find_next_line(curr_idx, lines)
        if not next_idx:
            break
        next_feat = lyr.GetFeature(next_idx)
        next_ord = next_feat.GetField(field_name)

        sibl_idx = find_sibling_line(curr_idx, lines)
        if sibl_idx is not None:
            sibl_feat = lyr.GetFeature(sibl_idx)
            sibl_ord = sibl_feat.GetField(field_name)

            if sibl_ord:
                if sibl_ord > curr_ord:
                    break
                elif sibl_ord < curr_ord:
                    if next_ord == curr_ord:
                        break
                else:
                    curr_ord += 1

        next_feat.SetField(field_name, curr_ord)
        lyr.SetFeature(next_feat)

        curr_idx = next_idx