#------------------------------------------------------------------------------- # Name: StreamNet.py (Stream Network) # Purpose: Create an object tree representing a river network # # Steps: # Get arcpy.Arrays for all lines (provides all necessary points through indexing) # Use setChildren() to # - Find endpoints matching the current object's start point # - count the number of children (and hence determine the streamState) # - deal with FourWays (store and remove extras from further handling) # - determine Left and Right # - create the appropriate number of children # Use buildTree() to call setChildren() # - (More or less builds its own tracks and then drives on them) # - Starts at a new node, then calls setChildren # - setChildren creates 0 or more children # - buildTree() recurses down into any new children # - buildTree() is now at a new node, repeat # The tree is now available for informational methods # # # Author: William F. Mellman # # Created: 2017-04-13 # Copyright: (c) 2017 # Licence: GNU # # HISTORY: # 0.0.0 2017-04 Mellman Development # 1.0.0 2017-04-14 Mellman First tree traversed as StreamIDs # 1.0.1 2017-04-15 Mellman First tree built as objects # 1.0.2 2017-04-16 Mellman Find orphans # 1.0.3 2017-04-16 Mellman getObject, finish # 1.1.0 2017-05-16 Mellman Impolment Depth and Stream Order, pretty print # 1.2.0 2017-05-16 Mellman Write results to table # 1.2.1 2017-05-21 Mellman Misc # 1.2.2 2017-05-22 Mellman whereClause, combine search cursors, collapse traverseNremove to findOrphans # 1.3.0 2017-06-02 Mellman read other data (EdgeType & Enabled) and flag EdgeType transitions # 1.3.0 2017-06-02 Mellman read other data (EdgeType & Enabled) and flag EdgeType transitions # 1.3.1 2017-06-05 Mellman tweaks, mostly to print-out # 2.0.0 2017-06-05 Mellman Add checks and inputs to allow running from ArcMap Script Tool # 2.0.1 2017-07-06 Mellman Add come comments, find and fix typo in length calculation #------------------------------------------------------------------------------- Version = '2.0.1' #------------------------------------------------------------------------------- import arcpy, math, os, sys, datetime from LeftRight import IsLeftRight #------------------------------------------------------------------------------- # A R C M A P I N T E R A C T I O N #------------------------------------------------------------------------------- VERBOSE = True SUPERQUIET = False ArcMap = True StartTime = datetime.datetime.now() TimeText = str(StartTime.year) + str(StartTime.month) + str(StartTime.day) + str(StartTime.hour) + str(StartTime.minute) + str(StartTime.second) if ArcMap: fc = arcpy.GetParameterAsText(0) if (fc == "") or (fc == "#"): #fc = r"C:\Projects\StreamNet\StreamNet.gdb\hyd_streams" #fc = r"streams" # replace with stream layer path fc = r"streams" # replace with stream layer path # End fc StartTrunk = arcpy.GetParameterAsText(1) if (StartTrunk == "") or (StartTrunk == "#"): StartTrunk = 44424 # middle: 40651 StartTrunk = int(StartTrunk) # Used for output table name watershed = arcpy.GetParameterAsText(2) if (watershed == "") or (watershed == "#"): curDate = datetime.datetime watershed = "StreamNet_" + TimeText # End watershed # Don't use whereclause in ArcMap since we'll be using a selection whereClause = '' else: fc = r"streams" # replace with stream layer path StartTrunk = 44424 # middle: 40651 watershed = "StreamNet_" + TimeText dataFile = r"C:\Projects\StreamNet\StreamNet.data" whereData = '' # PRECONDITION: data file has comma separated, undelimited object ID's on one line, with no trailing blank lines. with open(dataFile,"r") as f: whereData = f.next() whereBase = '("EdgeType"= 1 OR "EdgeType" = 3) AND ("Enabled" = 1) AND "OBJECTID" IN (' whereEnd = ')' whereClause = whereBase + whereData + whereEnd # End if ArcMap #------------------------------------------------------------------------------- # G L O B A L S #------------------------------------------------------------------------------- # Output to table #workPath = os.path.dirname(fc) workPath = r"C:\Projects\StreamNet\StreamNet.gdb" arcpy.env.workspace = workPath OutTabName = watershed + 'Results' FullOutTabName = workPath + '/' + OutTabName TemplateName = 'StreamNetTemplate' # Assumed to be in workPath GDB # Output to text file outFileName = r"C:\Projects\StreamNet\StreamNet.output.txt" outFile = open(outFileName,'w') # OVERWRITE ! ! ! outFile.write("Tree:\n") outFile.close outFile = open(outFileName,'a') dataStore = {} outFields = ['StreamID','Left_ID','Right_ID','Parent_ID','StreamState','PathDepth','StreamOrder','BranchPath','LineLength','Source2Start','End2Mouth'] START = 0 SECOND = 1 END = -1 SECONDEND = -2 A = 0 B = 1 FIRST = 0 SECOND = 1 LEFT = 0 RIGHT = 1 EDGETYPE = 2 ENABLED = 3 Leaf = 0 OneChild = 1 Normal = 2 FourWay = 3 NodeCnt = {0:'Leaf',1:'OneChild',2:'Normal',3:'FourWay'} oSID = 0 oLeft = 1 oRight = 2 oParent = 3 oState = 4 oDepth = 5 oOrder = 6 oPath = 7 oLength = 8 oSource2Start = 9 oEnd2Mouth = 10 #------------------------------------------------------------------------------- # F U N C T I O N S #------------------------------------------------------------------------------- def ThreeWrite(line): global outFile print line outFile.write(line+"\n") arcpy.AddMessage(line) # End def ThreeWrite() def TwoWrite(line): print line, arcpy.AddMessage(line) # End def TwoWrite() def getOrder(a,b): # http://www.cotf.edu/ete/modules/waterq3/WQassess4b.html # https://usgs-mrs.cr.usgs.gov/NHDHelp/WebHelp/NHD_Help_Portal.html#NHD_Help/Introduction_to_the_NHD/Feature_Attribution/Stream_Order.htm if a == b: return a + 1 else: if a > b: return a else: return b # End if no order change (return the higher order) # End if stream order increase # End def getOrder() #------------------------------------------------------------------------------- # C L A S S : StreamNet #------------------------------------------------------------------------------- class StreamNet(): StrmArrays = {} # Arrays of all the features StreamData = {} # Tuples of all stream data (EdgeType and Enabled) StreamCount = 0 # The number of stream objects created MasterList = [] OneList = [] TransList = [] FourList = [] Furthest = None # Running record of object with highest Depth featureCount = 0 # Total number of features even if objects not created for them Orphans = [] # A list of all features not currently in a built tree def __init__(self, ID): self.StreamID = ID self.Depth = 0 self.Order = -1 self.Left = None self.Right = None self.Parent = None self.ExtraChild = [] self.StrState = 0 self.EdgeType = StreamNet.StreamData[ID][0] self.Transition = False # Set to True if (left) child is a different EdgeType self.Path = "" self.Length = arcpy.Polyline(StreamNet.StrmArrays[ID]).length self.Source2Start = 0.0 # Maximum path distance to Source self.End2Mouth = 0.0 StreamNet.StreamCount += 1 if not SUPERQUIET: if not VERBOSE and (StreamNet.StreamCount % 100) == 0: # Medium quietness level, not too little, not too much TwoWrite('. ') # else # If VERBOSE is True then print out the fuller picture further down in setChildren() # This section is dots only so we know the program is running # End if "not even a peep" # End def __init__() def setChildren(self): # return self.StrState # Creates children for one object # ExtraChild List: # [] => Leaf Node # [Left] => OneChild # [Left,Right] => Normal junction # [Left,Right,...] => FourWay ChildList = [] for key,value in StreamNet.StrmArrays.items(): if StreamNet.StrmArrays[self.StreamID][START].equals(value[END]): ChildList.append(key) # End If a start-end match # End for all the endpoints # Look for loops for child in ChildList: if child in StreamNet.MasterList: ThreeWrite("[setChildren] NOT creating DUPLICATE Stream object. Removing " + str(child) + " from ChildList: " + str(ChildList)) ChildList.remove(child) else: StreamNet.MasterList.append(child) # End StreamNet.MasterList # End remove already created children from ChildList if len(ChildList) == 0: self.StrState = Leaf self.Order = 0 elif len(ChildList) == 1: self.StrState = OneChild # ASSERT: Only OneChild's can be transitions to/from type 3 streams. self.Left = StreamNet(ChildList[LEFT]) if self.Left.EdgeType != self.EdgeType: self.Transition = True StreamNet.TransList.append(self.StreamID) else: StreamNet.OneList.append(self.StreamID) # End if transition self.Left.Parent = self self.Left.Depth = self.Depth + 1 self.Left.Path = self.Path + 'L' self.Left.End2Mouth = self.End2Mouth + self.Length else: self.StrState = Normal # ASSERT: Transition not applicable since no stream fork will occurr exactly at the boundary of a pond/lake. if len(ChildList) > 2: self.StrState = FourWay self.ExtraChild = ChildList[2:] ChildList = ChildList[0:2] StreamNet.FourList.append(self.StreamID) # End If FourWay if IsLeftRight( [StreamNet.StrmArrays[self.StreamID][FIRST], StreamNet.StrmArrays[self.StreamID][SECOND]], [StreamNet.StrmArrays[ChildList[B]] [SECONDEND], StreamNet.StrmArrays[ChildList[B]] [END]], [StreamNet.StrmArrays[ChildList[A]] [SECONDEND], StreamNet.StrmArrays[ChildList[A]] [END]] ): ChildList = [ChildList[B], ChildList[A]] # End If B is the Left stream self.Left = StreamNet(ChildList[LEFT]) self.Left.Parent = self self.Left.Depth = self.Depth + 1 self.Left.Path = self.Path + 'L' self.Left.End2Mouth = self.End2Mouth + self.Length self.Right = StreamNet(ChildList[RIGHT]) self.Right.Parent = self self.Right.Depth = self.Depth + 1 self.Right.Path = self.Path + 'R' self.Right.End2Mouth = self.End2Mouth + self.Length return self.StrState # End def setChildren() def FollowPath(self, Path): # Given a path of form: "[LR]*" (e.g. 'LLRLLLRRRLRR') return the object at the end. if Path == '': Target = self elif self.StrState == Leaf: Target = None else: if Path[0] == 'L': if self.Left == None: Target = None else: Target = self.Left.FollowPath(Path[1:]) # End If empty branch else: if self.Right == None: Target = None else: Target = self.Right.FollowPath(Path[1:]) # End If empty branch # End If Left or Right # End If empty path return Target # Failure - No object found yet, keep looking # NOTE: this return is NOT redundant. We didn't do returns above so we're returning here # This is a slightly different style than getPath and getObject bwlow which do their returns inside # End def FollowPath() def GetPath(self, targetID): # Given an object, find the path of form "[LR]*" to that object from the current object. # EXITS: Two ways out: # Fail: exit by dropping out of the bottom of the function - Return value == None # Pass: exit via one of the inside returns - Return value == String tempPath = None if self.StreamID == targetID: return '' else: if self.Left: tempPath = self.Left.GetPath(targetID) if tempPath == '' or tempPath: return 'L' + tempPath # End If nothing found on the Left branch # End If Left exists if self.Right: tempPath = self.Right.GetPath(targetID) if tempPath == '' or tempPath: return 'R' + tempPath # End If nothing found on the Right branch # End If Right exists and we haven't had success yet # End if the target node return tempPath # Failure - no path found yet, keep looking # NOTE: this return is redundant. We'll never be here with a value other than None # End def GetPath() def getObject(self, ID): # PRECONDITION: ID must be reachable from self by normal traversal (i.e. NOT extra children) # EXITS: Two ways out: # Fail: exit by dropping out of the bottom of the function - Value == None # Pass: exit via one of the inside returns - Value == Object nextBranch = None if self.StreamID == ID: return self else: if self.Left: nextBranch = self.Left.getObject(ID) if nextBranch: return nextBranch # End If Left exists if self.Right: nextBranch = self.Right.getObject(ID) if nextBranch: return nextBranch # End If Right exists # End If this object is the target return None # Failure - No path found yet, keep looking # NOTE: this return is redundant. We'll never be here with a value other than None # End def getObject() #------------------------------------------------------------------------------- # StreamNet S T A T I C M E T H O D S #------------------------------------------------------------------------------- def buildTree(trunk): # Traversal only, no objects created at this method if not StreamNet.Furthest: # ASSERT: This is the GreatTrunk StreamNet.Furthest = trunk else: if trunk.Depth > StreamNet.Furthest.Depth: StreamNet.Furthest = trunk # End if new furthest node # End if not the first node thisState = trunk.setChildren() # ASSERT: Child objects have now been created with the preceding call to 'setChildren' if VERBOSE: if (thisState == OneChild) and (trunk.Transition): transitionFlag = 'T' else: transitionFlag = ' ' # End transition flag output = "{:>5} {:<8} {:<1} Depth: {:>3} ".format(trunk.StreamID, NodeCnt[thisState], transitionFlag, trunk.Depth) if len(trunk.ExtraChild) > 0: output = output + str(trunk.ExtraChild) # End tack on extra child data to output ThreeWrite(output) else: TwoWrite('.') # End printing if thisState != Leaf: StreamNet.buildTree(trunk.Left) # Recursive Branch Left trunk.Order = trunk.Left.Order trunk.Source2Start = trunk.Left.Source2Start + trunk.Left.Length if thisState != OneChild: StreamNet.buildTree(trunk.Right) # Recursive Branch Right trunk.Order = getOrder(trunk.Left.Order,trunk.Right.Order) if trunk.Source2Start < trunk.Right.Source2Start + trunk.Right.Length: trunk.Source2Start = trunk.Right.Source2Start + trunk.Right.Length # End if the right length is longer than the left length # NOTE: If implementing following FourWay children then it will probably go here. # End If's - How deep does the recursive rabbit hole go? # End def buildTree() def FindOrphans(trunk): # INCOMPLETE: currently does not traverse ExtraChild's # PRECONDITIONS: trunk is a built tree, all features not in this tree will be considered orphans. # StreamNet.Orphans list has been populated with all features in the fc (or whereClause). # We could remove the try/except if we assumed no objects exist without features try: StreamNet.Orphans.remove(trunk.StreamID) for Child in trunk.ExtraChild: StreamNet.Orphans.remove(Child) # End For extra children except: ThreeWrite("FindOrphans() found either a duplicate StreamID") ThreeWrite("or an object for which no Feature OID exists. StreamID = " + str(trunk.StreamID)) # End try to remove StreamID if trunk.Left: StreamNet.FindOrphans(trunk.Left) # End recurse Left branch if trunk.Right: StreamNet.FindOrphans(trunk.Right) # End recurse Right branch # End def FindOrphans() def TraverseNSave(ThisStream): # Save all node data to dictionary # ......... S A V E ..................................................................................................................... # ['StreamID','Left_ID','Right_ID','Parent_ID','StreamState','PathDepth','StreamOrder','BranchPath','LineLength','Source2Start','End2Mouth'] # ......................................................................................................................................... Row = [ThisStream.StreamID] # 'StreamID' if ThisStream.Left: # 'Left_ID' Row.append(ThisStream.Left.StreamID) else: Row.append(0) # End if Left if ThisStream.Right: # 'Right_ID' Row.append(ThisStream.Right.StreamID) else: Row.append(0) # End if Right if ThisStream.Parent: # 'Parent_ID' Row.append(ThisStream.Parent.StreamID) else: Row.append(0) # End if parent exists Row.append(ThisStream.StrState) # 'StreamState' Row.append(ThisStream.Depth) # 'PathDepth' Row.append(ThisStream.Order) # 'StreamOrder' Row.append(ThisStream.Path[:100]) # 'BranchPath' # Do I want to save long paths? Do I want to save paths at all? Row.append(ThisStream.Length) # 'LineLength' Row.append(ThisStream.Source2Start) # 'Source2Start' Row.append(ThisStream.End2Mouth) # 'End2Mouth' dataStore[ThisStream.StreamID] = Row # Save the row # ......... T R A V E R S E ......... if ThisStream.Left: StreamNet.TraverseNSave(ThisStream.Left) # ASSERT: Right cannot exist without Left (but Left *can* exist without Right) if ThisStream.Right: StreamNet.TraverseNSave(ThisStream.Right) # End recurse Right branch # End recurse Left branch # End def TraverseNSave() # ........................................................... # StreamNet C L A S S F U N C T I O N S # ........................................................... def getStreams(fc): # Returns a 4-Tuple Count = 0 StrmArrays = {} # StreamID = ( Array ) StreamData = {} # StreamID = ( tuple of all other columns from the search cursor ) featureList = [] StartPts = {} StartCnt = {} bifurcates = [] with arcpy.da.SearchCursor(fc,["OID@","SHAPE@","EdgeType","Enabled"],whereClause) as streams: for stream in streams: if ((stream[EDGETYPE] == 1) or (stream[EDGETYPE] == 3)) and stream[ENABLED]: Count += 1 if not SUPERQUIET and ((Count % 100) == 0): # Ignoring VERBOSE TwoWrite(str(stream[0]) + ",") StrmArrays[stream[0]] = (stream[1][0]) # A dictionary of arcpy.Arrays with OID as keys StreamData[stream[0]] = (stream[2],stream[3]) # A dictionary of tuples with OID as keys featureList.append(stream[0]) # List of all OIDs for orphan determination StartPts[stream[0]] = (stream[1][0][0].X,stream[1][0][0].Y) if StartPts[stream[0]] in StartCnt.keys(): StartCnt[StartPts[stream[0]]] += 1 else: StartCnt[StartPts[stream[0]]] = 1 # End stuffing count dict # End if stream is Type 1 or 3 and enabled # End for all streams in FC # End with the searchCursor for bifurcate,count in StartCnt.items(): if count > 1: bifurcates.append(bifurcate) # End counts above 1 # End for all stream counts return (StrmArrays,Count,featureList,StreamData) # End def getStreams() # ............................................. # Class initialization code (not definitions) # ............................................. print "\n" (StrmArrays,featureCount,Orphans,StreamData) = getStreams(fc) # Class Variables # ........................................................... # StreamNet E N D C L A S S # ........................................................... # ........................................................... # M A I N - T E S T S U I T E # ........................................................... print "\n" outFile.close() with open(outFileName,'a') as outFile: GreatTrunk = StreamNet(StartTrunk) Tree = StreamNet.buildTree(GreatTrunk) StreamNet.TraverseNSave(GreatTrunk) if arcpy.Exists(FullOutTabName): arcpy.Delete_management(FullOutTabName) # End delete existing ResultsTable = arcpy.CreateTable_management(workPath, OutTabName, TemplateName) with arcpy.da.InsertCursor(ResultsTable,outFields) as outputCursor: for node in dataStore.values(): outputCursor.insertRow(node) # End Inserting records StreamNet.FindOrphans(GreatTrunk) ThreeWrite("\nOrphan list: {}".format(StreamNet.Orphans)) ThreeWrite("\n") ThreeWrite("The furthest stream is {} at a depth of {}".format(StreamNet.Furthest.StreamID,StreamNet.Furthest.Depth)) ThreeWrite("\n") ThreeWrite("Total number of streams in the tree: {}".format(StreamNet.StreamCount)) ThreeWrite("Total number of features in the featureClass: {}".format(StreamNet.featureCount)) ThreeWrite("{} OneChilds: {}".format(len(StreamNet.OneList),StreamNet.OneList)) ThreeWrite("{} Transition: {}".format(len(StreamNet.TransList),StreamNet.TransList)) ThreeWrite("{} FourWays: {}".format(len(StreamNet.FourList),StreamNet.FourList)) print "\nDone"