diff --git a/experiments/rdmsol_generate.py b/experiments/rdmsol_generate.py index 1f22d2d..e586316 100644 --- a/experiments/rdmsol_generate.py +++ b/experiments/rdmsol_generate.py @@ -1,12 +1,11 @@ """ -Run a PCSE ensemble for different values of RDMSOL (maximum rooting depth) to determine its effect. +Run a PCSE ensemble for different values of input parameters to determine their effect. All available soil types are tested. -A single site, roughly central to the Netherlands, is used. +A specified number of sites, roughly central to the Netherlands, are used. Example: - python experiments/rdmsol_generate.py -v -c barley -n 9 + python experiments/parameter_sensitivity.py rdmsol wav -n 100 -v -c barley -s 16 """ -from itertools import product from functools import partial from tqdm import tqdm @@ -16,42 +15,39 @@ ### Parse command line arguments import argparse parser = argparse.ArgumentParser(description="Run a PCSE ensemble for different values of RDMSOL (maximum rooting depth) to determine its effect.") +parser.add_argument("parameter_names", help="parameter(s) to iterate over", type=str.upper, nargs="*") parser.add_argument("-d", "--data_dir", help="folder to load PCSE data from", type=fpcup.io.Path, default=fpcup.settings.DEFAULT_DATA) -parser.add_argument("-o", "--output_dir", help="folder to save PCSE outputs to", type=fpcup.io.Path, default=fpcup.settings.DEFAULT_OUTPUT / "rdmsol") +parser.add_argument("-o", "--output_dir", help="folder to save PCSE outputs to (default: generated from parameters)", type=fpcup.io.Path, default=None) +parser.add_argument("-n", "--number", help="number of values per parameter; note the exponential increase in runs when doing multiple parameters", type=int, default=100) parser.add_argument("-c", "--crop", help="crop to run simulations on", default="barley", choices=("barley", "maize", "sorghum", "soy", "wheat"), type=str.lower) -parser.add_argument("-n", "--number", help="number of locations; result may be lower due to rounding", type=int, default=25) +parser.add_argument("-s", "--number_sites", help="number of sites; result may be lower due to rounding", type=int, default=16) parser.add_argument("-v", "--verbose", help="increase output verbosity", action="store_true") args = parser.parse_args() +# Generate output folder name from parameters +if args.output_dir is None: + args.output_dir = fpcup.settings.DEFAULT_OUTPUT / "-".join(args.parameter_names) ### Load constants -parameter_name = "RDMSOL" -parameter = fpcup.parameters.pcse_inputs[parameter_name] -parameterrange = parameter.generate_space(n=100) - YEAR = 2022 -soiltypes = fpcup.soil.soil_types.values() -iterable = list(product(soiltypes, parameterrange)) crs = fpcup.constants.WGS84 cropdata = fpcup.crop.default sowdate = fpcup.agro.sowdate_range(args.crop, YEAR)[0] agromanagement = fpcup.agro.load_agrotemplate(args.crop, sowdate=sowdate) +soiltypes = list(fpcup.soil.soil_types.values()) ### Worker function; this runs PCSE once for one site -def run_pcse(soildata_and_depth: tuple[fpcup.soil.SoilType, int], *, - coordinates: fpcup._typing.Coordinates, sitedata: fpcup.site.PCSESiteDataProvider, weatherdata: fpcup.weather.WeatherDataProvider) -> bool | fpcup.model.RunDataBRP: +def run_pcse(overrides: fpcup._typing.Iterable[dict[str, fpcup._typing.Number]], *, + coordinates: fpcup._typing.Coordinates, sitedata: fpcup.site.PCSESiteDataProvider, weatherdata: fpcup.weather.WeatherDataProvider, soildata: fpcup.soil.SoilType) -> bool | fpcup.model.RunDataBRP: """ - For a single pair of soil type and depth (RDMSOL), wrap the data into a RunData object, and run PCSE on it. - Uses pre-made sitedata and weatherdata provided through a `partial` object. + For a single dictionary of override values (e.g. RDMSOL and WAV), wrap the data into a RunData object, and run PCSE on it. + The other kwargs (e.g. sitedata, weatherdata) must be provided beforehand through a `partial` object. Returns True if the results were succesfully written to file. Returns the corresponding RunData if a run failed. """ - # Setup - soildata, depth = soildata_and_depth - # Combine input data - run = fpcup.model.RunData(sitedata=sitedata, soildata=soildata, cropdata=cropdata, weatherdata=weatherdata, agromanagement=agromanagement, override={parameter_name: depth}, geometry=coordinates, crs=crs) + run = fpcup.model.RunData(override=overrides, sitedata=sitedata, soildata=soildata, cropdata=cropdata, weatherdata=weatherdata, agromanagement=agromanagement, geometry=coordinates, crs=crs) run.to_file(args.output_dir) # Run model @@ -82,22 +78,31 @@ def run_pcse(soildata_and_depth: tuple[fpcup.soil.SoilType, int], *, print(agromanagement) print() + # Generate the iterable + iterable = fpcup.parameters.generate_ensemble_space(*args.parameter_names, n=args.number) + if args.verbose: + print(f"Iterating over the following parameters: {args.parameter_names}") + print(f"Length of iterator: {len(iterable)}") + # Generate the coordinates - coords = fpcup.site.generate_sites_space(latitude=(51, 53), longitude=(4, 6), n=args.number) + coords = fpcup.site.generate_sites_space(latitude=(51, 53), longitude=(4, 6), n=args.number_sites) if args.verbose: print(f"Generated {len(coords)} coordinates.") ### Loop over sites (coordinates) + if args.verbose: + print(f"Expected total runs: {len(iterable) * len(coords) * len(soiltypes)}") model_statuses_combined = [] for c in tqdm(coords, desc="Sites", unit="site", leave=args.verbose): - ### Generate site-specific data - sitedata = fpcup.site.example(c) - weatherdata = fpcup.weather.load_weather_data_NASAPower(c) - run_pcse_site = partial(run_pcse, coordinates=c, sitedata=sitedata, weatherdata=weatherdata) - - ### Run the model - model_statuses = fpcup.model.multiprocess_pcse(run_pcse_site, iterable, leave_progressbar=False) - model_statuses_combined.extend(model_statuses) + for soildata in tqdm(soiltypes, desc="Soil types", unit="type", leave=False): + ### Generate site-specific data + sitedata = fpcup.site.example(c) + weatherdata = fpcup.weather.load_weather_data_NASAPower(c) + run_pcse_site = partial(run_pcse, coordinates=c, sitedata=sitedata, weatherdata=weatherdata, soildata=soildata) + + ### Run the model + model_statuses = fpcup.model.multiprocess_pcse(run_pcse_site, iterable, leave_progressbar=False) + model_statuses_combined.extend(model_statuses) failed_runs = fpcup.model.process_model_statuses(model_statuses, verbose=args.verbose) diff --git a/fpcup/parameters.py b/fpcup/parameters.py index 8cb62c4..5960e75 100644 --- a/fpcup/parameters.py +++ b/fpcup/parameters.py @@ -3,10 +3,12 @@ Includes a non-comprehensive list of parameters, sorted by subject and input/output status. """ from dataclasses import dataclass +from itertools import product import numpy as np from ._typing import Iterable, Optional, RealNumber +from .tools import dict_product ### Classes for dealing with parameters @dataclass @@ -94,7 +96,7 @@ def __str__(self) -> str: SOLNAM = PCSELabel(name="SOLNAM", description="Soil name") -n = PCSENumericParameter(name="area", description="Number of sites") +n = PCSENumericParameter(name="n", description="Number of sites") area = PCSENumericParameter(name="area", description="Total plot area", unit="ha") WAV = PCSENumericParameter(name="WAV", description="Initial amount of water in rootable zone in excess of wilting point", plotname="Initial amount of water", unit="cm", bounds=(0, 50)) NOTINF = PCSENumericParameter(name="NOTINF", description="Non-infiltrating fraction", bounds=(0, 1)) @@ -169,3 +171,14 @@ def parameterdict(*params: Iterable[PCSEParameter]) -> dict[str, PCSEParameter]: all_parameters = {**pcse_inputs, **pcse_outputs, **pcse_summary_outputs, **aggregate_parameters, **crop_parameters, **soil_parameters, **site_parameters} + + +### Functions related to generating ensembles +def generate_ensemble_space(*parameter_names, n: int=100) -> dict: + """ + Generate an iterable spanning the input space for any number of parameter_names. + """ + parameters = [all_parameters[name] for name in sorted(parameter_names)] + parameter_ranges = {p.name: p.generate_space(n) for p in parameters} + combined = dict_product(parameter_ranges) + return combined