Code utils (#25)

Code utils for:
- finding bgt_codes and wbh_codes
- creating a model_id from a wbh_code, wbh_code and object-code or x/y
coordinate

src/hydamo/hydamo/__init__.py

@@ -3,7 +3,6 @@
 from hydamo.datamodel import ExtendedGeoDataFrame, HyDAMO
 from hydamo.geometry import find_nearest_branch, possibly_intersecting
 from hydamo.styles import add_styles_to_geopackage, read_style
-from hydamo.utils import find_bgt_code
 
 __all__ = [
     "HyDAMO",
@@ -12,5 +11,4 @@
     "possibly_intersecting",
     "add_styles_to_geopackage",
     "read_style",
-    "find_bgt_code",
 ]

src/hydamo/hydamo/code_utils.py

@@ -0,0 +1,125 @@
+"""Utilities for generating unique codes."""
+from pathlib import Path
+from typing import Union
+
+import pandas as pd
+from pandas import DataFrame
+
+CODES_CSV = Path(__file__).parent.joinpath("data", "codes.csv")
+CODES_DF = None
+
+WBH_CODE_TEMPLATE = "NL.WBHCODE.{wbh_code}.{code}"
+BGT_CODE_TEMPLATE = "NL.BGTCODE.{bgt_code}.{code}"
+
+
+def get_codes_df():
+    """Get (and set) the global CODES_DF.
+
+    Returns
+    -------
+    CODES_DF : DataFrame
+
+    """
+    global CODES_DF
+
+    # read BGT_CSV if CODES_DF is None
+    if CODES_DF is None:
+        CODES_DF = pd.read_csv(CODES_CSV)
+        CODES_DF.set_index(CODES_DF.name.str.lower(), inplace=True)
+        CODES_DF["wbh_code"] = CODES_DF.wbh_code.apply(
+            lambda x: f"{int(x):02}" if not pd.isna(x) else None
+        )
+    return CODES_DF.copy()
+
+
+def bgt_code_exists(bgt_code: str):
+    """Check if bgt_code exists in CODES_DF"""
+    codes_df = get_codes_df()
+    return bgt_code in codes_df.bgt_code.to_numpy()
+
+
+def wbh_code_exists(wbh_code):
+    """Check if wbh_code exists in CODES_DF"""
+    codes_df = get_codes_df()
+    return wbh_code in codes_df.wbh_code.to_numpy()
+
+
+def bgt_to_wbh_code(bgt_code):
+    """Convert bgt_code to wbh_code if bgt_code exists"""
+    if bgt_code_exists(bgt_code):
+        codes_df = get_codes_df()
+        return (
+            codes_df.reset_index(drop=True).set_index("bgt_code").loc[bgt_code].wbh_code
+        )
+
+
+def find_codes(
+    organization: str,
+    administration_category: Union[str, None] = None,
+    to_dict: bool = True,
+):
+    """Find codes associated with an organization"""
+    codes_df = get_codes_df()
+
+    # filter on administration_category
+    if administration_category:
+        if (
+            administration_category.lower()
+            in codes_df.administration_category.to_numpy()
+        ):
+            codes_df = codes_df.loc[
+                codes_df.administration_category == administration_category.lower()
+            ]
+        else:
+            raise (
+                ValueError(
+                    f"invalid value for`administration_category`. {administration_category}` not in {codes_df.administration_category.unique()}"
+                )
+            )
+
+    # if 1 exact match we return it
+    if organization.lower() in codes_df.index:
+        df = codes_df.loc[organization.lower()]
+    # if more matches we return all matches
+    else:
+        df = codes_df.loc[
+            codes_df.name.apply(lambda x: organization.lower() in x.lower())
+        ]
+    if to_dict:
+        if isinstance(df, DataFrame):
+            return df.to_dict(orient="records")
+        else:
+            return df.to_dict()
+
+
+def generate_model_id(code, wbh_code=None, bgt_code=None, geometry=None):
+    """Generate a model_id from wbh_code or bgt_code and code or x/y coordinate"""
+    if code is None:
+        if geometry is not None:
+            code = f"loc={int(geometry.x+ 0.5)},{int(geometry.x + 0.5)}"
+        else:
+            raise ValueError(
+                f"""
+                             Specify 'code' ({code}), you have to specify 'geometry' ({geometry}).
+                """
+            )
+
+    result = None
+    if wbh_code:
+        if wbh_code_exists(wbh_code):
+            result = WBH_CODE_TEMPLATE.format(wbh_code=wbh_code, code=code)
+    elif bgt_code:
+        if bgt_code_exists(bgt_code):
+            wbh_code = bgt_to_wbh_code(bgt_code)
+            if wbh_code:
+                result = WBH_CODE_TEMPLATE.format(wbh_code=wbh_code, code=code)
+            else:
+                result = BGT_CODE_TEMPLATE.format(bgt_code=bgt_code, code=code)
+    if result is None:
+        raise ValueError(
+            f"""
+                         Specify a valid 'wbh_code' ({wbh_code}) or 'bgt_code' ({bgt_code}) to generate model_id.
+                         """
+        )
+    else:
+        return result