Merge pull request #216 from jinyan1214/master

Some UI/API edits for tranportation inventory generation
NHERI-SimCenter · Jul 17, 2024 · 10274e1 · 10274e1
2 parents 6b4dc6a + db47be3
commit 10274e1
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 7 deletions.
diff --git a/brails/TranspInventoryGenerator.py b/brails/TranspInventoryGenerator.py
@@ -74,19 +74,20 @@
 class TranspInventoryGenerator:
 
     def __init__(self, location='Berkeley, CA'):                
-        self.enabledElements = ['roads','bridges','tunnels','railroads']
+        self.enabledElements = ['roads','bridge','tunnel','railroads']
         self.location = location
         self.workDir = 'tmp'
         self.modelDir = 'tmp/models'
         self.inventory_files = ''
+        self.roadDataSource = 'TIGER' # OSM or TIGER
 
     def enabled_elements(self):
         print('Here is the list of elements currently enabled in TranspInventoryGenerator:\n')
         for element in self.enabledElements:
             print(f'       {element}')
 
     def generate(self):
-        tphandler = TransportationElementHandler()
+        tphandler = TransportationElementHandler(self.enabledElements)
         tphandler.fetch_transportation_elements(self.location) 
 
         self.inventory_files = tphandler.output_files

diff --git a/brails/workflow/TransportationElementHandler.py b/brails/workflow/TransportationElementHandler.py
@@ -48,9 +48,15 @@
 from brails.workflow.FootprintHandler import FootprintHandler 
 
 class TransportationElementHandler:
-    def __init__(self): 
+    def __init__(self, enabledElements, roadDataSource = 'TIGER'): 
         self.queryarea = ''
-        self.output_files = {'roads':'Roads.geojson'}
+        if 'roads' in enabledElements:
+            self.output_files = {'roads':'Roads.geojson'}
+        else:
+            self.output_files = {}
+        self.roadDataSource = roadDataSource
+        self.enabledElements = enabledElements
+
 
     def fetch_transportation_elements(self, queryarea:str):
 
@@ -146,8 +152,10 @@ def get_el_counts(bpoly:Polygon, eltype: str) -> int:
 
             # Download the element count for the bounding polygon using the 
             # defined retry strategy:
+            print("Querying element count for the bounding polygon")
             r = s.get(query)
             elcount = r.json()['count']
+            print(f"Querying finished with count {elcount}")
 
             return elcount                 
 
@@ -341,9 +349,13 @@ def combine_write_roadjsons(roadjsons,bpoly):
             bpoly,_,_ = fpHandler._FootprintHandler__fetch_roi(queryarea)   
 
         # Define supported element types:
-        eltypes = ['bridge', 'tunnel', 'railroad', 'primary_road',
-                   'secondary_road', 'local_road']
+        eltypes = self.enabledElements.copy()
         roadjsons = {'primary_road':[], 'secondary_road':[], 'local_road':[]}
+        if 'roads' in eltypes:
+            eltypes.remove('roads')
+            eltypes += ['primary_road',
+                   'secondary_road', 'local_road']
+
 
         # Write the GeoJSON output for each element:
         for eltype in eltypes:
@@ -354,4 +366,5 @@ def combine_write_roadjsons(roadjsons,bpoly):
             else:
                 print(f"Fetching {eltype.replace('_',' ')}s, may take some time...")
                 roadjsons[eltype] = (write2geojson(bpoly,eltype))
-        combine_write_roadjsons(roadjsons,bpoly)
+        if ['roads'] in self.enabledElements:
+            combine_write_roadjsons(roadjsons,bpoly)