create source/sink type attribute

CADWRDeltaModeling · Dec 3, 2024 · 907f1bd · 907f1bd
1 parent 391227e
commit 907f1bd
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/schimpy/schism_sources_sinks.py b/schimpy/schism_sources_sinks.py
@@ -410,13 +410,14 @@ def yaml2df(source_yaml):
     with open(source_yaml, 'r') as file:
         source_sink = yaml.safe_load(file)
     if {"sources", "sinks"} & source_sink.keys():
-        sites, x, y = [], [], []
+        sites, x, y, stype = [], [], [], []
         for v in ["sources","sinks"]:
             skeys = source_sink[v].keys()
             skeys = list(skeys)
             sites.extend(skeys)
             a = [source_sink[v][k] for k in skeys]
             a = np.array(a)
+            stype.extend([v] * len(skeys))
             x.extend(a[:,0])
             y.extend(a[:,1])
     else:
@@ -426,7 +427,8 @@ def yaml2df(source_yaml):
         a = np.array(a)
         x = a[:,0]
         y = a[:,1]
-    return pd.DataFrame({'sites': sites,'x':x,'y':y})
+        stype = None
+    return pd.DataFrame({'sites': sites,'x':x,'y':y,'stype':stype})
 
 if __name__ == "__main__":
     yaml_fn1 = 'source_sink.yaml'