Simplify BRP queries

demo.ipynb

@@ -412,7 +412,8 @@
    "metadata": {},
    "source": [
     "More interesting than a single WOFOST run is the ability to do multiple runs at the same time.\n",
-    "This can be used to simulate multiple locations at the same time, such as different plots in the BRP or random sites within one province."
+    "This can be used to simulate multiple locations at the same time, such as different plots in the BRP or random sites within one province.\n",
+    "In this example, we will run WOFOST for all barley-growing plots in the province of Zuid-Holland."
    ]
   },
   {
@@ -425,11 +426,12 @@
     "# Setup BRP\n",
     "year = 2020\n",
     "brp = fpcup.io.load_brp(year)\n",
-    "index = 1  # Change this to any index in the BRP table. Try some values and see what happens!\n",
     "\n",
-    "# Select plot\n",
-    "brp_plot = brp.iloc[index]\n",
-    "brp_plot"
+    "# Select plots\n",
+    "crop_species = \"barley\"\n",
+    "province = fpcup.geo.process_input_province(\"Zuid-Holland\") \n",
+    "brp_plots = fpcup.io.query_brp(brp, province=province, crop_species=crop_species)\n",
+    "brp_plots"
    ]
   },
   {

fpcup/io.py

@@ -7,14 +7,11 @@
 
 import geopandas as gpd
 gpd.options.io_engine = "pyogrio"
-from geopandas import GeoDataFrame
 
 from pyogrio import read_dataframe as read_geodataframe, write_dataframe as write_geodataframe
 from pyogrio.errors import DataSourceError
 
-from tqdm import tqdm
-
-from .constants import CRS_AMERSFOORT
+from .geo import process_input_province
 from .model import InputSummary, Summary, TimeSeries
 from .multiprocessing import multiprocess_file_io
 from .typing import Iterable, Optional, PathOrStr
@@ -33,6 +30,24 @@ def load_brp(year: int) -> gpd.GeoDataFrame:
     return brp
 
 
+def query_brp(brp: gpd.GeoDataFrame, **kwargs) -> gpd.GeoDataFrame:
+    """
+    Query the BRP (or another dataframe) on multiple columns.
+    Example:
+        query_brp(brp, province="Groningen", crop_species="barley")
+    """
+    # Pre-process common keys
+    if "province" in kwargs:
+        province = kwargs["province"]
+        if isinstance(province, str):
+            province = process_input_province(province)
+        kwargs["province"] = province.abbreviation
+
+    query = [f"{key} == '{value}'" for key, value in kwargs.items()]
+    query = " & ".join(query)
+    return brp.query(query)
+
+
 def save_ensemble_summary(output_dir: PathOrStr, *,
                           use_existing=True, verbose=False) -> None:
     """