Merge pull request #947 from akrherz/iemre_domains

IEMRE domain stuff

.github/workflows/etchosts.txt

@@ -3,6 +3,8 @@
 127.0.0.1 iemdb-mos.local
 127.0.0.1 iemdb-idep.local
 127.0.0.1 iemdb-iemre.local
+127.0.0.1 iemdb-iemre_china.local
+127.0.0.1 iemdb-iemre_europe.local
 127.0.0.1 iemdb-postgis.local
 127.0.0.1 iemdb-mesosite.local
 127.0.0.1 iemdb-afos.local

.pre-commit-config.yaml

@@ -2,7 +2,7 @@ ci:
   autoupdate_schedule: quarterly
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.5.5"
+    rev: "v0.5.6"
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

src/pyiem/iemre.py

@@ -99,17 +99,19 @@ def get_table(valid):
     return table
 
 
-def set_grids(valid, ds, table=None):
+def set_grids(valid, ds, table=None, domain: str = ""):
     """Update the database with a given ``xarray.Dataset``.
 
     Args:
       valid (datetime or date): If datetime, save hourly, if date, save daily
       ds (xarray.Dataset): The xarray dataset to save
       table (str,optional): hard coded database table to use to set the data
         on.  Usually dynamically computed.
+      domain (str,optional): IEMRE domain to save data to
     """
     table = Identifier(table if table is not None else get_table(valid))
-    pgconn = get_dbconn("iemre")
+    dom = DOMAINS[domain]
+    pgconn = get_dbconn(d2l(domain))
     cursor = pgconn.cursor()
     # Do we currently have database entries?
     cursor.execute(
@@ -141,17 +143,19 @@ def set_grids(valid, ds, table=None):
 
     # Implementation notes: seems quite fast
     pig = {v: ds[v].values.ravel().tolist() for v in ds}
-    pig["valid"] = [f"{valid:%Y-%m-%d}"] * (NX * NY)
-    pig["gid"] = list(range(NX * NY))
+    pig["valid"] = [f"{valid:%Y-%m-%d}"] * (dom["nx"] * dom["ny"])
+    pig["gid"] = list(range(dom["nx"] * dom["ny"]))
 
     sts = utc()
     cursor.executemany(query, zip(*pig.values()))
     cursor.close()
     pgconn.commit()
-    LOG.info("timing %.2f/s", NX * NY / (utc() - sts).total_seconds())
+    LOG.info(
+        "timing %.2f/s", dom["nx"] * dom["ny"] / (utc() - sts).total_seconds()
+    )
 
 
-def get_grids(valid, varnames=None, cursor=None, table=None):
+def get_grids(valid, varnames=None, cursor=None, table=None, domain: str = ""):
     """Fetch grid(s) from the database, returning xarray.
 
     Args:
@@ -161,12 +165,14 @@ def get_grids(valid, varnames=None, cursor=None, table=None):
       cursor (database cursor,optional): cursor to use for query
       table (str,optional): Hard coded table to fetch data from, useful in the
         case of forecast data.
+      domain (str,optional): IEMRE domain to fetch data from
 
     Returns:
       ``xarray.Dataset``"""
     table = table if table is not None else get_table(valid)
+    dom = DOMAINS[domain]
     if cursor is None:
-        pgconn = get_dbconn("iemre")
+        pgconn = get_dbconn(d2l(domain))
         cursor = pgconn.cursor()
     # rectify varnames
     if isinstance(varnames, str):
@@ -186,15 +192,20 @@ def get_grids(valid, varnames=None, cursor=None, table=None):
     cursor.execute(
         f"SELECT (gid / %s)::int as y, gid %% %s as x, {colsql} "
         f"from {table} WHERE valid = %s",
-        (NX, NX, valid),
+        (dom["nx"], dom["ny"], valid),
     )
-    data = dict((key, np.full((NY, NX), np.nan)) for key in use_columns)
+    data = {
+        key: np.full((dom["ny"], dom["nx"]), np.nan) for key in use_columns
+    }
     for row in cursor:
         for i, col in enumerate(use_columns):
             data[col][row[0], row[1]] = row[2 + i]
     ds = xr.Dataset(
         dict((key, (["y", "x"], data[key])) for key in data),
-        coords={"lon": (["x"], XAXIS), "lat": (["y"], YAXIS)},
+        coords={
+            "lon": (["x"], np.arange(dom["west"], dom["east"], DX)),
+            "lat": (["y"], np.arange(dom["south"], dom["north"], DY)),
+        },
     )
     return ds
 
@@ -270,6 +281,17 @@ def find_ij(lon, lat, domain: str = "") -> Tuple[Optional[int], Optional[int]]:
     return i, j
 
 
+def get_domain(lon: float, lat: float) -> Optional[str]:
+    """Compute the domain that contains the given point."""
+    for domain, dom in DOMAINS.items():
+        if (
+            dom["west"] <= lon < dom["east"]
+            and dom["south"] <= lat < dom["north"]
+        ):
+            return domain
+    return None
+
+
 def get_gid(lon, lat, domain: str = "") -> Optional[int]:
     """Compute the grid id for the given location."""
     i, j = find_ij(lon, lat, domain)

tests/test_iemre.py

@@ -95,6 +95,12 @@ def test_get_table():
     assert iemre.get_table(d2) == "iemre_hourly_200009"
 
 
+def test_get_domain():
+    """Test the get_domain method."""
+    assert iemre.get_domain(100, 30) == "china"
+    assert iemre.get_domain(-1000, 30) is None
+
+
 def test_get_gid():
     """Can we get a gid?"""
     assert iemre.get_gid(-96, 44) is not None
@@ -103,12 +109,13 @@ def test_get_gid():
 
 def test_writing_grids():
     """Test letting the API write data from the future."""
-    pgconn = database.get_dbconn("iemre")
+    domain = "china"
+    pgconn = database.get_dbconn(iemre.d2l(domain))
     cursor = pgconn.cursor()
     valid = datetime.date.today() + datetime.timedelta(days=120)
-    ds = iemre.get_grids(valid, varnames=["high_tmpk"])
-    iemre.set_grids(valid, ds)
-    ds = iemre.get_grids(valid, varnames=["high_tmpk"])
+    ds = iemre.get_grids(valid, varnames=["high_tmpk"], domain=domain)
+    iemre.set_grids(valid, ds, domain=domain)
+    ds = iemre.get_grids(valid, varnames=["high_tmpk"], domain=domain)
     assert ds["high_tmpk"].lat[0] > 0
     # Cleanup after ourself
     cursor.execute(