added rain_geom_table property and setUpClass

citycatpg/run.py

@@ -7,6 +7,7 @@
 from citycatio import Model
 import pandas as pd
 import rasterio as rio
+import geopandas as gpd
 
 
 @dataclass
@@ -43,6 +44,11 @@ class Run:
 
     model: Optional[Model] = None
 
+    @property
+    def rain_geom_table(self):
+        if self.rain_table is not None:
+            return self.rain_table + '_geom'
+
     def add(self, con: connection):
 
         query = sql.SQL("""
@@ -117,6 +123,7 @@ def get_model(self, con):
         self.model = Model(
             dem=self.get_dem(con),
             rainfall=self.get_rainfall(con),
+            rainfall_polygons=self.get_rainfall_polygons(con),
             duration=self.run_duration,
             output_interval=self.output_frequency
         )
@@ -152,6 +159,22 @@ def get_rainfall(self, con):
 
         return rain
 
+    def get_rainfall_polygons(self, con):
+        if self.rain_total:
+            return
+
+        else:
+            return gpd.GeoDataFrame.from_postgis(
+                sql.SQL(
+                """
+                SELECT {rain_geom_table}.gid, {rain_geom_table}.geom FROM {rain_geom_table}, {domain_table} 
+                WHERE ST_Intersects({rain_geom_table}.geom, {domain_table}.geom)
+                """).format(
+                    rain_geom_table=sql.Identifier(self.rain_geom_table),
+                    domain_table=sql.Identifier(self.domain_table)
+                ).as_string(con),
+                con)
+
 
 def fetch(con, run_id, run_table='runs'):
     query = sql.SQL("""

tests/test_run.py

@@ -8,47 +8,35 @@
 from psycopg2 import sql
 
 con = psycopg2.connect(database='postgres', user='postgres', password='password', host='localhost')
-dem_file = rio.MemoryFile()
-x_min, y_max, res, height, width = 100, 500, 5, 100, 200
-x_max, y_min = x_min + width * res, y_max - height * res
-array = np.round(np.random.random((height, width)), 3)
-
-with rio.open(
-        dem_file,
-        'w',
-        driver='GTiff',
-        height=height,
-        width=width,
-        count=1,
-        dtype=array.dtype,
-        transform=Affine.translation(x_min, y_max) * Affine.scale(res, -res),
-        nodata=-9999
-) as dst:
-    dst.write(array, 1)
 
 
 class TestRun(TestCase):
 
-    def test_add_run(self):
-        with con:
-            with con.cursor() as cur:
-                cur.execute('DROP TABLE IF EXISTS runs')
+    @classmethod
+    def setUpClass(cls):
+        """Insert sample data"""
+        dem_file = rio.MemoryFile()
+        x_min, y_max, res, height, width = 100, 500, 5, 100, 200
+        x_max, y_min = x_min + width * res, y_max - height * res
+        array = np.round(np.random.random((height, width)), 3)
 
-        Run(run_duration=500, srid=3035, resolution=90).add(con)
+        with rio.open(
+                dem_file,
+                'w',
+                driver='GTiff',
+                height=height,
+                width=width,
+                count=1,
+                dtype=array.dtype,
+                transform=Affine.translation(x_min, y_max) * Affine.scale(res, -res),
+                nodata=-9999
+        ) as dst:
+            dst.write(array, 1)
 
-    def test_fetch(self):
-        with con:
-            with con.cursor() as cur:
-                cur.execute('DROP TABLE IF EXISTS runs')
-
-        r = Run(run_duration=500, srid=3035, resolution=90)
-        r.add(con)
+        dem_file.seek(0)
 
-        fetch(con, run_id=r.run_id)
+        r = Run(100, 3035, 90, rain_table='rain')
 
-    def test_get_model(self):
-        dem_file.seek(0)
-        r = Run(100, 3035, 90, rain_total=100, rain_duration=500)
         with con:
             with con.cursor() as cursor:
                 cursor.execute("CREATE EXTENSION IF NOT EXISTS postgis_raster")
@@ -59,14 +47,48 @@ def test_get_model(self):
                         DROP TABLE IF EXISTS {domain_table};
                         CREATE TABLE {domain_table} (gid serial PRIMARY KEY, geom geometry);
                         INSERT INTO {domain_table} (geom) VALUES (ST_GeomFromText(%(geom)s));
+
+                        DROP TABLE IF EXISTS {rain_geom_table};
+                        CREATE TABLE {rain_geom_table} (gid serial PRIMARY KEY, geom geometry);
+                        INSERT INTO {rain_geom_table} (geom) VALUES (ST_GeomFromText(%(geom)s));
+
                         DROP TABLE IF EXISTS {dem_table};
                         CREATE TABLE {dem_table} (rast raster);
                         INSERT INTO {dem_table}(rast) VALUES (ST_FromGDALRaster(%(rast)s));
-                    """).format(domain_table=sql.Identifier(r.domain_table), dem_table=sql.Identifier(r.dem_table)),
+                    """).format(
+                        domain_table=sql.Identifier(r.domain_table),
+                        dem_table=sql.Identifier(r.dem_table),
+                        rain_geom_table=sql.Identifier(r.rain_geom_table)
+                    ),
                     dict(
                         geom=str(Polygon(
                             [[x_min, y_min], [x_min, y_max], [x_max / 2, y_max / 2], [x_max, y_min], [x_min, y_min]])),
                         rast=psycopg2.Binary(dem_file.read())))
 
+    def test_add_run(self):
+        with con:
+            with con.cursor() as cur:
+                cur.execute('DROP TABLE IF EXISTS runs')
+
+        Run(run_duration=500, srid=3035, resolution=90).add(con)
+
+    def test_fetch(self):
+        with con:
+            with con.cursor() as cur:
+                cur.execute('DROP TABLE IF EXISTS runs')
+
+        r = Run(run_duration=500, srid=3035, resolution=90)
+        r.add(con)
+
+        fetch(con, run_id=r.run_id)
+
+    def test_generate_rainfall(self):
+        r = Run(100, 3035, 90, rain_total=100, rain_duration=500)
+        r.get_model(con)
+        r.model.write('tests/test_model_generated')
+
+    def test_get_model(self):
+        r = Run(100, 3035, 90, rain_table='rain')
+
         r.get_model(con)
         r.model.write('tests/test_model')