diff --git a/agrifoodpy/__init__.py b/agrifoodpy/__init__.py index e69de29..9deb941 100644 --- a/agrifoodpy/__init__.py +++ b/agrifoodpy/__init__.py @@ -0,0 +1,3 @@ +from .food.food import FoodBalanceSheet +from .land.land import LandDataArray +from .array_accessor import XarrayAccessorBase \ No newline at end of file diff --git a/agrifoodpy/array_accessor.py b/agrifoodpy/array_accessor.py index a0cc1b2..dfb2f16 100644 --- a/agrifoodpy/array_accessor.py +++ b/agrifoodpy/array_accessor.py @@ -33,9 +33,9 @@ def add_items(self, items, copy_from=None): if np.isscalar(items): items = [items] - # Check for duplicates + # Check for duplicates, but keep the order of the input list indexes = np.unique(items, return_index=True)[1] - items = [items[index] for index in sorted(indexes)] #issues with np.unique + items = [items[index] for index in sorted(indexes)] new_items = xr.DataArray(data = np.ones(len(items)), coords = {"Item":items}) @@ -80,6 +80,7 @@ def add_regions(self, regions, copy_from=None): if np.isscalar(regions): regions = [regions] + # Check for duplicates, but keep the order of the input list indexes = np.unique(regions, return_index=True)[1] regions = [regions[index] for index in sorted(indexes)] @@ -100,18 +101,24 @@ def add_regions(self, regions, copy_from=None): return out - def add_years(self, years, projection="empty"): - """Extends the year range of an xarray object according to the defined - maximum year + def add_years( + self, + years, + pivot_year=None, + projection="empty"): + """Adds or extends the Year coordinate of an xarray object Parameters ---------- years : list, int list of years to be added to the data + pivot_year : int, optional + Year to use as a pivot for the projection. If not provided, the + last year of the input array is used. projection : string or array_like Projection mode. If "constant", the last year of the input array is copied to every new year. If "empty", values are initialized and - set to zero. If a float array is given, these are used to populate + set to NaN. If a float array is given, these are used to populate the new year using a scaling of the last year of the array Returns @@ -125,8 +132,14 @@ def add_years(self, years, projection="empty"): if np.isscalar(years): years = [years] - indexes = np.unique(years, return_index=True)[1] - years = [years[index] for index in sorted(indexes)] + # Sort and remove duplicates + years = sorted(np.unique(years)) + + # If no "Year" dimension, add one with the first year of the input + if "Year" not in fbs.dims: + fbs = fbs.expand_dims(Year=[years[0]]) + + years = [year for year in years if year not in fbs.Year.values] if isinstance(projection, str): if projection == "empty": @@ -141,12 +154,16 @@ def add_years(self, years, projection="empty"): # Should raise an exception if sizes do not match data = np.ones(len(years)) * projection - - new_years = xr.DataArray(data=data, - coords = {"Year":years}) + new_years = xr.DataArray(data=data, coords = {"Year":years}) - # Select last year as pivot - fbs_pivot = fbs.isel(Year=-1) + if pivot_year is not None: + if pivot_year not in fbs.Year: + raise ValueError("Pivot year not found in input array") + + fbs_pivot = fbs.sel(Year=pivot_year) + else: + # Select last year of data as pivot + fbs_pivot = fbs.isel(Year=-1) # Create DS or DA by multiplying along last value new_fbs = fbs_pivot*new_years diff --git a/agrifoodpy/food/food.py b/agrifoodpy/food/food.py index ab60f8f..c7ae89f 100644 --- a/agrifoodpy/food/food.py +++ b/agrifoodpy/food/food.py @@ -152,10 +152,8 @@ def scale_element( Input DataSet with FAOSTAT like elements scale : float, float array_like or xarray.Dataarray Scaling quantities for the element DataArray - element_in : str - Element DataArray to be scaled - element_out : str - Destination element DataArray to which the difference is added to + element : str + Food Balance Sheet element to be scaled items : list of int or list of str, optional List of items to be scaled. If not provided, all items are scaled. @@ -642,7 +640,7 @@ def plot_years( if size_cumsum == 1: ax.fill_between(years, cumsum, color=colors[0], alpha=0.5) ax.plot(years, cumsum, color=colors[0], linewidth=0.5, - label=labels) + label=labels[0]) else: ax.fill_between(years, cumsum.isel({show: 0}), color=colors[0], alpha=0.5) diff --git a/agrifoodpy/food/model.py b/agrifoodpy/food/model.py index d1b8021..55c7b63 100644 --- a/agrifoodpy/food/model.py +++ b/agrifoodpy/food/model.py @@ -4,8 +4,9 @@ import xarray as xr import numpy as np import copy -from ..pipeline import standalone +from ..pipeline import standalone, pipeline_node from ..utils.dict_utils import get_dict, set_dict, item_parser +import warnings @standalone(input_keys=["fbs"], return_keys=["out_key"]) @@ -15,11 +16,13 @@ def balanced_scaling( element, items=None, constant=False, + padding_items=None, origin=None, add_to_origin=True, elasticity=None, fallback=None, add_to_fallback=True, + conversion_arr=None, out_key=None, datablock=None ): @@ -45,6 +48,9 @@ def balanced_scaling( constant : bool, optional If set to True, the sum of element remains constant by scaling the non selected items accordingly + padding_items : list, optional + List of items to use for scaling when 'constant' is True. If None, all + non-selected items are used for scaling. origin : string, list, optional Names of the DataArrays which will be used as source for the quantity changes. Any change to the "element" DataArray will be reflected in @@ -60,6 +66,11 @@ def balanced_scaling( add_to_fallback : bool, optional Whether to add or subtract the difference below zero in the origin DataArray to the fallback array. + conversion_arr : string, xarray.DataArray, tuple or float + Conversion array to pre-scale quantities. If provided, the input food + balance sheet is first converted using the conversion array, then the + scaling is applied, and finally the results are converted back to the + original units using the inverse of the conversion array. out_key : string, tuple Output datablock path to write results to. If not given, input path is overwritten @@ -74,6 +85,25 @@ def balanced_scaling( # Pepare inputs data = copy.deepcopy(get_dict(datablock, fbs)) + + if conversion_arr is not None: + if isinstance(conversion_arr, xr.DataArray): + conversion_arr = conversion_arr.where( + np.isfinite(conversion_arr), other=1) + if (conversion_arr == 0).any(): + warnings.warn("Conversion array contains zero values, " + "which can lead to inaccurate scaling") + elif isinstance(conversion_arr, (int, float)): + conversion_arr = xr.full_like( + data[element], fill_value=conversion_arr) + else: + conversion_arr = get_dict(datablock, conversion_arr) + + else: + conversion_arr = xr.ones_like(data[element]) + + data = data*conversion_arr + out = copy.deepcopy(data) if out_key is None: @@ -118,19 +148,25 @@ def balanced_scaling( delta = out[element] - data[element] # Identify non selected items and scaling - non_sel_items = np.setdiff1d(data.Item.values, scaled_items) + if padding_items is None: + non_sel_items = np.setdiff1d(data.Item.values, scaled_items) + else: + non_sel_items = item_parser(data, padding_items) non_sel_scale = (data.sel(Item=non_sel_items)[element].sum(dim="Item") - delta.sum(dim="Item")) \ / data.sel(Item=non_sel_items)[element].sum(dim="Item") - # Make sure no scaling occurs on inf and nan - non_sel_scale = non_sel_scale.where( + # Identify where denominator is zero (non-finite scale) + non_finite_mask = ~np.isfinite(non_sel_scale) + + # Use multiplicative scaling where finite, no-op (scale=1) where not + non_sel_scale_finite = non_sel_scale.where( np.isfinite(non_sel_scale)).fillna(1.0) if origin is None: out = out.fbs.scale_element( element=element, - scale=non_sel_scale, + scale=non_sel_scale_finite, items=non_sel_items ) @@ -138,12 +174,32 @@ def balanced_scaling( out = out.fbs.scale_add( element_in=element, element_out=origin, - scale=non_sel_scale, + scale=non_sel_scale_finite, items=non_sel_items, add=add_to_origin, elasticity=elasticity ) + # For non-finite cases (zero original quantity), add delta directly + # distributed equally across non_sel_items + if bool(non_finite_mask.any()): + n_non_sel = len(non_sel_items) + per_item_additive = ( + (-delta.sum(dim="Item")).where(non_finite_mask).fillna(0) + / n_non_sel + ) + + sel = {"Item": non_sel_items} + + out[element].loc[sel] = out[element].loc[sel] + per_item_additive + if origin is not None: + for elmnt, add_el, elast in zip(origin, add_to_origin, elasticity): + out[elmnt].loc[sel] = ( + out[elmnt].loc[sel] + + np.where(add_el, -1, 1) * + (-per_item_additive) * elast + ) + # If a fallback DataArray is defined, transfer the excess negative # quantities to it if fallback is not None: @@ -152,6 +208,10 @@ def balanced_scaling( out[fallback] -= np.where(add_to_fallback, 1, -1)*dif out[orig] = out[orig].where(out[orig] > 0, 0) + # Convert back to original units if conversion array is provided + conversion_arr = conversion_arr.where(conversion_arr != 0, other=1) + out = out/conversion_arr + set_dict(datablock, out_key, out) return datablock @@ -328,3 +388,583 @@ def IDR( set_dict(datablock, out_key, idr) return datablock + + +@pipeline_node(["fbs", "conversion_arr"]) +def scale_above_threshold( + fbs, + scale, + element, + threshold=0, + items=None, + origin=None, + add_to_origin=True, + elasticity=None, + fallback=None, + add_to_fallback=True, + conversion_arr=None, +): + """Scales excess item quantities in a food balance sheet above a certain + threshold + + Parameters + ---------- + fbs : xarray.Dataset + Input food balance sheet Dataset + scale : float, xarray.DataArray + Scaling value or array to apply to the excess above the threshold. A + value of 1 means no scaling, while a value of 0 means scaling down to + the threshold. + element : string + Name of the DataArray to scale + threshold : float, xarray.DataArray, optional + Minimum value for the scaled element. Scaling of item quantities is + only applied to the excess above this threshold. + items : list, optional + List of items to scaled in the food balance sheet. If None, all items + are scaled. + origin : string, list, optional + Names of the DataArrays which will be used as source for the quantity + changes. Any change to the "element" DataArray will be reflected in + this DataArray + add_to_origin : bool, array, optional + Whether to add or subtract the difference from the respective origins + elasticity : float, array, optional + Relative fraction of the total difference to be assigned to each origin + element. Values are not normalized. + fallback : string + Name of the DataArray to use as fallback in case the origin quantities + fall below zero + add_to_fallback : bool, optional + Whether to add or subtract the difference below zero in the origin + DataArray to the fallback array. + conversion_arr : string, xarray.DataArray, tuple or float + Conversion array to pre-scale quantities. If provided, the input food + balance sheet is first converted using the conversion array, then the + scaling is applied, and finally the results are converted back to the + original units using the inverse of the conversion array. + """ + + if conversion_arr is not None: + if isinstance(conversion_arr, xr.DataArray): + conversion_arr = conversion_arr.where( + np.isfinite(conversion_arr), other=1) + if (conversion_arr == 0).any(): + warnings.warn("Conversion array contains zero values," \ + "which can lead to inaccurate scaling") + + elif isinstance(conversion_arr, (int, float)): + conversion_arr = xr.full_like( + fbs[element], fill_value=conversion_arr) + else: + conversion_arr = xr.ones_like(fbs[element]) + + fbs = fbs*conversion_arr + + if items is not None: + scaled_items = item_parser(fbs, items) + sel = {"Item": scaled_items} + else: + scaled_items = fbs.Item.values + sel = {} + + # Calculate the maximum scale threshold that would reduce the sum of the + # element to the threshold, and apply the input scale relative to this + max_scale_threshold = (fbs[element].sum(dim="Item") - threshold) \ + / fbs[element].sel(sel).sum(dim="Item") + + scale_threshold = max_scale_threshold * (1 - scale) + + if origin is not None and np.isscalar(origin): + origin = [origin] + + if np.isscalar(add_to_origin): + add_to_origin = [add_to_origin]*len(origin) + + if elasticity is None: + elasticity = [1/len(origin)]*len(origin) + + # Scale input + if origin is None: + out = fbs.fbs.scale_element( + element=element, + scale=1-scale_threshold, + items=scaled_items + ) + + else: + out = fbs.fbs.scale_add( + element_in=element, + element_out=origin, + scale=1-scale_threshold, + items=scaled_items, + add=add_to_origin, + elasticity=elasticity + ) + + # If a fallback DataArray is defined, transfer the excess negative + # quantities to it + if fallback is not None and origin is not None: + for orig in origin: + dif = out[orig].where(out[orig] < 0).fillna(0) + out[fallback] -= np.where(add_to_fallback, 1, -1)*dif + out[orig] = out[orig].where(out[orig] > 0, 0) + + # Convert back to original units if conversion array is provided + conversion_arr = conversion_arr.where(conversion_arr != 0, other=1) + out = out/conversion_arr + + return out + +@pipeline_node(["fbs", "population"]) +def project_by_population( + fbs, + population, + items=None, + food_element="food", + elasticity=0.5, + production_element="production", + imports_element="imports", + exports_element="exports", + scale_feed=False, + feed_element="feed", + feed_items=None, + scale_seed=False, + seed_element="seed", + seed_items=None, + scale_processing=False, + processing_element="processing", + processing_items=None, +): + """Project a food balance sheet forward using population and yield changes. + + The model scales selected food quantities by the relative population + change from the first available year and balances the resulting gap between + production and imports. Optionally, it adjusts feed, seed, and processing + quantities according to the resulting production change. + + Parameters + ---------- + fbs : xarray.Dataset + Input food balance sheet Dataset + population : xarray.DataArray or float + Population data array or scalar value to scale food quantities by. + items : list, optional + List of items to scale in the food balance sheet. If None, all items + are scaled. + food_element : string, optional + Name of the DataArray containing the food data + elasticity : float, optional + Relative fraction of the total difference to be assigned to production + and imports. + production_element : string, optional + Name of the DataArray containing the production data + imports_element : string, optional + Name of the DataArray containing the imports data + exports_element : string, optional + Name of the DataArray containing the exports data + scale_feed : bool, optional + Whether to scale feed quantities according to the production change + feed_element : string, optional + Name of the DataArray containing the feed data + feed_items : list, optional + List of items to monitor to do feed scaling. If None, all items are + monitored. + scale_seed : bool, optional + Whether to scale seed quantities according to the production change + seed_element : string, optional + Name of the DataArray containing the seed data + seed_items : list, optional + List of items to monitor to do seed scaling. If None, all items are + monitored. + scale_processing : bool, optional + Whether to scale processing quantities according to the production + change + processing_element : string, optional + Name of the DataArray containing the processing data + processing_items : list, optional + List of items to monitor to do processing scaling. If None, all items + are monitored. + """ + + out = copy.deepcopy(fbs).fillna(0) + fbs = fbs.fillna(0) + + if isinstance(population, xr.DataArray): + baseline_population = population.isel(Year=0) + pop_ratio = population / baseline_population + pop_ratio = pop_ratio.where(np.isfinite(pop_ratio), other=1.0) + + elif np.isscalar(population): + pop_ratio = population + + if items is None: + food_items = out.Item.values + else: + food_items = item_parser(out, items) + + out = out.fbs.scale_element( + element=food_element, + scale=pop_ratio, + items=food_items, + ) + + # Balance fbs from production change + baseline_gap = (fbs[food_element] + - (fbs[production_element] + + fbs[imports_element] + - fbs[exports_element])) + + projected_gap = (out[food_element] + - (out[production_element] + + out[imports_element] + - out[exports_element])) + + # Balance by keeping ingress/egress delta + partial_gap_delta = projected_gap - baseline_gap + + out[production_element] = out[production_element] + partial_gap_delta * elasticity + out[imports_element] = out[imports_element] + partial_gap_delta * (1 - elasticity) + + fbs_partial = copy.deepcopy(out) + + # Scales feed, seed, and processing according to total production change + def _scale_from_production_change(element_name, element_items): + """Scales the given element according to the production change ratio + for the given items. If element_items is None, all items are monitored. + """ + + if element_items is None: + parsed_items = out.Item.values + + else: + parsed_items = item_parser(out, element_items) + + production_ratio = ( + out[production_element].sel(Item=parsed_items).sum(dim="Item") + / fbs[production_element].sel(Item=parsed_items).sum(dim="Item") + ) + + # Handle cases where the denominator is zero (no production in baseline) + production_ratio = production_ratio.where( + np.isfinite(production_ratio), other=1.0) + + # Scale all items in the element by the production ratio + out[element_name] = out.fbs.scale_element( + element=element_name, + scale=production_ratio, + )[element_name] + + if scale_feed: + _scale_from_production_change(feed_element, feed_items) + if scale_seed: + _scale_from_production_change(seed_element, seed_items) + if scale_processing: + _scale_from_production_change(processing_element, processing_items) + + # Recalculate gap after scaling feed, seed, and processing + partial_baseline_gap = (fbs_partial[food_element] + + fbs_partial.get(feed_element, 0) + + fbs_partial.get(seed_element, 0) + + fbs_partial.get(processing_element, 0) + - (fbs_partial[production_element] + + fbs_partial[imports_element] + - fbs_partial[exports_element])) + + final_projected_gap = (out[food_element] + + out.get(feed_element, 0) + + out.get(seed_element, 0) + + out.get(processing_element, 0) + - (out[production_element] + + out[imports_element] + - out[exports_element])) + + final_gap_delta = final_projected_gap - partial_baseline_gap + + out[production_element] = out[production_element] + final_gap_delta * elasticity + out[imports_element] = out[imports_element] + final_gap_delta * (1 - elasticity) + + return out + + +@pipeline_node(["fbs"]) +def scale_production_yield( + fbs, + yield_scale, + items=None, + elasticity=0.5, + production_element="production", + imports_element="imports", + exports_element="exports", + scale_feed=False, + feed_element="feed", + feed_items=None, + scale_seed=False, + seed_element="seed", + seed_items=None, + scale_processing=False, + processing_element="processing", + processing_items=None, +): + + """Scale production quantities in a food balance sheet according to a yield + scale factor. + + Scales production by a scale factor, and balances the resulting + gap between exports and imports. Optionally, it adjusts feed, seed, and + processing quantities according to the resulting production change. + Land is assumed to remain constant, so the yield scale factor is applied to + production only. + + Parameters + ---------- + fbs : xarray.Dataset + Input food balance sheet Dataset + yield_scale : float, xarray.DataArray + Yield scale factor or array to apply to the production element. + A value of 1 means no scaling, while a value of 0 means scaling down + to zero. + items : list, optional + List of items to scale in the food balance sheet. If None, all items + are scaled. + elasticity : float, optional + Relative fraction of the total difference to be assigned to exports + and imports. + production_element : string, optional + Name of the DataArray containing the production data + imports_element : string, optional + Name of the DataArray containing the imports data + exports_element : string, optional + Name of the DataArray containing the exports data + scale_feed : bool, optional + Whether to scale feed quantities according to the production change + feed_element : string, optional + Name of the DataArray containing the feed data + feed_items : list, optional + List of items to monitor to do feed scaling. If None, all items are + monitored. + scale_seed : bool, optional + Whether to scale seed quantities according to the production change + seed_element : string, optional + Name of the DataArray containing the seed data + seed_items : list, optional + List of items to monitor to do seed scaling. If None, all items are + monitored. + scale_processing : bool, optional + Whether to scale processing quantities according to the production + change + processing_element : string, optional + Name of the DataArray containing the processing data + processing_items : list, optional + List of items to monitor to do processing scaling. If None, all items + are monitored. + """ + + from ..utils.scaling import linear_scale + + out = copy.deepcopy(fbs).fillna(0) + fbs = fbs.fillna(0) + + if items is None: + prod_items = out.Item.values + else: + prod_items = item_parser(out, items) + + if isinstance(yield_scale, xr.DataArray): + yield_scale = yield_scale.where(np.isfinite(yield_scale), other=1.0) + + out = out.fbs.scale_element( + element=production_element, + scale=yield_scale, + items=prod_items, + ) + + # Balance fbs from production change + partial_gap = ( + out[production_element] + + out[imports_element] + - out[exports_element] + ) + + baseline_gap = ( + fbs[production_element] + + fbs[imports_element] + - fbs[exports_element] + ) + + partial_gap_delta = partial_gap - baseline_gap + + out[exports_element] = ( + out[exports_element] + partial_gap_delta * elasticity) + out[imports_element] = ( + out[imports_element] - partial_gap_delta * (1 - elasticity)) + + fbs_partial = copy.deepcopy(out) + + # Scales feed, seed, and processing according to total production change + def _scale_from_production_change(element_name, element_items): + """Scales the given element according to the production change ratio + for the given items. If element_items is None, all items are monitored. + """ + + if element_items is None: + parsed_items = out.Item.values + + else: + parsed_items = item_parser(out, element_items) + + production_ratio = ( + out[production_element].sel(Item=parsed_items).sum(dim="Item") + / fbs[production_element].sel(Item=parsed_items).sum(dim="Item") + ) + + # Handle cases where the denominator is zero (0 production in baseline) + production_ratio = production_ratio.where( + np.isfinite(production_ratio), other=1.0) + + # Scale all items in the element by the production ratio + out[element_name] = out.fbs.scale_element( + element=element_name, + scale=production_ratio, + )[element_name] + + if scale_feed: + _scale_from_production_change(feed_element, feed_items) + if scale_seed: + _scale_from_production_change(seed_element, seed_items) + if scale_processing: + _scale_from_production_change(processing_element, processing_items) + + # Recalculate gap after scaling feed, seed, and processing + partial_gap = (fbs_partial.get(feed_element, 0) + + fbs_partial.get(seed_element, 0) + + fbs_partial.get(processing_element, 0) + - (fbs_partial.get(production_element, 0) + + fbs_partial.get(imports_element, 0) + - fbs_partial.get(exports_element, 0))) + + projected_gap = (out.get(feed_element, 0) + + out.get(seed_element, 0) + + out.get(processing_element, 0) + - (out.get(production_element, 0) + + out.get(imports_element, 0) + - out.get(exports_element, 0))) + + gap_delta = projected_gap - partial_gap + + out[production_element] = out[production_element] + gap_delta * elasticity + out[imports_element] = out[imports_element] + gap_delta * (1 - elasticity) + + return out + + + +@pipeline_node(["fbs", "land_current", "land_reference"]) +def food_scaling_from_land( + fbs, + land_current, + land_reference, + categories, + element, + items=None, + category_dim=None, + keep_elements_constant=False, + target_items=None, + origin=None, + add_to_origin=True, + elasticity=None, + fallback=None, + add_to_fallback=True, + conversion_arr=None, + ): + ''' + This function changes food quantities in a Food Balance Sheet array element + by scaling them relative to the change in a Land Data Array category + quantities. + + Parameters + ---------- + fbs : xarray.Dataset + Food balance sheet dataset containing the current food quantities. + land_current : xarray.DataArray + Land array containing the current land quantities. + land_reference : xarray.DataArray + Land array containing the reference land quantities. + categories : string, list + Name or list of land use category names to be monitored for relative + food scaling. + element : string + Name of the DataArray to scale. + items : list, optional + List of items to scaled in the food balance sheet. If None, all items + are scaled and 'constant' is ignored. + category_dim : string, optional + Name of the category dimension in the land datasets. If None, the first + non-spatial dimension is used. + keep_elements_constant : bool, optional + If set to True, the sum of element remains constant by scaling the + non selected items accordingly + categories : string, list, tuple, optional + Item or list of items to be used for scaling. If not provided, all + items are used. + target_items : list, optional + List of items to use for scaling when 'constant' is True. If None, all + non-selected items are used for scaling. + origin : string, list, optional + Names of the DataArrays which will be used as source for the quantity + changes. Any change to the "element" DataArray will be reflected in + this DataArray + add_to_origin : bool, array, optional + Whether to add or subtract the difference from the respective origins + elasticity : float, array, optional + Relative fraction of the total difference to be assigned to each origin + element. Values are not normalized. + fallback : string + Name of the DataArray to use as fallback in case the origin quantities + fall below zero + add_to_fallback : bool, optional + Whether to add or subtract the difference below zero in the origin + DataArray to the fallback array. + conversion_arr : string, xarray.DataArray, tuple or float + Conversion array to pre-scale quantities. If provided, the input food + balance sheet is first converted using the conversion array, then the + scaling is applied, and finally the results are converted back to the + original units using the inverse of the conversion array. + + Returns + ------- + scales_fbs : + ''' + # Use the first non spatial dimension if category dimension not provided + if category_dim is None: + category_dim = [d for d in land_current.dims if d not in ["Year", "x", "y"]][0] + + # Make categories into a list if not already + if np.isscalar(categories): + categories = [categories] + + # Obtain reference and current land quantities + ref_quantities = land_reference.sel({category_dim: categories})\ + .sum(dim=[category_dim, 'x', 'y']) + obs_quantities = land_current.sel({category_dim: categories})\ + .sum(dim=[category_dim, 'x', 'y']) + + # Compute scaling factor + scaling_factor = obs_quantities / ref_quantities + + scaled_fbs = balanced_scaling( + fbs=fbs, + scale=scaling_factor, + element=element, + items=items, + constant=keep_elements_constant, + padding_items=target_items, + origin=origin, + add_to_origin=add_to_origin, + elasticity=elasticity, + fallback=fallback, + add_to_fallback=add_to_fallback, + conversion_arr=conversion_arr, + ) + + return scaled_fbs \ No newline at end of file diff --git a/agrifoodpy/food/tests/test_food.py b/agrifoodpy/food/tests/test_food.py index e374f02..ff0c37e 100644 --- a/agrifoodpy/food/tests/test_food.py +++ b/agrifoodpy/food/tests/test_food.py @@ -471,9 +471,11 @@ def test_plot_bars(): # Test with a "show" dimension not in the coordinate list with pytest.raises(ValueError): ax_notincoord = fbs.plot_bars(show="NotInCoord") - -def test_plot_years(): - + +# Test for plot_years method + +@pytest.fixture +def sample_da(): years = [2019, 2010, 2011, 2012, 2013, 2014] regions = ['A', 'B', 'C'] @@ -482,33 +484,48 @@ def test_plot_years(): coords=[('Year', years), ('Region', regions)], dims=['Year', 'Region']) + + return da - fbs = FoodElementSheet(da) +@pytest.fixture +def sample_fbs(sample_da): + """Provides a FoodElementSheet initialized with sample_da.""" + return FoodElementSheet(sample_da) +def test_plot_years_default_call(sample_fbs, sample_da): + # Test default call - ax_default = fbs.plot_years() + ax_default = sample_fbs.plot_years() assert len(ax_default.lines) == 1 assert np.array_equal(ax_default.lines[0].get_ydata(), - da.sum(dim="Region").values) - assert np.array_equal(ax_default.lines[0].get_xdata(), da.Year.values) + sample_da.sum(dim="Region").values) + assert np.array_equal(ax_default.lines[0].get_xdata(), sample_da.Year.values) + + +def test_plot_years_with_show_coordinate(sample_fbs, sample_da): # Test with explicit "show" coordinate - ax_show = fbs.plot_years(show="Region") + ax_show = sample_fbs.plot_years(show="Region") assert len(ax_show.lines) == 3 for il, line in enumerate(ax_show.lines): assert np.array_equal(line.get_ydata(), - da.cumsum(dim="Region").isel(Region=il).values) - assert np.array_equal(line.get_xdata(), da.Year.values) + sample_da.cumsum(dim="Region").isel(Region=il).values) + assert np.array_equal(line.get_xdata(), sample_da.Year.values) + + +def test_plot_years_with_stack_option(sample_fbs, sample_da): # Test without stacking - ax_no_stack = fbs.plot_years(show="Region", stack=False) + ax_no_stack = sample_fbs.plot_years(show="Region", stack=False) assert len(ax_no_stack.lines) == 3 for il, line in enumerate(ax_no_stack.lines): - assert np.array_equal(line.get_ydata(), da.isel(Region=il).values) - assert np.array_equal(line.get_xdata(), da.Year.values) + assert np.array_equal(line.get_ydata(), sample_da.isel(Region=il).values) + assert np.array_equal(line.get_xdata(), sample_da.Year.values) + +def test_plot_years_without_year_dimension(sample_da): # Test with array wihtout "Year" dimension - da_noyear = da.sum(dim="Year") + da_noyear = sample_da.sum(dim="Year") fbs_noyear = FoodElementSheet(da_noyear) with pytest.raises(TypeError): ax_noyear = fbs_noyear.plot_years() \ No newline at end of file diff --git a/agrifoodpy/food/tests/test_model.py b/agrifoodpy/food/tests/test_model.py index 1da1db8..b01178b 100644 --- a/agrifoodpy/food/tests/test_model.py +++ b/agrifoodpy/food/tests/test_model.py @@ -1,6 +1,6 @@ import numpy as np import xarray as xr - +import pytest import warnings def test_balanced_scaling(): @@ -30,6 +30,27 @@ def test_balanced_scaling(): xr.testing.assert_equal(result_basic, fbs) + # Test without year dimension + fbs_no_year = fbs.isel(Year=0).drop_vars("Year") + + result_no_year = balanced_scaling( + fbs_no_year, + scale=2.0, + element="food" + ) + + ex_result_no_year = xr.Dataset( + data_vars=dict( + imports=(["Item"], [10., 20.]), + production=(["Item"], [50., 60.]), + exports=(["Item"], [5., 10.]), + food=(["Item"], [110., 140.]) + ), + coords=dict(Item=("Item", items)) + ) + + xr.testing.assert_equal(result_no_year, ex_result_no_year) + # Test result with scalar scaling factor result_scalar = balanced_scaling( fbs, @@ -224,7 +245,6 @@ def test_balanced_scaling(): imports=(["Year", "Item"], [[10., 20.], [30., 40.]]), production=(["Year", "Item"], [[50., 60.], [70., 80.]]), exports=(["Year", "Item"], [[5., 10.], [15., 20.]]), - # food=(["Year", "Item"], [[55., 70.], [85., 100.]]) food=(["Year", "Item"], [[55., 70.], [170., 100.]]) ), @@ -368,3 +388,634 @@ def test_fbs_convert(): ) xr.testing.assert_allclose(result_factor, ex_result_factor) + + +def test_scale_above_threshold(): + + from agrifoodpy.food.model import scale_above_threshold + from agrifoodpy.pipeline import Pipeline + + items = ["Beef", "Apples"] + years = [2020, 2021] + + fbs = xr.Dataset( + data_vars=dict( + imports=(["Year", "Item"], [[10., 20.], [30., 40.]]), + production=(["Year", "Item"], [[50., 60.], [70., 80.]]), + exports=(["Year", "Item"], [[5., 10.], [15., 20.]]), + food=(["Year", "Item"], [[55., 70.], [85., 100.]]) + ), + + coords=dict(Item=("Item", items), Year=("Year", years)) + ) + + # Test basic result with unity scaling factor + result_basic = scale_above_threshold( + fbs, + scale=1.0, + element="food", + ) + + xr.testing.assert_allclose(result_basic, fbs) + + # Test basic in a pipeline + test_pipeline = Pipeline({"data": fbs}) + + test_pipeline.add_node( + scale_above_threshold, + { + "fbs": "data", + "scale": 1.0, + "element": "food", + }) + + test_pipeline.run() + + assert "scale_above_threshold" in test_pipeline.datablock + xr.testing.assert_equal( + test_pipeline.datablock["scale_above_threshold"], fbs) + + # Test result with positive scaling factor less than 1 + result_scaled = scale_above_threshold( + fbs, + scale=0.5, + element="food", + threshold=60.0 + ) + ex_result = (fbs["food"].sum(dim="Item") - 60.0) * 0.5 + 60.0 + + xr.testing.assert_allclose( + result_scaled["food"].sum(dim="Item"), + ex_result) + + # Test result with positive scaling factor greater than 1 + result_scaled = scale_above_threshold( + fbs, + scale=2.0, + element="food", + threshold=60.0 + ) + ex_result = (fbs["food"].sum(dim="Item") - 60.0) * 2.0 + 60.0 + + xr.testing.assert_allclose(result_scaled["food"].sum(dim="Item"), + ex_result) + + # Test result with zero scaling factor + result_zero = scale_above_threshold( + fbs, + scale=0.0, + element="food", + threshold=60.0 + ) + ex_result_zero = (fbs["food"].sum(dim="Item") - 60.0) * 0.0 + 60.0 + xr.testing.assert_allclose( + result_zero["food"].sum(dim="Item"), + ex_result_zero) + + # Test result with array scale + scale_arr = xr.DataArray( + [0.1, 0.5], + dims=["Year"], + coords={"Year": years}, + ) + + result_array_scale = scale_above_threshold( + fbs, + scale=scale_arr, + element="food", + threshold=60.0 + ) + + ex_result_arr = (fbs["food"].sum(dim="Item") - 60.0) * scale_arr + 60.0 + xr.testing.assert_allclose( + result_array_scale["food"].sum(dim="Item"), + ex_result_arr) + + # Test with selected items + result_array_items = scale_above_threshold( + fbs, + scale=0.0, + element="food", + threshold=60.0, + items="Beef" + ) + + ex_result_sel = (fbs["food"].sum(dim="Item") - 60.0) * 0.0 + 60.0 + + xr.testing.assert_allclose( + result_array_items["food"].sum(dim="Item"), + ex_result_sel) + + # Test with array threshold + threshold_arr = xr.DataArray( + [50.0, 70.0], + dims=["Year"], + coords={"Year": years}, + ) + + result_array_threshold = scale_above_threshold( + fbs, + scale=0.5, + element="food", + threshold=threshold_arr + ) + + ex_result_threshold = (fbs["food"].sum(dim="Item") - threshold_arr) * 0.5 \ + + threshold_arr + + xr.testing.assert_allclose( + result_array_threshold["food"].sum(dim="Item"), + ex_result_threshold) + + # Test with origin + result_origin = scale_above_threshold( + fbs, + scale=0.5, + element="food", + threshold=60.0, + origin="production" + ) + + excess = fbs["food"] - result_origin["food"] + ex_result_origin = fbs["production"] - excess + xr.testing.assert_allclose( + result_origin["production"], ex_result_origin) + + # Test with multiple origins and separate elasticity values + result_elasticity = scale_above_threshold( + fbs, + scale=0.5, + element="food", + threshold=60.0, + origin=["production", "imports"], + elasticity=[0.8, 0.2] + ) + + excess = fbs["food"] - result_elasticity["food"] + ex_result_prod = fbs["production"] - excess * 0.8 + ex_result_imports = fbs["imports"] - excess * 0.2 + + xr.testing.assert_allclose( + result_elasticity["production"], ex_result_prod) + xr.testing.assert_allclose( + result_elasticity["imports"], ex_result_imports) + + # Test with conversion array + conversion_arr = xr.DataArray( + [1.0, 2.0], + dims=["Item"], + coords={"Item": items}) + + result_conversion = scale_above_threshold( + fbs, + scale=0.5, + element="food", + threshold=60.0, + conversion_arr=conversion_arr + ) + + conv_arr = fbs * conversion_arr + ex_result_conv = (conv_arr["food"].sum(dim="Item") - 60.0) * 0.5 + 60.0 + + xr.testing.assert_allclose( + (result_conversion["food"]*conversion_arr).sum(dim="Item"), + ex_result_conv) + + +def test_food_scaling_from_land_basic(): + from agrifoodpy.food.model import food_scaling_from_land + + # Food data + items = ["Beef", "Apples"] + food_coords = {"Item": items} + food_data = np.random.rand(len(food_coords["Item"])) + + fbs = xr.Dataset( + data_vars={"production": (["Item"], food_data)}, + coords=food_coords + ) + + # Land data + categories = ["Pasture", "Arable", "Forest"] + land_coords = { + "x": [0, 1, 2], + "y": [0, 1], + "category": categories, + } + + land_data = np.random.rand( + len(land_coords["x"]), + len(land_coords["y"]), + len(land_coords["category"])) + land_ref = xr.DataArray(land_data, coords=land_coords) + + land_obs = land_ref.copy(deep=True) + print(land_obs) + print() + print(fbs) + land_obs *= 2 + + result = food_scaling_from_land( + fbs=fbs, + land_current=land_obs, + land_reference=land_ref, + categories='Pasture', + element='production' + ) + + truth = fbs.copy(deep=True) + truth["production"] *= 2.0 + xr.testing.assert_allclose(result, truth) + + +def test_food_scaling_from_land_basic_change(): + from agrifoodpy.food.model import food_scaling_from_land + + # Food data + items = ["Beef", "Apples"] + food_coords = {"Item": items} + food_data = np.random.rand(len(food_coords["Item"])) + + fbs = xr.Dataset( + data_vars={"production": (["Item"], food_data)}, + coords=food_coords + ) + + # Land data + categories = ["Pasture", "Arable", "Forest"] + land_coords = { + "x": [0, 1, 2], + "y": [0, 1], + "category": categories, + } + + land_data = np.random.rand( + len(land_coords["x"]), + len(land_coords["y"]), + len(land_coords["category"])) + land_ref = xr.DataArray(land_data, coords=land_coords) + + land_obs = land_ref.copy(deep=True) + land_obs *= 2 + + result = food_scaling_from_land( + fbs=fbs, + land_current=land_obs, + land_reference=land_ref, + categories='Pasture', + element='production' + ) + + truth = fbs.copy(deep=True) + truth["production"] *= 2.0 + xr.testing.assert_allclose(result, truth) + + +def test_food_scaling_from_land_with_item(): + from agrifoodpy.food.model import food_scaling_from_land + + # Food data + items = ["Beef", "Apples"] + food_coords = {"Item": items} + food_data = np.random.rand(len(food_coords["Item"])) + + fbs = xr.Dataset( + data_vars={"production": (["Item"], food_data)}, + coords=food_coords + ) + + # Land data + categories = ["Pasture", "Arable", "Forest"] + land_coords = { + "x": [0, 1, 2], + "y": [0, 1], + "category": categories, + } + + land_data = np.random.rand( + len(land_coords["x"]), + len(land_coords["y"]), + len(land_coords["category"])) + land_ref = xr.DataArray(land_data, coords=land_coords) + + land_obs = land_ref.copy(deep=True) + print(land_obs) + land_obs *= 2 + + result = food_scaling_from_land( + fbs=fbs, + land_current=land_obs, + land_reference=land_ref, + categories='Pasture', + element='production', + items='Beef' + ) + + truth = fbs.copy(deep=True) + truth['production'] = truth['production'].where(truth['Item'] != 'Beef', + other=truth['production'] * 2) + xr.testing.assert_allclose(result, truth) + + +def test_food_scaling_from_land_target_items(): + from agrifoodpy.food.model import food_scaling_from_land + + # Food data + items = ["Beef", "Apples"] + food_coords = {"Item": items} + food_data = np.random.rand(len(food_coords["Item"])) + + fbs = xr.Dataset( + data_vars={"production": (["Item"], food_data)}, + coords=food_coords + ) + + # Land data + categories = ["Pasture", "Arable", "Forest"] + land_coords = { + "x": [0, 1, 2], + "y": [0, 1], + "category": categories, + } + + land_data = np.random.rand( + len(land_coords["x"]), + len(land_coords["y"]), + len(land_coords["category"])) + land_ref = xr.DataArray(land_data, coords=land_coords) + + land_scale = np.random.rand() * 0.5 + 1 + land_obs = land_ref.copy(deep=True) + land_obs *= land_scale + + result = food_scaling_from_land( + fbs=fbs, + land_current=land_obs, + land_reference=land_ref, + categories='Pasture', + element='production', + items = ['Beef'], + keep_elements_constant=True, + target_items=['Apples'] + ) + + truth = fbs.copy(deep=True) + truth["production"].loc[{'Item':'Beef'}] *= land_scale + truth["production"].loc[{"Item": "Apples"}] -= fbs["production"].loc[{"Item": "Beef"}] \ + * (land_scale - 1) + xr.testing.assert_allclose(result, truth) + + +def test_food_scaling_from_land_with_time_dependent_fbs(): + from agrifoodpy.food.model import food_scaling_from_land + + # Food data + items = ["Beef", "Apples"] + years = [2000, 2010] + food_coords = {"Item": items, "Year": years} + food_data = np.random.rand( + len(food_coords["Item"]), + len(food_coords["Year"]) + ) + fbs = xr.Dataset( + data_vars={"production": (["Item", "Year"], food_data)}, + coords=food_coords + ) + + # Land data + categories = ["Pasture", "Arable", "Forest"] + land_coords = { + "x": [0, 1, 2], + "y": [0, 1], + "category": categories, + } + + land_data = np.random.rand( + len(land_coords["x"]), + len(land_coords["y"]), + len(land_coords["category"])) + land_ref = xr.DataArray(land_data, coords=land_coords) + + land_scale = np.random.rand() * 0.5 + 1 + land_obs = land_ref.copy(deep=True) + land_obs *= land_scale + + result = food_scaling_from_land( + fbs=fbs, + land_current=land_obs, + land_reference=land_ref, + categories='Pasture', + element='production', + items='Beef' + ) + + truth = fbs.copy(deep=True) + truth["production"].loc[{'Item':'Beef'}] *= land_scale + xr.testing.assert_allclose(result, truth) + + +@pytest.fixture +def year_scale_data(): + years = [2020, 2021, 2022, 2023] + pop_data = np.array([1.0, 1.5, 2.0, 2.5]) + pop_coords = {"Year": years} + population = xr.DataArray(pop_data, coords=pop_coords, dims=["Year"]) + return population + + +@pytest.fixture +def food_data(): + items = ["Beef", "Wheat"] + years = [2020] + + fbs = xr.Dataset( + data_vars=dict( + imports=(["Year", "Item"], [[10., 20.]]), + production=(["Year", "Item"], [[50., 60.]]), + exports=(["Year", "Item"], [[5., 10.]]), + food=(["Year", "Item"], [[45., 50.]]), + feed=(["Year", "Item"], [[0., 10.]]), + seed=(["Year", "Item"], [[0., 4.]]), + processing=(["Year", "Item"], [[10., 6.]]) + ), + + coords=dict(Item=("Item", items), Year=("Year", years)) + ) + + fbs = fbs.fbs.add_years(years=[2021, 2022, 2023], projection="constant") + + return fbs + + +# project_by_population tests + +def test_project_by_population_basic(year_scale_data, food_data): + from agrifoodpy.food.model import project_by_population + + result = project_by_population( + fbs=food_data, + population=year_scale_data, + items='Beef' + ) + + expected_scale_factor = year_scale_data / year_scale_data.sel(Year=2020) + expected_result = food_data.copy(deep=True) + expected_result["food"].loc[{'Item':'Beef'}] *= expected_scale_factor + delta_food = expected_result["food"] - food_data["food"] + expected_result["production"] += delta_food * 0.5 + expected_result["imports"] += delta_food * 0.5 + + xr.testing.assert_allclose(result, expected_result) + + +def test_project_by_population_with_elasticity(year_scale_data, food_data): + + from agrifoodpy.food.model import project_by_population + elasticity = np.random.rand() + + result = project_by_population( + fbs=food_data, + population=year_scale_data, + items='Beef', + elasticity=elasticity + ) + + expected_scale_factor = year_scale_data / year_scale_data.sel(Year=2020) + expected_result = food_data.copy(deep=True) + expected_result["food"].loc[{'Item':'Beef'}] *= expected_scale_factor + delta_food = expected_result["food"] - food_data["food"] + expected_result["production"] += delta_food * elasticity + expected_result["imports"] += delta_food * (1-elasticity) + + xr.testing.assert_allclose(result, expected_result) + + +def test_project_by_population_with_feed_scaling(year_scale_data, food_data): + from agrifoodpy.food.model import project_by_population + + result = project_by_population( + fbs=food_data, + population=year_scale_data, + items='Beef', + scale_feed=True, + feed_items="Beef" + ) + + expected_food_factor = year_scale_data / year_scale_data.sel(Year=2020) + expected_result = food_data.copy(deep=True) + expected_result["food"].loc[{'Item':'Beef'}] *= expected_food_factor + + delta_food = expected_result["food"] - food_data["food"] + + expected_result["production"] += delta_food * 0.5 + expected_result["imports"] += delta_food * 0.5 + + expected_feed_factor = ( + expected_result["production"].loc[{'Item':'Beef'}] + / food_data["production"].loc[{'Item':'Beef'}]) + + expected_result["feed"] *= expected_feed_factor + delta_feed = expected_result["feed"] - food_data["feed"] + + expected_result["production"] += delta_feed * 0.5 + expected_result["imports"] += delta_feed * 0.5 + + xr.testing.assert_allclose(result, expected_result) + + +# scale_production_yield tests + +def test_scale_production_yield_basic(food_data): + from agrifoodpy.food.model import scale_production_yield + + yield_scale = np.random.rand() + 0.5 + + result = scale_production_yield( + fbs=food_data, + yield_scale=yield_scale, + ) + + expected_result = food_data.copy(deep=True) + expected_result["production"] *= yield_scale + delta_production = expected_result["production"] - food_data["production"] + + expected_result["exports"] += delta_production * 0.5 + expected_result["imports"] -= delta_production * 0.5 + + xr.testing.assert_allclose( + result["production"], expected_result["production"]) + + +def test_scale_production_yield_with_elasticity(food_data): + from agrifoodpy.food.model import scale_production_yield + + elasticity = np.random.rand() + yield_scale = np.random.rand() + 0.5 + + result = scale_production_yield( + fbs=food_data, + yield_scale=yield_scale, + elasticity=elasticity + ) + + expected_result = food_data.copy(deep=True) + expected_result["production"] *= yield_scale + delta_production = expected_result["production"] - food_data["production"] + + expected_result["exports"] += delta_production * elasticity + expected_result["imports"] -= delta_production * (1 - elasticity) + + xr.testing.assert_allclose(result, expected_result) + + +def test_scale_production_yield_with_feed_scaling(year_scale_data, food_data): + from agrifoodpy.food.model import scale_production_yield + + yield_scale = np.random.rand() + 0.5 + + result = scale_production_yield( + fbs=food_data, + yield_scale=yield_scale, + scale_feed=True, + feed_items="Beef" + ) + + expected_result = food_data.copy(deep=True) + expected_result["production"] *= yield_scale + + delta_production = expected_result["production"] - food_data["production"] + + expected_result["exports"] += delta_production * 0.5 + expected_result["imports"] -= delta_production * 0.5 + + expected_feed_factor = ( + expected_result["production"].loc[{'Item':'Beef'}] + / food_data["production"].loc[{'Item':'Beef'}]) + + expected_result["feed"] *= expected_feed_factor + delta_feed = expected_result["feed"] - food_data["feed"] + + expected_result["production"] += delta_feed * 0.5 + expected_result["imports"] += delta_feed * 0.5 + + xr.testing.assert_allclose(result, expected_result) + + +def test_scale_production_yield_with_array_scaling(year_scale_data, food_data): + from agrifoodpy.food.model import scale_production_yield + + yield_scale = year_scale_data + + result = scale_production_yield( + fbs=food_data, + yield_scale=yield_scale, + ) + + expected_result = food_data.copy(deep=True) + expected_result["production"] *= yield_scale + delta_production = expected_result["production"] - food_data["production"] + + expected_result["exports"] += delta_production * 0.5 + expected_result["imports"] -= delta_production * 0.5 + + xr.testing.assert_allclose(result, expected_result) diff --git a/agrifoodpy/land/land.py b/agrifoodpy/land/land.py index 4fa02d3..80f60b8 100644 --- a/agrifoodpy/land/land.py +++ b/agrifoodpy/land/land.py @@ -6,7 +6,7 @@ import matplotlib.pyplot as plt from matplotlib import colors as mcolors import matplotlib.patches as mpatches - +from warnings import warn @xr.register_dataarray_accessor("land") class LandDataArray: @@ -16,17 +16,26 @@ def __init__(self, xarray_obj): @staticmethod def _validate(obj): - """Validate land xarray, checking it is a DataSet and has the minimum + """Validate land xarray, checking it is a DataArray and has the minimum set of coordinates """ if not isinstance(obj, xr.DataArray): - raise TypeError("Land array must be an xarray.Dataarray") + raise TypeError("Land array must be an xarray.DataArray") if "x" not in obj.dims or "y" not in obj.dims: raise AttributeError("Land array must have 'x' and 'y' dimensions") - - def plot(self, ax=None, class_coord=None, colors=None, labels=None, - **kwargs): + + + def plot( + self, + ax=None, + category_dim=None, + colors=None, + labels=None, + legend=False, + force_categorical=False, + **kwargs + ): """Plot a LandDataArray Generates a plot of a LandDataArray using matplotlib imshow, without @@ -39,16 +48,20 @@ def plot(self, ax=None, class_coord=None, colors=None, labels=None, ---------- ax : matplotlib.pyplot.Artist Axes on which to draw the plot - class_coord : string - Name of the coordinate to use as land classes. If not provided, the - first non spatial coordinate is used. + category_dim : string + Name of the dimension to use as land category. If not provided, the + first non spatial dimension is used. colors : list of strings Dictionary of colors to use for each land class. If not provided, the default matplotlib colour map is used. - labels : list of strings - Dictionary of labels to use for each land class. If not provided - and the map is a class percentage map, the coordinate values are - used as labels. + labels : list of strings, dict + List or dictionary of labels to use for each land class. If not + provided and the map is a class percentage map, the coordinate + values are used as labels. + legend : bool + If True, and data is determined to be categorical, a legend is + added to the plot. If data is not categorical, a colorbar is + added instead. **kwargs : dict Style options to be passed to the imshow function. @@ -57,6 +70,14 @@ def plot(self, ax=None, class_coord=None, colors=None, labels=None, ax : matplotlib axes instance """ + if "class_coord" in kwargs: + warn("class_coord is deprecated as a keyword argument. " + "Please use category_dim instead.", + DeprecationWarning, + stacklevel=2) + + category_dim = kwargs.pop("class_coord") + map = self._obj if ax is None: @@ -65,48 +86,87 @@ def plot(self, ax=None, class_coord=None, colors=None, labels=None, # Check the map type by checking for additional coordinates extra_coords = [dim for dim in map.dims if dim not in ["x", "y"]] if len(extra_coords) >= 1: + map = self.dominant_category(category_dim=category_dim) + + unique_vals = np.unique(map.values[~np.isnan(map.values)]) + + is_discrete = True if len(unique_vals) < 20 else False + if force_categorical: + is_discrete = True + + if is_discrete: + if colors is None: + colors = [f"C{i}" for i in np.arange(len(unique_vals))] + if labels is None: - labels = map[extra_coords[0]].values - map = self.dominant_class(class_coord=class_coord) - else: - labels = np.unique(map.values) + # labels = [str(val) for val in unique_vals] + labels = {val: str(val) for val in unique_vals} - if colors is None: - colors = [f"C{i}" for i in np.arange(len(labels))] + # Create a discrete colour map + cmap = mcolors.ListedColormap(colors) + bounds = np.linspace(-0.5, len(colors), len(colors) + 1) + norm = mcolors.BoundaryNorm(bounds, cmap.N) - # Create a colour map - cmap = mcolors.ListedColormap(colors) - bounds = np.linspace(-0.5, len(colors), len(colors) + 1) - norm = mcolors.BoundaryNorm(bounds, cmap.N) + # Get plot ranges + dx_low = (map.x.values[1] - map.x.values[0])/2 + dx_high = (map.x.values[-1] - map.x.values[-2])/2 + dy_low = (map.y.values[1] - map.y.values[0])/2 + dy_high = (map.y.values[-1] - map.y.values[-2])/2 - # Get plot ranges - dx_low = (map.x.values[1] - map.x.values[0])/2 - dx_high = (map.x.values[-1] - map.x.values[-2])/2 - dy_low = (map.y.values[1] - map.y.values[0])/2 - dy_high = (map.y.values[-1] - map.y.values[-2])/2 + xmin, xmax = map.x.values[[0, -1]] + ymin, ymax = map.y.values[[0, -1]] - xmin, xmax = map.x.values[[0, -1]] - ymin, ymax = map.y.values[[0, -1]] + patches = [mpatches.Patch(color=colors[i], + label=labels[val]) + for i, val in enumerate(unique_vals)] - ax.imshow(map, interpolation="none", origin="lower", - extent=[xmin-dx_low, - xmax+dx_high, - ymin-dy_low, - ymax+dy_high], - cmap=cmap, norm=norm) + else: + cmap = kwargs.pop("cmap", "viridis") + norm = kwargs.pop("norm", None) + + # Get plot ranges + dx_low = (map.x.values[1] - map.x.values[0])/2 + dx_high = (map.x.values[-1] - map.x.values[-2])/2 + dy_low = (map.y.values[1] - map.y.values[0])/2 + dy_high = (map.y.values[-1] - map.y.values[-2])/2 - patches = [mpatches.Patch(color=colors[i], - label=labels[i]) - for i in np.arange(len(labels))] + xmin, xmax = map.x.values[[0, -1]] + ymin, ymax = map.y.values[[0, -1]] - ax.legend(handles=patches, loc="best") + ax.imshow(map, interpolation="none", origin="lower", + extent=[xmin-dx_low, + xmax+dx_high, + ymin-dy_low, + ymax+dy_high], + cmap=cmap, norm=norm) + + if legend: + if is_discrete: + ax.legend(handles=patches, loc="best") + else: + plt.colorbar(ax.imshow(map, interpolation="none", origin="lower", + extent=[xmin-dx_low, + xmax+dx_high, + ymin-dy_low, + ymax+dy_high], + cmap=cmap, norm=norm), ax=ax) return ax + - def area_by_type( + def area_by_type(self, values=None, dim=None): + warn("The area_by_type method is deprecated and will be removed in a " + "future version. Use area_by_category instead.", + DeprecationWarning, + stacklevel=2) + + return self.area_by_category(categories=values, dim=dim) + + + def area_by_category( self, - values=None, - dim=None + categories=None, + dim=None, ): """Area per map category in a LandDataArray @@ -115,9 +175,9 @@ def area_by_type( Parameters ---------- - values : int, array - List of category types to return the total area for. If not set, - the function returns areas for all values found on the map, + categories : int, array + List of categories to return the total area for. If not set, + the function returns areas for all categories found on the map, excluding nan values. dim : string Name to assign to the categories coordinate. If not set, the input @@ -136,27 +196,29 @@ def area_by_type( if dim is None: dim = map.name - if values is None: - values = np.unique(map) + if categories is None: + categories = np.unique(map) else: - values = np.array(values) + categories = np.array(categories) # Prevent nan values from being counted - nan_indices = np.isnan(values) - values = values[~nan_indices] - area = [ones.where(map == value).sum() for value in values] + nan_indices = np.isnan(categories) + categories = categories[~nan_indices] + area = [ones.where(map == value).sum() for value in categories] - area_arr = xr.DataArray(area, dims=dim, coords={dim: values}) + area_arr = xr.DataArray(area, dims=dim, coords={dim: categories}) return area_arr + def area_overlap( self, map_right, - values_left=None, - values_right=None, + categories_left=None, + categories_right=None, dim_left=None, - dim_right=None + dim_right=None, + **kwargs ): """Area overlap of selected categories between two maps @@ -168,14 +230,14 @@ def area_overlap( ---------- map_right : xarray.DataArray LandDataArray style DataArray to compare overlapping areas with - values_left - List of category types from the left map to return the total area + categories_left : int, array + List of land categories from the left map to return the total area overlaps for. - If not set, all category types are used, except nan values. - values_right : int, array - List of category types from the right map to return the total area + If not set, all categories are used, except nan values. + categories_right : int, array + List of land categories from the right map to return the total area overlaps for. - If not set, all category types are used, except nan values. + If not set, all categories are used, except nan values. dim_left : string Names to assign to the category coordinates on the output DataArray. @@ -191,6 +253,25 @@ def area_overlap( Array with the corresponding areas for each category type """ + + if "values_left" in kwargs: + warn("values_left is deprecated and will be removed in a future " + "version. Please use categories_left instead.", + DeprecationWarning, + stacklevel=2) + + categories_left = kwargs.pop("values_left") + + if "values_right" in kwargs: + warn("values_right is deprecated and will be removed in a future " + "version. Please use categories_right instead.", + DeprecationWarning, + stacklevel=2) + + categories_right = kwargs.pop("values_right") + + + map_left = self._obj # Check that both maps have the same dimensions and coordinates. @@ -203,40 +284,42 @@ def area_overlap( if dim_right is None: dim_right = map_right.name - if values_left is None: - values_left = np.unique(map_left) + if categories_left is None: + categories_left = np.unique(map_left) else: - values_left = np.array(values_left) + categories_left = np.array(categories_left) - if values_right is None: - values_right = np.unique(map_right) + if categories_right is None: + categories_right = np.unique(map_right) else: - values_right = np.array(values_right) + categories_right = np.array(categories_right) # Prevent nan values from being counted - nan_indices_left = np.isnan(values_left) - nan_indices_right = np.isnan(values_right) + nan_indices_left = np.isnan(categories_left) + nan_indices_right = np.isnan(categories_right) - values_left = values_left[~nan_indices_left] - values_right = values_right[~nan_indices_right] + categories_left = categories_left[~nan_indices_left] + categories_right = categories_right[~nan_indices_right] ones = xr.ones_like(map_left) area = [[ones.where(map_left == vl).where(map_right == vr).sum().values - for vr in values_right] for vl in values_left] + for vr in categories_right] for vl in categories_left] area_arr = xr.DataArray(area, dims=[dim_left, dim_right], - coords={dim_left: values_left, - dim_right: values_right}) + coords={dim_left: categories_left, + dim_right: categories_right}) return area_arr + def category_match( self, map_right, - values_left=None, - values_right=None, - join="left" + categories_left=None, + categories_right=None, + join="left", + **kwargs ): """Returns a land Dataarray with values where a selected overlap occurs between categories from two maps. This returns the values from the left @@ -259,42 +342,77 @@ def category_match( Land DataArray with values from the left map where overlap occurs. All other positions are assign a nan value. """ + + if "values_left" in kwargs: + warn("values_left is deprecated and will be removed in a future " + "version. Please use categories_left instead.", + DeprecationWarning, + stacklevel=2) + + categories_left = kwargs.pop("values_left") + + if "values_right" in kwargs: + warn("values_right is deprecated and will be removed in a future " + "version. Please use categories_right instead.", + DeprecationWarning, + stacklevel=2) + + categories_right = kwargs.pop("values_right") + map_left = self._obj # Check input values - if values_left is None: - values_left = np.unique(map_left) + if categories_left is None: + categories_left = np.unique(map_left) else: - values_left = np.array(values_left) + categories_left = np.array(categories_left) - if values_right is None: - values_right = np.unique(map_right) + if categories_right is None: + categories_right = np.unique(map_right) else: - values_right = np.array(values_right) + categories_right = np.array(categories_right) # Align maps to the left so they have the same dimensions map_left, map_right = xr.align(map_left, map_right, join=join) shape = map_left.shape - left_match = np.isin(map_left, values_left).reshape(shape) - right_match = np.isin(map_right, values_right).reshape(shape) + left_match = np.isin(map_left, categories_left).reshape(shape) + right_match = np.isin(map_right, categories_right).reshape(shape) category_match = map_left.where(left_match).where(right_match) return category_match + def dominant_class( self, class_coord=None, - return_index=False + return_index=False, + **kwargs + ): + warn("The dominant_class method is deprecated and will be removed in" \ + " a future version. Please use dominant_category instead.", + DeprecationWarning, + stacklevel=2) + + return self.dominant_category( + category_dim=class_coord, + return_index=return_index, + **kwargs) + + def dominant_category( + self, + category_dim=None, + return_index=False, + **kwargs ): """Returns a land DataArray with the dominant land class for each pixel. Parameters ---------- - class_coord : string + category_dim : string Name of the land class coordinate. If not set, the first coordinate is used. return_index : bool @@ -307,15 +425,134 @@ class value. Land DataArray with the dominant land class for each pixel. """ + if "class_coord" in kwargs: + warn("class_coord is deprecated as a keyword argument. " + "Please use category_dim instead.", + DeprecationWarning, + stacklevel=2) + + category_dim = kwargs.pop("class_coord") + map = self._obj - if class_coord is None: - class_coord = [dim for dim in map.dims if dim not in ["x", "y"]][0] + if category_dim is None: + category_dim = [dim for dim in map.dims if dim not in ["Year", "x", "y"]][0] if return_index: - len_class = len(map[class_coord].values) - map = map.assign_coords({class_coord: np.arange(len_class)}) + len_class = len(map[category_dim].values) + map = map.assign_coords({category_dim: np.arange(len_class)}) - map = map.idxmax(dim=class_coord, skipna=True) + map = map.idxmax(dim=category_dim, skipna=True) return map + + + def add_category( + self, + category, + category_value=0, + mask=None, + category_dim=None, + ): + """Add a new land category to a LandDataArray + + Parameters + ---------- + category : string, list of strings + Name of the land class coordinate to add. + Cannot contain values already present in the LandDataArray, + or duplicated within the input list. + category_value : int, xarray.DataArray + Value of the new land class to add. If an array is provided, it + must have the same shape as the spatial dimensions of the + LandDataArray. + mask : xarray.DataArray + Boolean DataArray with the same spatial dimensions as the + LandDataArray indicating where to add the new category. + If not provided, the new category is added to all pixels. + category_dim : string + Name of the land class dimension. If not set, the first non spatial + dimension is used. + + Returns + ------- + xarray.DataArray + Land DataArray with the new land category added. + """ + + map = self._obj + + # Use the first non spatial dimension if category dimension not provided + if category_dim is None: + category_dim = [d for d in map.dims if d not in ["x", "y", "Year"]][0] + + + # Validate the mask input + if mask is None: + mask = xr.ones_like(map.isel({category_dim:0}), dtype=bool) + else: + if not isinstance(mask, xr.DataArray): + raise TypeError("'mask' must be an xarray.DataArray") + + if set(mask.dims) not in ({"x", "y"}, {"x", "y", "Year"}): + raise ValueError( + "'mask' must have spatial dimensions 'x' and 'y', with " + "an optional 'Year' dimension" + ) + + spatial_template = map.isel({category_dim: 0}) + try: + xr.align(spatial_template, mask, join="exact") + except ValueError as err: + raise ValueError( + "'mask' must be aligned with the map spatial coordinates " + "'x' and 'y'" + ) from err + + # Validate the category_value input + if isinstance(category_value, (int, float)): + category_value = xr.ones_like( + map.isel({category_dim:0}))*category_value + + elif isinstance(category_value, xr.DataArray): + if set(category_value.dims) not in ({"x", "y"}, {"x", "y", "Year"}): + raise ValueError( + "'category_value' must have spatial dimensions 'x' and " + "'y', with an optional 'Year' dimension" + ) + + spatial_template = map.isel({category_dim: 0}) + try: + xr.align(spatial_template, category_value, join="exact") + except ValueError as err: + raise ValueError( + "'category_value' must be aligned with the map spatial " \ + "coordinates 'x' and 'y'" + ) from err + else: + raise TypeError("'category_value' must be either a scalar or an " \ + "xarray.DataArray with 'x' and 'y' dimensions, and an optional " + "'Year' dimension") + + new_map = map.copy() + + if isinstance(category, str): + category = [category] + + # Check that there are no duplicate category values in the input list + if len(category) != len(set(category)): + raise ValueError("Duplicate category values found in input list") + + # Check that the new category does not already exist in the map + for cat in category: + if cat in map[category_dim].values: + raise ValueError(f"Category {cat} already exists in the" + " LandDataArray") + + # Append new category values to the map + new_class = xr.ones_like(map.isel({category_dim:0}))*category_value + new_class = new_class.where(mask) + new_class[category_dim] = cat + new_map = xr.concat([new_map, new_class], dim=category_dim) + + return new_map diff --git a/agrifoodpy/land/model.py b/agrifoodpy/land/model.py index 0134455..11bb18e 100644 --- a/agrifoodpy/land/model.py +++ b/agrifoodpy/land/model.py @@ -1,10 +1,12 @@ """ Module for land intervention models - """ import numpy as np +import xarray as xr from ..land.land import LandDataArray - +from ..pipeline import pipeline_node +from ..utils.dict_utils import item_parser +import warnings def land_sequestration( land_da, @@ -66,7 +68,7 @@ def land_sequestration( if not (fraction >= 0).all() and (fraction <= 1).all(): raise ValueError("Input fraction values must be between 0 and 1") - pixel_count_category = land_da.land.area_by_type(values=use_id) + pixel_count_category = land_da.land.area_by_category(categories=use_id) # area in hectares area_category = pixel_count_category * ha_per_pixel @@ -97,3 +99,246 @@ def land_sequestration( total_seq = area.sum() * scale * max_seq return total_seq + +@pipeline_node("land") +def scale_land( + land, + origin, + fraction, + keep_land_constant=False, + target_category=None, + target_distribution=None, + category_dim=None, +): + """Convert land from one category to another. + + Given a Land Data Array map with pixel values for the different land use + types, the model computes the new land use map after conversion of a + fraction of the per-pixel utilisation from a set of categories to a target + category. + + Parameters + ---------- + land : xarray.Dataarray + Input land array containing the id id value with the dominant land use + type on that pixel. + origin : str, array + Land category identifiers for the land uses to be converted. + fraction : float, xarray.DataArray + Conversion fraction of each repurposed land category. Can be a single + scalar value, or an xarray DataArray with the same length as the + ``origin`` land categories, and optionally, include a time dimension + for dynamic conversion. + keep_land_constant : bool, optional + If True, the total land area is kept constant by scaling the other + categories proportionally to the change in the target category. If + False, only the target category is scaled, allowing the total land area + to change. + target_category : str, optional + If keep_land_constant is True, this defines the land category + identifiers to convert land to in order to keep land constant. If not + provided, the converted land is distributed across all non-origin + categories. + target_distribution : array-like, xarray.DataArray, optional + Relative distribution used to allocate converted land across target + categories. If array-like, it must have the same length as ``target``. + If a DataArray, it must include ``category_dim`` with coordinates equal + to ``target_category`` and can optionally include a ``Year`` dimension + for time-varying distribution. If not provided, the converted land is + distributed equally across target categories. + category_dim : str + Name of the land category dimension in the input land DataArray. + + Returns + ------- + new_land : xarray.DataArray + DataArray with the new land use map after conversion. + """ + + if isinstance(origin, str): + origin = [origin] + + + # Use the first non spatial dimension if category dimension not provided + if category_dim is None: + category_dim = [d for d in land.dims if d not in ["Year", "x", "y"]][0] + + for orig in origin: + if orig not in land[category_dim].values: + raise ValueError(f"Origin category {orig} not found in land" + " categories") + + new_land = land.copy() + + # Validate fraction values + if np.isscalar(fraction): + if fraction < 0: + warnings.warn("Fraction values must be equal to or greater than 0") + else: + if np.any((fraction < 0)): + warnings.warn("Fraction values must be equal to or greater than 0") + + # Scale land + delta_land = land.sel({category_dim: origin}) * (1 - fraction) + + new_land.loc[{category_dim: origin}] -= delta_land + + # If keeping land constant, distribute the converted land across targets + if keep_land_constant: + if target_category is None: + target_category = [c for c in land[category_dim].values if c not in origin] + + elif isinstance(target_category, str): + target_category = [target_category] + + # Check that target class exist. If not, add them to the land categories + for targ in target_category: + if targ not in land[category_dim].values: + new_land = new_land.land.add_category(targ, category_dim=category_dim) + + # Validate and normalize target distribution + if target_distribution is None: + target_distribution = np.ones(len(target_category), dtype=float) + + if isinstance(target_distribution, xr.DataArray): + if category_dim not in target_distribution.dims: + raise ValueError( + "target_distribution DataArray must include the category " + f"dimension '{category_dim}'" + ) + + extra_dims = set(target_distribution.dims) - {category_dim, "Year"} + if extra_dims: + raise ValueError( + "target_distribution DataArray supports only category and " + f"optional Year dimensions, got extra dimensions: {extra_dims}" + ) + + dist_categories = set(target_distribution[category_dim].values.tolist()) + target_categories = set(target_category) + if dist_categories != target_categories: + raise ValueError( + "target_distribution category coordinates must match target " + f"categories exactly. Expected {target_category}, got " + f"{list(target_distribution[category_dim].values)}" + ) + + target_distribution = target_distribution.sel({category_dim: target_category}) + + if ((target_distribution < 0) | (~np.isfinite(target_distribution))).any(): + raise ValueError("target_distribution values must be finite and" \ + " non-negative") + + dist_sum = target_distribution.sum(dim=category_dim) + if (dist_sum <= 0).any(): + raise ValueError( + "target_distribution must have a strictly positive total " + "across target categories" + ) + + target_distribution = target_distribution / dist_sum + else: + target_distribution = np.asarray(target_distribution, dtype=float) + if target_distribution.ndim != 1 or target_distribution.size != len(target_category): + raise ValueError( + "target_distribution must be a 1D array with the same length " + "as target" + ) + + if np.any((target_distribution < 0) | (~np.isfinite(target_distribution))): + raise ValueError("target_distribution values must be finite and" \ + "non-negative") + + dist_sum = target_distribution.sum() + if dist_sum <= 0: + raise ValueError( + "target_distribution must have a strictly positive total " + "across target categories" + ) + + target_distribution = xr.DataArray( + target_distribution / dist_sum, + dims=[category_dim], + coords={category_dim: target_category}, + ) + + total_delta = delta_land.sum(dim=category_dim) + distributed_delta = total_delta * target_distribution + + new_land.loc[{category_dim: target_category}] += distributed_delta.sel({category_dim: target_category}) + + return new_land + + +@pipeline_node(["land", "fbs", "fbs_reference"]) +def land_scaling_from_food( + land, + fbs, + fbs_reference, + element, + category, + items=None, + keep_land_constant=False, + target_category=None, + target_distribution=None, + category_dim=None, +): + """Scale land categories based on relative changes to food quantities. + + Given two food balance sheets, this model uses the relative change in + selected items quantities to scale land categories in a land data array. + + Parameters + ---------- + land : xarray.DataArray + Input land array containing the land categories to be scaled. + fbs : xarray.DataSet + Food balance sheet dataset containing the current food quantities. + fbs_reference : xarray.DataSet + Food balance sheet dataset containing the reference food quantities. + element : string + Name of the fbs element to obtain the food quantities from. + category : string, list + Name or list of names of the land use categories to be scaled. + items : string, list, tuple, optional + Item or list of items to be used for scaling. If not provided, all + items are used. + keep_land_constant : bool, optional + If True, the total land area is kept constant by scaling the other + categories proportionally to the change in the target category. If False, + only the target category is scaled, allowing the total land area to change. + target_category : string, list, optional + Name of the land category to be used as padding when keep_land_constant + is True. + target_distribution : array-like, xarray.DataArray, optional + Relative distribution used to allocate land changes across target + categories when keep_land_constant is True. + category_dim : string, optional + Name of the dimension along which the land use categories are defined. + If not provided, the first non spatial dimension is used. + """ + + # Obtain reference and current food quantities + if items is not None: + items = item_parser(fbs_reference, items) + else: + items = fbs_reference.Item.values + + ref_quantities = fbs_reference[element].sel(Item=items).sum(dim="Item") + obs_quantities = fbs[element].sel(Item=items).sum(dim="Item") + + # Compute scaling factor + scaling_factor = obs_quantities / ref_quantities + + out_land = scale_land( + land, + origin=category, + fraction=scaling_factor, + keep_land_constant=keep_land_constant, + target_category=target_category, + target_distribution=target_distribution, + category_dim=category_dim, + ) + + return out_land + \ No newline at end of file diff --git a/agrifoodpy/land/tests/test_land.py b/agrifoodpy/land/tests/test_land.py index 5d71e12..577ba74 100644 --- a/agrifoodpy/land/tests/test_land.py +++ b/agrifoodpy/land/tests/test_land.py @@ -4,7 +4,7 @@ import pytest import matplotlib.pyplot as plt -def test_area_by_type(): +def test_area_by_category(): data = np.tile((np.arange(3, dtype=float)), (4,1)) @@ -12,36 +12,36 @@ def test_area_by_type(): coords={"x": [0, 1, 2, 3], "y": [0, 1, 2]}) land = LandDataArray(da) - # Test area by type without explicit values - expected_result_no_values = np.array([4,4,4]) - result_no_values = land.area_by_type() + # Test area by category without explicit categories + expected_result_no_categories = np.array([4,4,4]) + result_no_categories = land.area_by_category() - assert(np.array_equal(expected_result_no_values, result_no_values)) - assert(list(result_no_values.coords) == ["land"]) + assert(np.array_equal(expected_result_no_categories, result_no_categories)) + assert(list(result_no_categories.coords) == ["land"]) - # Test area by type with explicit values - expected_result_values = np.array([4,4]) - result_values = land.area_by_type(values=[1,2]) + # Test area by category with explicit categories + expected_result_categories = np.array([4,4]) + result_categories = land.area_by_category(categories=[1,2]) - assert(np.array_equal(expected_result_values, result_values)) - assert(list(result_values.coords) == ["land"]) + assert(np.array_equal(expected_result_categories, result_categories)) + assert(list(result_categories.coords) == ["land"]) - # Test area by type with explicit dimension name + # Test area by category with explicit dimension name dim_name = "area" expected_result_dims = np.array([4,4,4]) - result_dims = land.area_by_type(dim=dim_name) + result_dims = land.area_by_category(dim=dim_name) assert(np.array_equal(expected_result_dims, result_dims)) assert(list(result_dims.coords) == [dim_name]) - # Test area by type with nan values + # Test area by category with nan categories da_nan = da.copy(deep=True) da_nan[-1, -1] = np.nan da_nan.name = "land_nan" land_nan = LandDataArray(da_nan) expected_result_nan = np.array([4,4,3]) - result_nan = land_nan.area_by_type() + result_nan = land_nan.area_by_category() assert(np.array_equal(expected_result_nan, result_nan)) assert(list(result_nan.coords) == ["land_nan"]) @@ -62,27 +62,27 @@ def test_area_overlap(): land_left = LandDataArray(da_left) land_right = LandDataArray(da_right) - # Test area overlap without explicit values + # Test area overlap without explicit categories - expected_results_no_values = np.ones((3,4)) - result_no_values = land_left.area_overlap(da_right) + expected_results_no_cats = np.ones((3,4)) + result_no_cats = land_left.area_overlap(da_right) - assert(np.array_equal(expected_results_no_values, result_no_values)) - assert(list(result_no_values.coords) == [name_left, name_right]) + assert(np.array_equal(expected_results_no_cats, result_no_cats)) + assert(list(result_no_cats.coords) == [name_left, name_right]) - # Test area overlap with explicit values - input_values_left = [1,2] - input_values_right = [2,3] + # Test area overlap with explicit categories + input_cat_left = [1,2] + input_cat_right = [2,3] - expected_results_values = np.ones((2,2)) - result_values = land_left.area_overlap(da_right, - values_left=input_values_left, - values_right=input_values_right) + expected_results_categories = np.ones((2,2)) + result_cats = land_left.area_overlap(da_right, + categories_left=input_cat_left, + categories_right=input_cat_right) - assert(np.array_equal(expected_results_no_values, result_no_values)) - assert(np.array_equal(result_no_values.coords, [name_left, name_right])) - assert(np.array_equal(result_values[name_left].values, input_values_left)) - assert(np.array_equal(result_values[name_right].values, input_values_right)) + assert(np.array_equal(expected_results_categories, result_cats)) + assert(np.array_equal(result_cats.coords, [name_left, name_right])) + assert(np.array_equal(result_cats[name_left].values, input_cat_left)) + assert(np.array_equal(result_cats[name_right].values, input_cat_right)) # Test area overlap with explicit dimension names @@ -137,17 +137,18 @@ def test_category_match(): assert result_basic.name == name_left # Example with values on left map - result_values_left = land_left.category_match(da_right, values_left=1) - assert result_values_left.where(result_values_left==1).equals(da_left.where(da_left==1)) + result_values_left = land_left.category_match(da_right, categories_left=1) + assert result_values_left.where(result_values_left==1).equals( + da_left.where(da_left==1)) assert np.all(result_values_left.where(result_values_left!=1).isnull()) # Example with values on right map - result_values_left = land_left.category_match(da_right, values_right=1) + result_values_left = land_left.category_match(da_right, categories_right=1) # Example with multiple values on both maps result_multivar = land_left.category_match(da_right, - values_left=[0,1], - values_right=[1,2]) + categories_left=[0,1], + categories_right=[1,2]) non_nan_values = result_multivar.where(~np.isnan(result_multivar), drop=True) @@ -158,7 +159,7 @@ def test_category_match(): assert result_multivar.name == name_left # Example with non-matching values - result_non_matching = land_left.category_match(da_right, values_left=4) + result_non_matching = land_left.category_match(da_right, categories_left=4) assert np.all(result_non_matching.isnull()) def test_plot(): @@ -184,20 +185,234 @@ def test_dominant_class(): land = LandDataArray(da) # Test dominant class without coord name - result_no_coord = land.dominant_class() + result_no_coord = land.dominant_category() assert result_no_coord.equals(da.idxmax(dim="class")) # Test dominant class with coord name - result_coord = land.dominant_class(class_coord="class") + result_coord = land.dominant_category(category_dim="class") assert result_coord.equals(da.idxmax(dim="class")) # Test dominant class with non-matching coord name with pytest.raises(KeyError): - result_non_matching = land.dominant_class(class_coord="non_matching") + result_non_matching = land.dominant_category(category_dim="non_matching") # Test with return index set to True - result_return_index = land.dominant_class(return_index=True) + result_return_index = land.dominant_category(return_index=True) result_return_index_truth = xr.DataArray(np.argmax(data, axis=2), coords=coords_index) - assert result_return_index.equals(result_return_index_truth) \ No newline at end of file + assert result_return_index.equals(result_return_index_truth) + +def test_add_category_basic(): + categories = ["a", "b"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "categories": categories, + } + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(categories)) + da = xr.DataArray(data, coords=coords) + land = LandDataArray(da) + + result = land.add_category("c") + assert "c" in result["categories"].values + assert np.all(result.sel(categories="c") == 0) + +def test_add_category_with_value(): + categories = ["a", "b"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "categories": categories, + } + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(categories)) + da = xr.DataArray(data, coords=coords) + land = LandDataArray(da) + category_value = 1 + + result = land.add_category("c", category_value=category_value) + assert "c" in result["categories"].values + assert np.all(result.sel(categories="c") == category_value) + +def test_add_category_with_array_value(): + categories = ["a", "b"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "categories": categories, + } + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(categories)) + da = xr.DataArray(data, coords=coords) + land = LandDataArray(da) + + value_array_data = np.random.rand(len(coords["x"]), len(coords["y"])) + + value_array = xr.DataArray( + value_array_data, + dims=["x", "y"], + coords={"x": coords["x"], "y": coords["y"]}, + ) + + result = land.add_category("c", category_value=value_array) + assert "c" in result["categories"].values + xr.testing.assert_allclose( + result.sel(categories="c").drop_vars("categories"), + value_array) + +def test_add_category_with_year_array_value(): + categories = ["a", "b"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "Year": [2000, 2001], + "categories": categories, + } + data = np.random.rand( + len(coords["x"]), + len(coords["y"]), + len(coords["Year"]), + len(categories), + ) + da = xr.DataArray(data, coords=coords) + land = LandDataArray(da) + + value_array = xr.DataArray( + np.random.rand(len(coords["x"]), len(coords["y"]), len(coords["Year"])), + dims=["x", "y", "Year"], + coords={"x": coords["x"], "y": coords["y"], "Year": coords["Year"]}, + ) + + result = land.add_category("c", category_value=value_array) + assert "c" in result["categories"].values + xr.testing.assert_allclose( + result.sel(categories="c").drop_vars("categories"), + value_array, + ) + +def test_add_category_with_mask(): + categories = ["a", "b"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "categories": categories, + } + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(categories)) + da = xr.DataArray(data, coords=coords) + land = LandDataArray(da) + + mask_array = xr.DataArray( + np.array([[True, False], [False, True], [True, True]]), + dims=["x", "y"], + coords={"x": coords["x"], "y": coords["y"]}, + ) + + ex_result = xr.ones_like(mask_array).where(mask_array) + + result = land.add_category("c", category_value=1, mask=mask_array) + assert "c" in result["categories"].values + assert result.sel(categories="c").drop_vars("categories").equals(ex_result) + +def test_add_category_with_year_mask(): + categories = ["a", "b"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "Year": [2000, 2001], + "categories": categories, + } + data = np.random.rand( + len(coords["x"]), + len(coords["y"]), + len(coords["Year"]), + len(categories), + ) + da = xr.DataArray(data, coords=coords) + land = LandDataArray(da) + + mask_array = xr.DataArray( + np.array( + [ + [[True, False], [False, True]], + [[False, True], [True, False]], + [[True, True], [False, False]], + ] + ), + dims=["x", "y", "Year"], + coords={"x": coords["x"], "y": coords["y"], "Year": coords["Year"]}, + ) + + ex_result = xr.ones_like(mask_array, dtype=float).where(mask_array) + + result = land.add_category("c", category_value=1, mask=mask_array) + assert "c" in result["categories"].values + xr.testing.assert_allclose( + result.sel(categories="c").drop_vars("categories"), + ex_result, + ) + +def test_add_category_where_validation(): + classes = ["a", "b"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "class": classes, + } + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(classes)) + da = xr.DataArray(data, coords=coords) + land = LandDataArray(da) + + # Valid x/y boolean mask should pass + valid_where = xr.DataArray( + np.array([[True, False], [False, True], [True, True]]), + dims=["x", "y"], + coords={"x": coords["x"], "y": coords["y"]}, + ) + result = land.add_category("c", category_value=1, mask=valid_where) + assert "c" in result["class"].values + + # Missing spatial dimensions should fail + invalid_where_dims = xr.DataArray( + np.array([True, False]), + dims=["x"], + coords={"x": coords["x"][:2]}) + + with pytest.raises(ValueError): + land.add_category("d", category_value=1, mask=invalid_where_dims) + + # Misaligned x coordinate should fail + invalid_where_coords = xr.DataArray( + np.array([[True, False], [False, True], [True, True]]), + dims=["x", "y"], + coords={"x": [10, 11, 12], "y": coords["y"]}, + ) + with pytest.raises(ValueError): + land.add_category("e", category_value=1, mask=invalid_where_coords) + +def test_add_category_existing_categories_error(): + categories = ["a", "b"] + coords = { + "x": [0, 1], + "y": [0, 1], + "class": categories, + } + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(categories)) + da = xr.DataArray(data, coords=coords) + land = LandDataArray(da) + + with pytest.raises(ValueError): + land.add_category(["b"], category_value=1) + +def test_add_category_duplicate_input_categories_error(): + classes = ["a", "b"] + coords = { + "x": [0, 1], + "y": [0, 1], + "class": classes, + } + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(classes)) + da = xr.DataArray(data, coords=coords) + land = LandDataArray(da) + + + with pytest.raises(ValueError): + land.add_category(["c", "c"], category_value=1) \ No newline at end of file diff --git a/agrifoodpy/land/tests/test_model.py b/agrifoodpy/land/tests/test_model.py index c6d7c7a..27d5a94 100644 --- a/agrifoodpy/land/tests/test_model.py +++ b/agrifoodpy/land/tests/test_model.py @@ -47,4 +47,564 @@ def test_land_sequestration(): 9.52574127, 9.8201379 , 10.], coords={"Year":np.arange(11)}) - xr.testing.assert_allclose(result, truth) \ No newline at end of file + xr.testing.assert_allclose(result, truth) + +# ----------------------------- +# Tests for scale_land function +# ----------------------------- + +def test_scale_land_basic(): + # basic test with one category and one fraction + + from agrifoodpy.land.model import scale_land + + categories = ["Pasture", "Arable"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "categories": categories, + } + + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(categories)) + land = xr.DataArray(data, coords=coords) + + fraction = np.random.rand() + + result = scale_land(land, origin="Pasture", fraction=fraction) + + truth = land.copy(deep=True) + truth.loc[{"categories": "Pasture"}] *= fraction + + xr.testing.assert_allclose(result, truth) + +def test_scale_land_multiple_categories(): + # test with multiple categories and one fraction + + from agrifoodpy.land.model import scale_land + + categories = ["Pasture", "Arable", "Forest"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "categories": categories, + } + + fraction = np.random.rand() + + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(categories)) + land = xr.DataArray(data, coords=coords) + + result = scale_land(land, origin=["Pasture", "Arable"], fraction=fraction) + + truth = land.copy(deep=True) + truth.loc[{"categories": ["Pasture", "Arable"]}] *= fraction + + xr.testing.assert_allclose(result, truth) + +def test_scale_land_multiple_categories_multiple_fractions(): + # test with multiple categories and multiple fractions + + from agrifoodpy.land.model import scale_land + + categories = ["Pasture", "Arable", "Forest"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "categories": categories, + } + + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(categories)) + land = xr.DataArray(data, coords=coords) + + fraction = xr.DataArray( + np.random.rand(2), + coords={"categories": ["Pasture", "Arable"]}) + + result = scale_land(land, origin=["Pasture", "Arable"], fraction=fraction) + + truth = land.copy(deep=True) + truth.loc[{"categories": "Pasture"}] *= fraction.loc[{"categories": "Pasture"}] + truth.loc[{"categories": "Arable"}] *= fraction.loc[{"categories": "Arable"}] + + xr.testing.assert_allclose(result, truth) + +def test_scale_land_single_target_category(): + # test with target category + + from agrifoodpy.land.model import scale_land + + categories = ["Pasture", "Arable", "Forest"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "categories": categories, + } + + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(categories)) + land = xr.DataArray(data, coords=coords) + + fraction = np.random.rand() + + result = scale_land( + land, + origin="Pasture", + fraction=fraction, + keep_land_constant=True, + target_category="Forest") + + truth = land.copy(deep=True) + truth.loc[{"categories": "Pasture"}] *= fraction + truth.loc[{"categories": "Forest"}] += \ + land.loc[{"categories": "Pasture"}] * (1-fraction) + + xr.testing.assert_allclose(result, truth) + +def test_scale_land_multiple_target_categories(): + # Test with multiple target categories + + from agrifoodpy.land.model import scale_land + + categories = ["Pasture", "Arable", "Forest", "Urban"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "categories": categories, + } + + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(categories)) + land = xr.DataArray(data, coords=coords) + + fraction = np.random.rand() + + result = scale_land( + land, + origin="Pasture", + fraction=fraction, + keep_land_constant=True, + target_category=["Forest", "Urban"]) + + truth = land.copy(deep=True) + truth.loc[{"categories": "Pasture"}] *= fraction + truth.loc[{"categories": ["Forest", "Urban"]}] += \ + land.loc[{"categories": "Pasture"}] * (1-fraction) / 2 + + xr.testing.assert_allclose(result, truth) + +def test_scale_land_multiple_target_categories_with_distribution(): + # Test with multiple target categories and target distribution + + from agrifoodpy.land.model import scale_land + + categories = ["Pasture", "Arable", "Forest", "Urban"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "categories": categories, + } + + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(categories)) + land = xr.DataArray(data, coords=coords) + + fraction = np.random.rand() + target_distr_arr = np.random.rand(2) + target_distr_arr /= target_distr_arr.sum() + target_distribution = xr.DataArray( + target_distr_arr, + coords={"categories": ["Forest", "Urban"]}) + + result = scale_land( + land, + origin="Pasture", + fraction=fraction, + keep_land_constant=True, + target_category=["Forest", "Urban"], + target_distribution=target_distribution) + + truth = land.copy(deep=True) + truth.loc[{"categories": "Pasture"}] *= fraction + truth.loc[{"categories": "Forest"}] += land.loc[{"categories": "Pasture"}] \ + * (1-fraction) * target_distribution.loc[{"categories": "Forest"}] + truth.loc[{"categories": "Urban"}] += land.loc[{"categories": "Pasture"}] \ + * (1-fraction) * target_distribution.loc[{"categories": "Urban"}] + + xr.testing.assert_allclose(result, truth) + +def test_scale_land_time_dependent_scale(): + # Test with time dependent scale factor + + from agrifoodpy.land.model import scale_land + + categories = ["Pasture", "Arable", "Forest"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "categories": categories, + "Year": [2000, 2010] + } + + data = np.random.rand( + len(coords["x"]), + len(coords["y"]), + len(categories), + len(coords["Year"]) + ) + land = xr.DataArray(data, coords=coords) + + fraction = xr.DataArray( + np.random.rand(len(coords["Year"])), + coords={"Year": coords["Year"]}) + + result = scale_land( + land, + origin="Pasture", + fraction=fraction, + keep_land_constant=True, + target_category=["Arable", "Forest"], + target_distribution=[0.75, 0.25]) + + truth = land.copy(deep=True) + for t in coords["Year"]: + truth.loc[{"categories": "Pasture", "Year": t}] \ + *= fraction.loc[{"Year": t}] + + truth.loc[{"categories": "Arable", "Year": t}] \ + += land.loc[{"categories": "Pasture", "Year": t}] \ + * (1-fraction.loc[{"Year": t}]) * 0.75 + + truth.loc[{"categories": "Forest", "Year": t}] \ + += land.loc[{"categories": "Pasture", "Year": t}] \ + * (1-fraction.loc[{"Year": t}]) * 0.25 + + xr.testing.assert_allclose(result, truth) + +def test_scale_land_custom_category_dim(): + # Test with custom category dimension + + from agrifoodpy.land.model import scale_land + + categories = ["Pasture", "Arable", "Forest"] + coords = { + "x": [0, 1, 2], + "y": [0, 1], + "land_category": categories, + } + + data = np.random.rand(len(coords["x"]), len(coords["y"]), len(categories)) + land = xr.DataArray(data, coords=coords) + + fraction = np.random.rand() + + result = scale_land( + land, + origin="Pasture", + fraction=fraction, + keep_land_constant=True, + target_category=["Arable", "Forest"], + target_distribution=[0.75, 0.25], + category_dim="land_category") + + truth = land.copy(deep=True) + truth.loc[{"land_category": "Pasture"}] \ + *= fraction + + truth.loc[{"land_category": "Arable"}] \ + += land.loc[{"land_category": "Pasture"}] * (1-fraction) * 0.75 + + truth.loc[{"land_category": "Forest"}] \ + += land.loc[{"land_category": "Pasture"}] * (1-fraction) * 0.25 + + xr.testing.assert_allclose(result, truth) + +# -------------------------------- +# Tests for land_scaling_from_food +# -------------------------------- + +def test_land_scaling_from_food_basic(): + # Basic test with unchanged food balance sheet and one category + + from agrifoodpy.land.model import land_scaling_from_food + + items = ["Beef", "Apples"] + food_coords = {"Item": items} + food_data = np.random.rand(len(food_coords["Item"])) + + fbs_obs = xr.Dataset( + data_vars={"production": (["Item"], food_data)}, + coords=food_coords + ) + + fbs_ref = fbs_obs.copy(deep=True) + + categories = ["Pasture", "Arable", "Forest"] + land_coords = { + "x": [0, 1, 2], + "y": [0, 1], + "category": categories, + } + + land_data = np.random.rand( + len(land_coords["x"]), + len(land_coords["y"]), + len(land_coords["category"])) + + land = xr.DataArray(land_data, coords=land_coords) + + result = land_scaling_from_food( + land, + fbs_obs, + fbs_ref, + element="production", + category="Pasture") + + xr.testing.assert_allclose(result, land) + +def test_land_scaling_from_food_basic_change(): + # Basic test with changed food balance sheet and one category + + from agrifoodpy.land.model import land_scaling_from_food + + items = ["Beef", "Apples"] + food_coords = {"Item": items} + food_data = np.random.rand(len(food_coords["Item"])) + + fbs_ref = xr.Dataset( + data_vars={"production": (["Item"], food_data)}, + coords=food_coords + ) + + fbs_obs = fbs_ref.copy(deep=True) + fbs_obs["production"] *= 2 + + categories = ["Pasture", "Arable", "Forest"] + land_coords = { + "x": [0, 1, 2], + "y": [0, 1], + "category": categories, + } + + land_data = np.random.rand( + len(land_coords["x"]), + len(land_coords["y"]), + len(land_coords["category"])) + + land = xr.DataArray(land_data, coords=land_coords) + + result = land_scaling_from_food( + land, + fbs_obs, + fbs_ref, + element="production", + category="Pasture") + + truth = land.copy(deep=True) + truth.loc[{"category":"Pasture"}] *= 2 + + xr.testing.assert_allclose(result, truth) + +def test_land_scaling_from_food_with_items(): + # Test with specified items + + from agrifoodpy.land.model import land_scaling_from_food + + items = ["Beef", "Apples"] + food_coords = {"Item": items} + food_data = np.random.rand(len(food_coords["Item"])) + + fbs_ref = xr.Dataset( + data_vars={"production": (["Item"], food_data)}, + coords=food_coords + ) + + fbs_obs = fbs_ref.copy(deep=True) + fbs_obs["production"].loc[{"Item":"Beef"}] *= 2 + + categories = ["Pasture", "Arable", "Forest"] + land_coords = { + "x": [0, 1, 2], + "y": [0, 1], + "category": categories, + } + + land_data = np.random.rand( + len(land_coords["x"]), + len(land_coords["y"]), + len(land_coords["category"])) + + land = xr.DataArray(land_data, coords=land_coords) + + result = land_scaling_from_food( + land, + fbs_obs, + fbs_ref, + element="production", + category="Pasture", + items=["Beef"]) + + truth = land.copy(deep=True) + truth.loc[{"category":"Pasture"}] *= 2 + + xr.testing.assert_allclose(result, truth) + + +def test_land_scaling_from_food_with_target_categories(): + # Test with constant land and target categories + + from agrifoodpy.land.model import land_scaling_from_food + + items = ["Beef", "Apples"] + food_coords = {"Item": items} + food_data = np.random.rand(len(food_coords["Item"])) + + fbs_ref = xr.Dataset( + data_vars={"production": (["Item"], food_data)}, + coords=food_coords + ) + + food_scale = np.random.rand() * 0.5 + 1 + fbs_obs = fbs_ref.copy(deep=True) + fbs_obs["production"].loc[{"Item":"Beef"}] *= food_scale + + + categories = ["Pasture", "Arable", "Forest"] + land_coords = { + "x": [0, 1, 2], + "y": [0, 1], + "category": categories, + } + + land_data = np.random.rand( + len(land_coords["x"]), + len(land_coords["y"]), + len(land_coords["category"])) + + land = xr.DataArray(land_data, coords=land_coords) + + result = land_scaling_from_food( + land, + fbs_obs, + fbs_ref, + element="production", + category="Pasture", + keep_land_constant=True, + target_category="Forest", + items=["Beef"]) + + truth = land.copy(deep=True) + truth.loc[{"category":"Pasture"}] *= food_scale + truth.loc[{"category":"Forest"}] -= land.loc[{"category":"Pasture"}] \ + * (food_scale - 1) + + xr.testing.assert_allclose(result, truth) + + +def test_land_scaling_from_food_with_target_categories_and_distribution(): + # Test with constant land, target categories and target distribution + + from agrifoodpy.land.model import land_scaling_from_food + + items = ["Beef", "Apples"] + food_coords = {"Item": items} + food_data = np.random.rand(len(food_coords["Item"])) + + fbs_ref = xr.Dataset( + data_vars={"production": (["Item"], food_data)}, + coords=food_coords + ) + + food_scale = np.random.rand() * 0.5 + 1 + fbs_obs = fbs_ref.copy(deep=True) + fbs_obs["production"].loc[{"Item":"Beef"}] *= food_scale + + + categories = ["Pasture", "Arable", "Forest"] + land_coords = { + "x": [0, 1, 2], + "y": [0, 1], + "category": categories, + } + + land_data = np.random.rand( + len(land_coords["x"]), + len(land_coords["y"]), + len(land_coords["category"])) + + land = xr.DataArray(land_data, coords=land_coords) + + result = land_scaling_from_food( + land, + fbs_obs, + fbs_ref, + element="production", + category="Pasture", + keep_land_constant=True, + target_category=["Forest", "Arable"], + target_distribution=[0.75, 0.25], + items=["Beef"]) + + truth = land.copy(deep=True) + truth.loc[{"category":"Pasture"}] *= food_scale + truth.loc[{"category":"Forest"}] -= land.loc[{"category":"Pasture"}] \ + * (food_scale - 1) * 0.75 + truth.loc[{"category":"Arable"}] -= land.loc[{"category":"Pasture"}] \ + * (food_scale - 1) * 0.25 + + xr.testing.assert_allclose(result, truth) + +def test_land_scaling_from_food_with_time_dependent_fbs(): + + from agrifoodpy.land.model import land_scaling_from_food + + items = ["Beef", "Apples"] + years = [2000, 2010] + food_coords = {"Item": items, "Year": years} + food_data = np.random.rand( + len(food_coords["Item"]), + len(food_coords["Year"]) + ) + + fbs_ref = xr.Dataset( + data_vars={"production": (["Item", "Year"], food_data)}, + coords=food_coords + ) + + food_scale = xr.DataArray( + np.random.rand(len(years)), + coords={"Year": years} + ) * 0.5 + 1 + + fbs_obs = fbs_ref.copy(deep=True) + fbs_obs["production"].loc[{"Item":"Beef"}] *= food_scale + + categories = ["Pasture", "Arable", "Forest"] + land_coords = { + "x": [0, 1, 2], + "y": [0, 1], + "category": categories, + "Year": years + } + + land_data = np.random.rand( + len(land_coords["x"]), + len(land_coords["y"]), + len(land_coords["category"]), + len(land_coords["Year"]) + ) + + land = xr.DataArray(land_data, coords=land_coords) + + result = land_scaling_from_food( + land, + fbs_obs, + fbs_ref, + element="production", + category="Pasture", + keep_land_constant=True, + target_category="Forest", + items=["Beef"]) + + truth = land.copy(deep=True) + truth.loc[{"category":"Pasture"}] *= food_scale + truth.loc[{"category":"Forest"}] -= land.loc[{"category":"Pasture"}] \ + * (food_scale - 1) + + xr.testing.assert_allclose(result, truth) diff --git a/agrifoodpy/pipeline/pipeline.py b/agrifoodpy/pipeline/pipeline.py index 10c0e41..8288a85 100644 --- a/agrifoodpy/pipeline/pipeline.py +++ b/agrifoodpy/pipeline/pipeline.py @@ -16,6 +16,18 @@ class Pipeline(): '''Class for constructing and running pipelines of functions with individual sets of parameters.''' + _SCALE_YAML_FUNCTIONS = { + "linear": "agrifoodpy.utils.scaling.linear_scale", + "step": "agrifoodpy.utils.scaling.step_scale", + "pulse": "agrifoodpy.utils.scaling.pulse_scale", + "logistic": "agrifoodpy.utils.scaling.logistic_scale", + "smoothstep": "agrifoodpy.utils.scaling.smoothstep_scale", + + "piecewise_linear": "agrifoodpy.utils.scaling.piecewise_linear_scale", + "piecewise_constant": "agrifoodpy.utils.scaling.piecewise_constant_scale", + "piecewise_smoothstep": "agrifoodpy.utils.scaling.piecewise_smoothstep_scale", + } + def __init__(self, datablock=None): self.nodes = [] self.params = [] @@ -34,10 +46,13 @@ def _load_function(path): @staticmethod def _is_supported_yaml_function(path): - """Return True for dotted numpy/xarray function paths.""" + """Return True for dotted numpy/xarray and scale function paths.""" if not isinstance(path, str) or "." not in path: return False + if path in Pipeline._SCALE_YAML_FUNCTIONS.values(): + return True + module_path, _ = path.rsplit(".", 1) return ( module_path == "numpy" @@ -46,6 +61,17 @@ def _is_supported_yaml_function(path): or module_path.startswith("xarray.") ) + @staticmethod + def _resolve_yaml_function_path(namespace, suffix): + """Resolve YAML constructor suffixes to supported function paths.""" + if namespace == "scale": + return Pipeline._SCALE_YAML_FUNCTIONS.get(suffix) + + if suffix: + return f"{namespace}.{suffix}" + + return namespace + @classmethod def read(cls, filename): """Read a pipeline configuration from a YAML file @@ -65,14 +91,23 @@ def dynamic_call_constructor(package_name): """Build a multi-constructor for supported package functions.""" def constructor(loader, suffix, node): - func_path = f"{package_name}.{suffix}" if suffix else package_name + func_path = cls._resolve_yaml_function_path(package_name, suffix) + tag_path = f"{package_name}.{suffix}" if suffix else package_name + + if func_path is None: + raise yaml.constructor.ConstructorError( + None, + None, + f"Unsupported YAML function tag '!{tag_path}'.", + node.start_mark, + ) # Check if the function path is supported if not cls._is_supported_yaml_function(func_path): raise yaml.constructor.ConstructorError( None, None, - f"Unsupported YAML function tag '!{func_path}'.", + f"Unsupported YAML function tag '!{tag_path}'.", node.start_mark, ) @@ -90,7 +125,7 @@ def constructor(loader, suffix, node): raise yaml.constructor.ConstructorError( None, None, - f"Unsupported YAML node type for '!{func_path}'.", + f"Unsupported YAML node type for '!{tag_path}'.", node.start_mark, ) @@ -106,6 +141,11 @@ def constructor(loader, suffix, node): dynamic_call_constructor("xarray"), Loader=yaml.FullLoader, ) + yaml.add_multi_constructor( + "!scale.", + dynamic_call_constructor("scale"), + Loader=yaml.FullLoader, + ) with open(filename, "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) @@ -363,9 +403,10 @@ def pipeline_node(input_keys=None): be interpreted as datablock lookup keys. The corresponding objects will be extracted from the datablock and passed to the function. Unregistered keyword arguments will be passed directly to the function. The decorated - function takes a "return_key" kwarg to specify the key under which the - function output will be stored in the datablock. If not provided, the - function name will be used as the return key. + function takes "return_key" or "return_keys" kwargs to specify the keys + under which the function outputs will be stored in the datablock. + If not provided, the function name will be used as the prefix for the + return keys. Parameters ---------- @@ -386,7 +427,52 @@ def pipeline_node(input_keys=None): input_keys = [] def pipeline_decorator(func): - reserved = {"datablock", "return_key"} + + def normalize_return_keys(return_keys): + if isinstance(return_keys, (str, tuple)): + return [return_keys] + + if isinstance(return_keys, list): + return return_keys + + raise TypeError( + "return_key/return_keys must be a str, tuple, or list." + ) + + def validate_return_keys(return_keys): + if len(return_keys) == 0: + raise ValueError("return_keys cannot be empty.") + + for i, key in enumerate(return_keys): + if isinstance(key, str): + continue + + if isinstance(key, tuple): + if len(key) == 0: + raise ValueError( + f"return_keys[{i}] tuple path cannot be empty." + ) + + for j, key_part in enumerate(key): + if not isinstance(key_part, str): + raise TypeError( + "return_keys[{}][{}] path component must be " + "str, got {}.".format( + i, + j, + type(key_part).__name__, + ) + ) + continue + + raise TypeError( + "return_keys[{}] must be str or tuple, got {}.".format( + i, + type(key).__name__, + ) + ) + + reserved = {"datablock", "return_key", "return_keys"} if reserved & set(signature(func).parameters): raise ValueError(f"Function {func.__name__} has reserved parameter" f" names {reserved & set(signature(func).parameters)}." @@ -404,7 +490,29 @@ def wrapper(*args, **kwargs): # Pop wrapper-specific kwargs datablock = kwargs.pop("datablock", None) - return_key = kwargs.pop("return_key", func.__name__) + + has_return_key = ( + "return_key" in kwargs and kwargs["return_key"] is not None + ) + has_return_keys = ( + "return_keys" in kwargs and kwargs["return_keys"] is not None + ) + + if has_return_key and has_return_keys: + raise ValueError( + "Ambiguous return mapping: both 'return_key' and " + "'return_keys' were provided." + ) + + return_keys = kwargs.pop("return_keys", None) + return_key = kwargs.pop("return_key", None) + + if return_keys is None and return_key is not None: + return_keys = return_key + + if return_keys is not None: + return_keys = normalize_return_keys(return_keys) + validate_return_keys(return_keys) # Bind positional and keyword args to their parameter names func_sig = signature(func) @@ -428,8 +536,35 @@ def wrapper(*args, **kwargs): bound.arguments[key]) result = func(*bound.args, **bound.kwargs) - set_dict(datablock, return_key, result) - + if return_keys is None: + if isinstance(result, tuple): + return_keys = [ + f"{func.__name__}_{i}" for i in range(len(result)) + ] + else: + return_keys = [func.__name__] + + if isinstance(result, tuple): + if len(return_keys) != len(result): + raise ValueError( + f"Function '{func.__name__}' returned tuple of " + f"length {len(result)}, but {len(return_keys)} " + "return keys were provided." + ) + + for rk, rs in zip(return_keys, result): + set_dict(datablock, rk, rs) + + else: + if len(return_keys) != 1: + raise ValueError( + f"Function '{func.__name__}' returned non-tuple " + f"{type(result).__name__}, but {len(return_keys)} " + "return keys were provided." + ) + + set_dict(datablock, return_keys[0], result) + return datablock return wrapper return pipeline_decorator \ No newline at end of file diff --git a/agrifoodpy/pipeline/tests/data/test_config_scale_linear.yaml b/agrifoodpy/pipeline/tests/data/test_config_scale_linear.yaml new file mode 100644 index 0000000..707b656 --- /dev/null +++ b/agrifoodpy/pipeline/tests/data/test_config_scale_linear.yaml @@ -0,0 +1,6 @@ +nodes: + - function: agrifoodpy.utils.nodes.write_to_datablock + name: Scale Linear + params: + key: "test_scale_linear" + value: !scale.linear [2020, 2022, 2024, 2026, 1.0, 3.0] diff --git a/agrifoodpy/pipeline/tests/data/test_config_scale_linear_kwargs.yaml b/agrifoodpy/pipeline/tests/data/test_config_scale_linear_kwargs.yaml new file mode 100644 index 0000000..75047e7 --- /dev/null +++ b/agrifoodpy/pipeline/tests/data/test_config_scale_linear_kwargs.yaml @@ -0,0 +1,6 @@ +nodes: + - function: agrifoodpy.utils.nodes.write_to_datablock + name: Scale Linear + params: + key: "test_scale_linear" + value: !scale.linear {y0: 2020, y1: 2022, y2: 2024, y3: 2026, c_init: 1.0, c_end: 3.0} diff --git a/agrifoodpy/pipeline/tests/data/test_config_scale_logistic.yaml b/agrifoodpy/pipeline/tests/data/test_config_scale_logistic.yaml new file mode 100644 index 0000000..147240a --- /dev/null +++ b/agrifoodpy/pipeline/tests/data/test_config_scale_logistic.yaml @@ -0,0 +1,6 @@ +nodes: + - function: agrifoodpy.utils.nodes.write_to_datablock + name: Scale Logistic + params: + key: "test_scale_logistic" + value: !scale.logistic [2020, 2022, 2024, 2026, 1.0, 3.0] diff --git a/agrifoodpy/pipeline/tests/data/test_config_scale_logistic_kwargs.yaml b/agrifoodpy/pipeline/tests/data/test_config_scale_logistic_kwargs.yaml new file mode 100644 index 0000000..cfa3e0d --- /dev/null +++ b/agrifoodpy/pipeline/tests/data/test_config_scale_logistic_kwargs.yaml @@ -0,0 +1,6 @@ +nodes: + - function: agrifoodpy.utils.nodes.write_to_datablock + name: Scale Logistic + params: + key: "test_scale_logistic" + value: !scale.logistic {y0: 2020, y1: 2022, y2: 2024, y3: 2026, c_init: 1.0, c_end: 3.0} diff --git a/agrifoodpy/pipeline/tests/data/test_config_scale_unsupported_function.yaml b/agrifoodpy/pipeline/tests/data/test_config_scale_unsupported_function.yaml new file mode 100644 index 0000000..298f42c --- /dev/null +++ b/agrifoodpy/pipeline/tests/data/test_config_scale_unsupported_function.yaml @@ -0,0 +1,6 @@ +nodes: + - function: agrifoodpy.utils.nodes.write_to_datablock + name: Scale Unsupported + params: + key: "test_scale_unsupported" + value: !scale.cubic [2020, 2022, 2024, 2026, 1.0, 3.0] diff --git a/agrifoodpy/pipeline/tests/test_pipeline.py b/agrifoodpy/pipeline/tests/test_pipeline.py index 6f05ad2..461489c 100644 --- a/agrifoodpy/pipeline/tests/test_pipeline.py +++ b/agrifoodpy/pipeline/tests/test_pipeline.py @@ -1,4 +1,5 @@ from agrifoodpy.pipeline import Pipeline, standalone +from agrifoodpy.utils.scaling import linear_scale, logistic_scale import numpy as np import xarray as xr import pytest @@ -364,6 +365,119 @@ def unknown_input_node(right_key): pass +def test_pipeline_node_multi_output_and_return_keys(): + + from agrifoodpy.pipeline.pipeline import Pipeline, pipeline_node + + @pipeline_node([]) + def split_values(): + return 10, 20 + + @pipeline_node([]) + def scalar_value(): + return 7 + + # Explicit string keys for tuple output + pipeline_string_keys = Pipeline() + pipeline_string_keys.add_node( + split_values, + params={"return_keys": ["out1", "out2"]}, + ) + pipeline_string_keys.run() + assert pipeline_string_keys.datablock["out1"] == 10 + assert pipeline_string_keys.datablock["out2"] == 20 + + # Explicit tuple paths for tuple output + pipeline_tuple_paths = Pipeline() + pipeline_tuple_paths.add_node( + split_values, + params={"return_keys": [("nested", "a"), ("nested", "b")]}, + ) + pipeline_tuple_paths.run() + assert pipeline_tuple_paths.datablock["nested"]["a"] == 10 + assert pipeline_tuple_paths.datablock["nested"]["b"] == 20 + + # Mixed str/tuple paths for tuple output + pipeline_mixed_paths = Pipeline() + pipeline_mixed_paths.add_node( + split_values, + params={"return_keys": ["out", ("nested", "out")]}, + ) + pipeline_mixed_paths.run() + assert pipeline_mixed_paths.datablock["out"] == 10 + assert pipeline_mixed_paths.datablock["nested"]["out"] == 20 + + # Auto-generated keys for tuple outputs when no keys are passed + pipeline_auto_keys = Pipeline() + pipeline_auto_keys.add_node(split_values) + pipeline_auto_keys.run() + assert pipeline_auto_keys.datablock["split_values_0"] == 10 + assert pipeline_auto_keys.datablock["split_values_1"] == 20 + + # Single tuple path through return_key for scalar output + pipeline_return_key_tuple = Pipeline() + pipeline_return_key_tuple.add_node( + scalar_value, + params={"return_key": ("nested", "scalar")}, + ) + pipeline_return_key_tuple.run() + assert pipeline_return_key_tuple.datablock["nested"]["scalar"] == 7 + + # Mismatch: tuple output with fewer keys + pipeline_tuple_mismatch = Pipeline() + pipeline_tuple_mismatch.add_node( + split_values, + params={"return_keys": ["only_one"]}, + ) + with pytest.raises(ValueError, match="returned tuple of length"): + pipeline_tuple_mismatch.run() + + # Mismatch: scalar output with multiple keys + pipeline_scalar_mismatch = Pipeline() + pipeline_scalar_mismatch.add_node( + scalar_value, + params={"return_keys": ["a", "b"]}, + ) + with pytest.raises(ValueError, match="returned non-tuple"): + pipeline_scalar_mismatch.run() + + # Empty return_keys is invalid + pipeline_empty_keys = Pipeline() + pipeline_empty_keys.add_node( + scalar_value, + params={"return_keys": []}, + ) + with pytest.raises(ValueError, match="cannot be empty"): + pipeline_empty_keys.run() + + # Invalid key type is rejected + pipeline_bad_key_type = Pipeline() + pipeline_bad_key_type.add_node( + scalar_value, + params={"return_keys": [123]}, + ) + with pytest.raises(TypeError, match="must be str or tuple"): + pipeline_bad_key_type.run() + + # Invalid tuple path component is rejected + pipeline_bad_tuple_component = Pipeline() + pipeline_bad_tuple_component.add_node( + scalar_value, + params={"return_keys": [("nested", 123)]}, + ) + with pytest.raises(TypeError, match="path component must be str"): + pipeline_bad_tuple_component.run() + + # Providing both return_key and return_keys is ambiguous + pipeline_ambiguous_keys = Pipeline() + pipeline_ambiguous_keys.add_node( + scalar_value, + params={"return_key": "a", "return_keys": ["b"]}, + ) + with pytest.raises(ValueError, match="both 'return_key' and 'return_keys'"): + pipeline_ambiguous_keys.run() + + # Test reading YAML config with numpy array parameters and values def test_read_yaml_numpy_array(): @@ -408,10 +522,62 @@ def test_read_yaml_xarray_dataarray_kwargs(): xr.testing.assert_equal(pipeline.params[0]['value'], expected_array) xr.testing.assert_equal(pipeline.datablock["test_value"], expected_array) +def test_read_yaml_scale_linear(): + script_dir = os.path.dirname(__file__) + config_path = os.path.join(script_dir, "data/test_config_scale_linear.yaml") + + pipeline = Pipeline.read(str(config_path)) + pipeline.run() + + expected_array = linear_scale(2020, 2022, 2024, 2026, 1.0, 3.0) + xr.testing.assert_equal(pipeline.params[0]['value'], expected_array) + xr.testing.assert_equal(pipeline.datablock["test_scale_linear"], expected_array) + +def test_read_yaml_scale_linear_kwargs(): + script_dir = os.path.dirname(__file__) + config_path = os.path.join(script_dir, "data/test_config_scale_linear_kwargs.yaml") + + pipeline = Pipeline.read(str(config_path)) + pipeline.run() + + expected_array = linear_scale(2020, 2022, 2024, 2026, 1.0, 3.0) + xr.testing.assert_equal(pipeline.params[0]['value'], expected_array) + xr.testing.assert_equal(pipeline.datablock["test_scale_linear"], expected_array) + +def test_read_yaml_scale_logistic(): + script_dir = os.path.dirname(__file__) + config_path = os.path.join(script_dir, "data/test_config_scale_logistic.yaml") + + pipeline = Pipeline.read(str(config_path)) + pipeline.run() + + expected_array = logistic_scale(2020, 2022, 2024, 2026, 1.0, 3.0) + xr.testing.assert_equal(pipeline.params[0]['value'], expected_array) + xr.testing.assert_equal(pipeline.datablock["test_scale_logistic"], expected_array) + +def test_read_yaml_scale_logistic_kwargs(): + script_dir = os.path.dirname(__file__) + config_path = os.path.join(script_dir, "data/test_config_scale_logistic_kwargs.yaml") + + pipeline = Pipeline.read(str(config_path)) + pipeline.run() + + expected_array = logistic_scale(2020, 2022, 2024, 2026, 1.0, 3.0) + xr.testing.assert_equal(pipeline.params[0]['value'], expected_array) + xr.testing.assert_equal(pipeline.datablock["test_scale_logistic"], expected_array) + def test_read_yaml_unsupported_function(): from yaml.constructor import ConstructorError script_dir = os.path.dirname(__file__) config_path = os.path.join(script_dir, "data/test_config_unsupported_function.yaml") + with pytest.raises(ConstructorError): + pipeline = Pipeline.read(str(config_path)) + +def test_read_yaml_unsupported_scale_function(): + from yaml.constructor import ConstructorError + script_dir = os.path.dirname(__file__) + config_path = os.path.join(script_dir, "data/test_config_scale_unsupported_function.yaml") + with pytest.raises(ConstructorError): pipeline = Pipeline.read(str(config_path)) \ No newline at end of file diff --git a/agrifoodpy/tests/test_base_class.py b/agrifoodpy/tests/test_base_class.py index f2a062a..fda7674 100644 --- a/agrifoodpy/tests/test_base_class.py +++ b/agrifoodpy/tests/test_base_class.py @@ -2,44 +2,80 @@ import xarray as xr from ..array_accessor import XarrayAccessorBase -def test_add_years(): +# Tests for add_years +def test_add_years_single(): + # Test adding single year years = np.arange(2010, 2013) data = np.random.rand(3, 2, 2) - new_years = [2013, 2014, 2015] - expected_years = np.concatenate([years, new_years]) + new_year = 2014 ds = xr.Dataset({"data": (("Year", "X", "Y"), data)}, coords={"Year": years, "X": [0, 1], "Y": [0, 1]}) fbs = XarrayAccessorBase(ds) - # Test adding single year - - result_single = fbs.add_years(new_years[0]) - expected_years_single = np.concatenate([years, [new_years[0]]]) + result_single = fbs.add_years(new_year) + expected_years_single = np.concatenate([years, [new_year]]) assert np.array_equal(result_single["Year"], expected_years_single) - assert np.isnan(result_single["data"].loc[{"Year":new_years[0]}].to_numpy() - ).all() + assert np.isnan( + result_single["data"].loc[{"Year":new_year}].to_numpy()).all() + + +def test_add_years_multiple(): + # Test adding multiple years + + years = np.arange(2010, 2013) + data = np.random.rand(3, 2, 2) + + ds = xr.Dataset({"data": (("Year", "X", "Y"), data)}, + coords={"Year": years, "X": [0, 1], "Y": [0, 1]}) + fbs = XarrayAccessorBase(ds) + + new_years = [2013, 2014, 2015] + expected_years = np.unique(np.concatenate([years, new_years])) - # Test adding years with "empty" projection - result_empty = fbs.add_years(new_years, projection="empty") + result_empty = fbs.add_years(new_years) assert np.array_equal(result_empty["Year"], expected_years) assert np.isnan(result_empty["data"].loc[{"Year":new_years}].to_numpy() ).all() + +def test_add_years_constant_projection(): # Test adding years with "constant" projection + + years = np.arange(2010, 2013) + data = np.random.rand(3, 2, 2) + + ds = xr.Dataset({"data": (("Year", "X", "Y"), data)}, + coords={"Year": years, "X": [0, 1], "Y": [0, 1]}) + fbs = XarrayAccessorBase(ds) + + new_years = [2013, 2014, 2015] + expected_years = np.unique(np.concatenate([years, new_years])) + result_constant = fbs.add_years(new_years, projection="constant") last_year_data = ds["data"].isel(Year=-1).values assert np.array_equal(result_constant["Year"], expected_years) for year in new_years: - assert np.array_equal(result_constant["data"].loc[{"Year":year}].values, - last_year_data) + assert np.array_equal( + result_constant["data"].loc[{"Year":year}].values, last_year_data) +def test_add_years_with_projection_array(): # Test adding years with specific projection + + years = np.arange(2010, 2013) + data = np.random.rand(3, 2, 2) + new_years = [2013, 2014, 2015] + expected_years = np.unique(np.concatenate([years, new_years])) + + ds = xr.Dataset({"data": (("Year", "X", "Y"), data)}, + coords={"Year": years, "X": [0, 1], "Y": [0, 1]}) + fbs = XarrayAccessorBase(ds) + proj = [0.5, 0.6, 0.7] # Scaling factors result_projection = fbs.add_years(new_years, projection=proj) last_year_data = result_projection["data"].loc[dict(Year=years[-1])].values @@ -49,12 +85,37 @@ def test_add_years(): assert np.allclose(result_projection["data"].loc[{"Year":new_years}].values, expected_data) +def test_add_years_duplicates(): # Test for duplicate years + years = np.arange(2010, 2013) + data = np.random.rand(3, 2, 2) + new_years = [2013, 2014, 2015] + expected_years = np.unique(np.concatenate([years, new_years])) + + ds = xr.Dataset({"data": (("Year", "X", "Y"), data)}, + coords={"Year": years, "X": [0, 1], "Y": [0, 1]}) + fbs = XarrayAccessorBase(ds) + new_years_duplicate = [2013, 2013, 2014, 2015] result_duplicate = fbs.add_years(new_years_duplicate) assert np.array_equal(result_duplicate["Year"], expected_years) +def test_add_years_no_year_dim(): + # Test for adding years to an array without "Year" dimension + + ds_no_year = xr.Dataset({"data": (("X", "Y"), np.random.rand(2, 2))}, + coords={"X": [0, 1], "Y": [0, 1]}) + + fbs_no_year = XarrayAccessorBase(ds_no_year) + + new_years = [2013, 2014, 2015] + result_no_year = fbs_no_year.add_years(new_years) + + assert np.array_equal(result_no_year["Year"], new_years) + + +# Tests for add_items def test_add_items(): items = ["Beef", "Apples", "Poultry"] diff --git a/agrifoodpy/utils/nodes.py b/agrifoodpy/utils/nodes.py index 5e3cefb..76381cb 100644 --- a/agrifoodpy/utils/nodes.py +++ b/agrifoodpy/utils/nodes.py @@ -7,6 +7,7 @@ from ..pipeline import standalone from ..utils.dict_utils import get_dict, set_dict +from ..food.food import FoodBalanceSheet # noqa: F401 @standalone(["dataset"], ["dataset"]) @@ -75,6 +76,7 @@ def add_items( def add_years( dataset, years, + pivot_year=None, projection='empty', datablock=None ): @@ -92,13 +94,20 @@ def add_years( projection : str Projection mode. If "constant", the last year of the input array is copied to every new year. If "empty", values are initialized and - set to zero. If a float array is given, these are used to populate + set to NaN. If a float array is given, these are used to populate the new year using a scaling of the last year of the array + pivot_year : int, optional + The year to use as the pivot for projections. If not specified, the last + year of the dataset is used. """ data = get_dict(datablock, dataset) - data = data.fbs.add_years(years, projection) + data = data.fbs.add_years( + years, + pivot_year=pivot_year, + projection=projection + ) set_dict(datablock, dataset, data) diff --git a/agrifoodpy/utils/scaling.py b/agrifoodpy/utils/scaling.py index a98ff9d..f6de6a8 100644 --- a/agrifoodpy/utils/scaling.py +++ b/agrifoodpy/utils/scaling.py @@ -4,10 +4,22 @@ import xarray as xr +def _validate_non_decreasing_years(years): + """Validate that years are in non-decreasing order.""" + + if not all(earlier <= later for earlier, later in zip(years, years[1:])): + raise ValueError("Years must be in ascending order.") + + def logistic_scale(y0, y1, y2, y3, c_init, c_end): """ Create an xarray DataArray with a logistic growth interval + The returned DataArray contains values that remain constant at c_init from + year y0 to y1, then transition following a logistic curve from c_init to + c_end between years y1 and y2, and finally remain constant at c_end from + year y2 to y3. + Parameters: y0 : (int) The initial year. @@ -23,19 +35,22 @@ def logistic_scale(y0, y1, y2, y3, c_init, c_end): The final constant value. Returns: xarray DataArray - An xarray DataArray object with 'year' as the coordinate and values set + An xarray DataArray object with 'Year' as the coordinate and values set by a logistic growth between the user defined intervals. """ + _validate_non_decreasing_years([y0, y1, y2, y3]) + # Create arrays and set values between y0 and y1 to c_init years = np.arange(y0, y3+1) values = np.ones_like(years, dtype=float) * c_init # Set values between y1 and y2 using a logistic curve - var_segment = np.logical_and(years >= y1, years < y2) - t = (years[var_segment] - y1) / (y2 - y1) - values[var_segment] = c_init \ - + (c_end - c_init)*(1 / (1 + np.exp(-10 * (t - 0.5)))) + if y1 < y2: + var_segment = np.logical_and(years >= y1, years < y2) + t = (years[var_segment] - y1) / (y2 - y1) + values[var_segment] = c_init \ + + (c_end - c_init)*(1 / (1 + np.exp(-10 * (t - 0.5)))) # Set values between y2 and y3 to c_end values[years >= y2] = c_end @@ -46,12 +61,11 @@ def logistic_scale(y0, y1, y2, y3, c_init, c_end): def linear_scale(y0, y1, y2, y3, c_init, c_end): """ - Create an xarray DataArray with a single coordinate called 'year'. + Create an xarray DataArray with a linear growth interval. - The values from the first year 'y0' up to a given year 'y1' will be - constant, then they will vary linearly between 'y1' and another given year - 'y2', and from 'y2' until the end of the array 'y3', the values will - continue with a constant value. + The returned DataArray contains values that remain constant at c_init from + year y0 to y1, then transition linearly from c_init to c_end between years + y1 and y2, and finally remain constant at c_end from year y2 to y3. Parameters: y0 : (int) @@ -67,25 +81,273 @@ def linear_scale(y0, y1, y2, y3, c_init, c_end): c_end : (float) Returns: xr.DataArray - An xarray DataArray object with 'year' as the coordinate and values set + An xarray DataArray object with 'Year' as the coordinate and values set by a linear growth between the user defined intervals. """ - # Create arrays and set values between y0 and y1 to c_init - years = np.arange(y0, y3 + 1) - values = np.ones_like(years, dtype=float) * c_init + return piecewise_linear_scale( + years=[y0, y1, y2, y3], + values=[c_init, c_init, c_end, c_end] + ) - # Set values between y1 and y2 using straight line. - var_segment = np.logical_and(years >= y1, years < y2) - if y2 == y1: - slope = c_end - c_init - else: - slope = float((c_end - c_init) / (y2 - y1)) - values[var_segment] = slope * (years[var_segment] - y1) + c_init +def piecewise_linear_scale(years, values): + """ + Create an xarray DataArray with a piecewise linear scale. - # Set values between y2 and y3 to c_end - values[years >= y2] = c_end + The returned DataArray contains values that transition linearly between + specified years and values. + + Parameters: + years : (list of int) + A list of years where the values change. + values : (list of float) + A list of values corresponding to each year in the 'years' list. + + Returns: xr.DataArray + An xarray DataArray object with 'Year' as the coordinate and values set + by a piecewise linear scale between the user defined intervals. + """ + + if len(years) != len(values): + raise ValueError("Length of 'years' and 'values' must be the same.") + + _validate_non_decreasing_years(years) + + # Create arrays for the full range of years + full_years = np.arange(years[0], years[-1] + 1) + full_values = np.zeros_like(full_years, dtype=float) + + # Fill in the values for each segment + for i in range(len(years) - 1): + var_segment = np.logical_and( + full_years >= years[i], + full_years < years[i + 1] + ) + + delta_years = years[i + 1] - years[i] + if delta_years == 0: + continue + + slope = (values[i + 1] - values[i]) / delta_years + + full_values[var_segment] = (slope + *(full_years[var_segment] - years[i]) + +values[i]) + + # Set the value for the last year + full_values[full_years >= years[-1]] = values[-1] + + data_array = xr.DataArray( + full_values, + coords={'Year': full_years}, + dims=['Year'] + ) + + return data_array + +def piecewise_constant_scale(years, values): + """ + Create an xarray DataArray with a piecewise constant scale. + + The returned DataArray contains values that remain constant between + specified years and values. + + Parameters: + years : (list of int) + A list of years where the values change. + values : (list of float) + A list of values corresponding to each year in the 'years' list. + + Returns: xr.DataArray + An xarray DataArray object with 'Year' as the coordinate and values set + by a piecewise constant scale between the user defined intervals. + """ + + if len(years) != len(values): + raise ValueError("Length of 'years' and 'values' must be the same.") + + _validate_non_decreasing_years(years) + + # Create arrays for the full range of years + full_years = np.arange(years[0], years[-1] + 1) + full_values = np.zeros_like(full_years, dtype=float) + + # Fill in the values for each segment + for i in range(len(years) - 1): + var_segment = np.logical_and( + full_years >= years[i], + full_years < years[i + 1] + ) + + full_values[var_segment] = values[i] + + # Set the value for the last year + full_values[full_years >= years[-1]] = values[-1] + + data_array = xr.DataArray( + full_values, + coords={'Year': full_years}, + dims=['Year'] + ) + + return data_array + +def step_scale(y0, y1, y2, c_init, c_end): + """ + Create an xarray DataArray with a step function. + + The returned DataArray contains values that remain constant at c_init from + year y0 to y1, then transition immediately to c_end at year y1 and remain + constant at c_end thereafter. + + Parameters: + y0 : (int) + Starting year. + y1 : (int) + Year where the step change occurs. + y2 : (int) + Last year in the array. + c_init : (float) + Value to use for initial constant scale segment. + c_end : (float) + Value to use for final constant scale segment. - data_array = xr.DataArray(values, coords={'Year': years}, dims=['Year']) + Returns: xr.DataArray + An xarray DataArray object with 'Year' as the coordinate and values set + by a step change between the user defined intervals. + """ + + return piecewise_constant_scale( + years=[y0, y1, y2], + values=[c_init, c_end, c_end] + ) + +def pulse_scale(y0, y1, y2, y3, c_init, c_pulse): + """ + Create an xarray DataArray with a pulse function. + + The returned DataArray contains values that remain constant at c_init from + year y0 to y1, then transition immediately to c_pulse at year y1 and remain + constant at c_pulse until year y2, after which it transitions back to + c_init and remains constant at c_init until year y3. + + Parameters: + y0 : (int) + Starting year. + y1 : (int) + Year where the pulse starts. + y2 : (int) + Year where the pulse ends. + y3 : (int) + Last year in the array. + c_init : (float) + Value to use for initial and final constant scale segments. + c_pulse : (float) + Value to use for the pulse segment. + + Returns: xr.DataArray + An xarray DataArray object with 'Year' as the coordinate and values set + by a pulse function between the user defined intervals. + """ + + return piecewise_constant_scale( + years=[y0, y1, y2, y3], + values=[c_init, c_pulse, c_init, c_init] + ) + + +def piecewise_smoothstep_scale(years, values, smoother=False): + """ + Create an xarray DataArray with a piecewise smoothstep scale. + + The returned DataArray contains values that transition smoothly between + specified years and values using a smoothstep function. + + Parameters: + years : (list of int) + A list of years where the values change. + values : (list of float) + A list of values corresponding to each year in the 'years' list. + smoother : (bool), optional + If True, use a smoother 5th-degree polynomial smoothstep; + otherwise use the standard smoothstep. + + Returns: xr.DataArray + An xarray DataArray object with 'Year' as the coordinate and values set + by a piecewise smoothstep scale between the user defined intervals. + """ + + if len(years) != len(values): + raise ValueError("Length of 'years' and 'values' must be the same.") + + _validate_non_decreasing_years(years) + + # Create arrays for the full range of years + full_years = np.arange(years[0], years[-1] + 1) + full_values = np.zeros_like(full_years, dtype=float) + + # Fill in the values for each segment + for i in range(len(years) - 1): + var_segment = np.logical_and( + full_years >= years[i], + full_years < years[i + 1] + ) + + delta_years = years[i + 1] - years[i] + if delta_years == 0: + continue + + t = (full_years[var_segment] - years[i]) / delta_years + + if smoother: + # Use a smoother smoothstep function (5th degree polynomial) + smoothstep = t**3 * (t * (t * 6 - 15) + 10) + else: + smoothstep = t**2 * (3 - 2 * t) + + full_values[var_segment] = (values[i] + + (values[i + 1] - values[i]) * smoothstep) + + # Set the value for the last year + full_values[full_years >= years[-1]] = values[-1] + + data_array = xr.DataArray( + full_values, + coords={'Year': full_years}, + dims=['Year'] + ) return data_array + +def smoothstep_scale(y0, y1, y2, y3, c_init, c_end, smoother=False): + """ + Create an xarray DataArray with a smoothstep function. + + The returned DataArray contains values that remain constant at c_init from + year y0 to y1, then transition smoothly to c_end between years y1 and y2, + and finally remain constant at c_end from year y2 to y3. + + Parameters: + y0 : (int) + Starting year. + y1 : (int) + Year where the smooth transition starts. + y2 : (int) + Year where the smooth transition ends. + y3 : (int) + Last year in the array. + c_init : (float) + Value to use for initial constant scale segment. + c_end : (float) + Value to use for final constant scale segment. + + Returns: xr.DataArray + An xarray DataArray object with 'Year' as the coordinate and values set + by a smoothstep function between the user defined intervals. + """ + + return piecewise_smoothstep_scale( + years=[y0, y1, y2, y3], + values=[c_init, c_init, c_end, c_end], + smoother=smoother + ) \ No newline at end of file diff --git a/agrifoodpy/utils/tests/test_add_years.py b/agrifoodpy/utils/tests/test_add_years.py index ed57e54..930c72a 100644 --- a/agrifoodpy/utils/tests/test_add_years.py +++ b/agrifoodpy/utils/tests/test_add_years.py @@ -2,7 +2,7 @@ import xarray as xr -def test_add_items(): +def test_add_years(): from agrifoodpy.utils.nodes import add_years diff --git a/agrifoodpy/utils/tests/test_scaling.py b/agrifoodpy/utils/tests/test_scaling.py index 5b6e86e..c83e43f 100644 --- a/agrifoodpy/utils/tests/test_scaling.py +++ b/agrifoodpy/utils/tests/test_scaling.py @@ -1,9 +1,21 @@ import numpy as np -from agrifoodpy.utils.scaling import linear_scale, logistic_scale - - -def test_logistic_scale(): - +import pytest +from agrifoodpy.utils.scaling import ( + linear_scale, + logistic_scale, + piecewise_constant_scale, + piecewise_linear_scale, + piecewise_smoothstep_scale, + pulse_scale, + smoothstep_scale, + step_scale, +) + +# --------------------------------- +# Tests for logistic_scale function +# --------------------------------- + +def test_logistic_scale_basic(): # Basic functionality test y0, y1, y2, y3 = 2000, 2005, 2010, 2015 c_init, c_end = 0, 10 @@ -16,6 +28,7 @@ def test_logistic_scale(): assert np.allclose(basic_result, truth) assert np.array_equal(basic_result["Year"].values, np.arange(2000, 2016)) +def test_logistic_scale_negative(): # Negative values test y0, y1, y2, y3 = 2000, 2005, 2010, 2015 c_init, c_end = -1, -10 @@ -26,14 +39,7 @@ def test_logistic_scale(): assert np.allclose(negative_result, truth) - # Instant change test - y0, y1, y2, y3 = 2000, 2001, 2001, 2015 - c_init, c_end = 0, 10 - instant_change_result = logistic_scale(y0, y1, y2, y3, c_init, c_end) - - assert instant_change_result[0] == 0 - assert instant_change_result[1] == 10 - +def test_logistic_scale_from_first_year(): # Change from first year y0, y1, y2, y3 = 2000, 2000, 2000, 2015 c_init, c_end = 0, 10 @@ -41,6 +47,7 @@ def test_logistic_scale(): assert np.array_equal(change_first_year, c_end * np.ones(y3+1-y0)) +def test_logistic_scale_constant_value(): # Constant value y0, y1, y2, y3 = 2000, 2000, 2000, 2015 c_init, c_end = 5.5, 5.5 @@ -49,9 +56,24 @@ def test_logistic_scale(): assert np.array_equal(constant_value, c_init * np.ones(y3+1-y0)) assert np.array_equal(constant_value, c_end * np.ones(y3+1-y0)) +def test_logistic_scale_immediate_transition(): + # y1 == y2 means no transition interval, so values jump to c_end at y1 + result = logistic_scale(2000, 2005, 2005, 2010, 0, 10) + truth = [0, 0, 0, 0, 0, 10, 10, 10, 10, 10, 10] -def test_linear_scale(): + assert np.allclose(result, truth) + assert np.array_equal(result["Year"].values, np.arange(2000, 2011)) + +def test_logistic_scale_raises_for_decreasing_years(): + with pytest.raises(ValueError, match="Years must be in ascending order"): + logistic_scale(2000, 2005, 2004, 2010, 0, 10) + +# ------------------------------- +# Tests for linear_scale function +# ------------------------------- + +def test_linear_scale(): # Basic functionality test basic_result = linear_scale(2000, 2005, 2010, 2015, 0, 10) truth = [0, 0, 0, 0, 0, 0, 2, 4, 6, 8, 10, 10, 10, 10, 10, 10] @@ -59,6 +81,7 @@ def test_linear_scale(): assert np.allclose(basic_result, truth) assert np.array_equal(basic_result["Year"].values, np.arange(2000, 2016)) +def test_linear_scale_negative(): # Negative values test y0, y1, y2, y3 = 2000, 2005, 2010, 2015 c_init, c_end = -1, -10 @@ -68,14 +91,7 @@ def test_linear_scale(): assert np.allclose(negative_result, truth) - # Instant change test - y0, y1, y2, y3 = 2000, 2001, 2001, 2015 - c_init, c_end = 0, 10 - instant_change_result = linear_scale(y0, y1, y2, y3, c_init, c_end) - - assert instant_change_result[0] == 0 - assert instant_change_result[1] == 10 - +def test_linear_scale_from_first_year(): # Change from first year y0, y1, y2, y3 = 2000, 2000, 2000, 2015 c_init, c_end = 0, 10 @@ -83,6 +99,7 @@ def test_linear_scale(): assert np.array_equal(change_first_year, c_end * np.ones(y3+1-y0)) +def test_linear_scale_constant_value(): # Constant value y0, y1, y2, y3 = 2000, 2000, 2000, 2015 c_init, c_end = 5.5, 5.5 @@ -90,3 +107,147 @@ def test_linear_scale(): assert np.array_equal(constant_value, c_init * np.ones(y3+1-y0)) assert np.array_equal(constant_value, c_end * np.ones(y3+1-y0)) + +def test_linear_scale_start_immediately(): + # y0 == y1 should start linear growth immediately at y0 + result = linear_scale(2000, 2000, 2005, 2010, 0, 10) + truth = [0, 2, 4, 6, 8, 10, 10, 10, 10, 10, 10] + + assert np.allclose(result, truth) + assert np.array_equal(result["Year"].values, np.arange(2000, 2011)) + +def test_linear_scale_end_last_year(): + # y2 == y3 should still return a valid series ending at c_end + result = linear_scale(2000, 2005, 2010, 2010, 0, 10) + truth = [0, 0, 0, 0, 0, 0, 2, 4, 6, 8, 10] + + assert np.allclose(result, truth) + assert np.array_equal(result["Year"].values, np.arange(2000, 2011)) + +def test_linear_scale_raises_for_decreasing_years(): + with pytest.raises(ValueError, match="Years must be in ascending order"): + linear_scale(2000, 2005, 2004, 2010, 0, 10) + + +# -------------------------------------- +# Tests for piecewise_linear_scale +# -------------------------------------- + +def test_piecewise_linear_scale_basic(): + result = piecewise_linear_scale([2000, 2003, 2005], [0, 6, 10]) + truth = [0, 2, 4, 6, 8, 10] + + assert np.allclose(result, truth) + assert np.array_equal(result["Year"].values, np.arange(2000, 2006)) + + +def test_piecewise_linear_scale_raises_for_mismatched_lengths(): + with pytest.raises(ValueError, match="Length of 'years' and 'values' must be the same"): + piecewise_linear_scale([2000, 2005], [0]) + + +def test_piecewise_linear_scale_raises_for_decreasing_years(): + with pytest.raises(ValueError, match="Years must be in ascending order"): + piecewise_linear_scale([2000, 2005, 2004], [0, 1, 2]) + + +# -------------------------------------- +# Tests for piecewise_constant_scale +# -------------------------------------- + +def test_piecewise_constant_scale_basic(): + result = piecewise_constant_scale([2000, 2003, 2005], [1, 3, 7]) + truth = [1, 1, 1, 3, 3, 7] + + assert np.allclose(result, truth) + assert np.array_equal(result["Year"].values, np.arange(2000, 2006)) + + +def test_piecewise_constant_scale_raises_for_mismatched_lengths(): + with pytest.raises(ValueError, match="Length of 'years' and 'values' must be the same"): + piecewise_constant_scale([2000, 2005], [1]) + + +def test_piecewise_constant_scale_raises_for_decreasing_years(): + with pytest.raises(ValueError, match="Years must be in ascending order"): + piecewise_constant_scale([2000, 2005, 2004], [1, 2, 3]) + + +# -------------------------------------- +# Tests for step_scale and pulse_scale +# -------------------------------------- + +def test_step_scale_basic(): + result = step_scale(2000, 2003, 2005, 2, 9) + truth = [2, 2, 2, 9, 9, 9] + + assert np.allclose(result, truth) + assert np.array_equal(result["Year"].values, np.arange(2000, 2006)) + + +def test_step_scale_immediate_transition(): + result = step_scale(2000, 2000, 2003, 2, 9) + truth = [9, 9, 9, 9] + + assert np.allclose(result, truth) + assert np.array_equal(result["Year"].values, np.arange(2000, 2004)) + + +def test_pulse_scale_basic(): + result = pulse_scale(2000, 2002, 2004, 2006, 1, 5) + truth = [1, 1, 5, 5, 1, 1, 1] + + assert np.allclose(result, truth) + assert np.array_equal(result["Year"].values, np.arange(2000, 2007)) + + +# -------------------------------------- +# Tests for piecewise_smoothstep_scale +# and smoothstep_scale +# -------------------------------------- + +def test_piecewise_smoothstep_scale_basic(): + result = piecewise_smoothstep_scale([2000, 2002], [0.0, 1.0]) + truth = [0.0, 0.5, 1.0] + + assert np.allclose(result, truth) + assert np.array_equal(result["Year"].values, np.arange(2000, 2003)) + + +def test_piecewise_smoothstep_scale_smoother_basic(): + result = piecewise_smoothstep_scale([2000, 2002], [0.0, 1.0], smoother=True) + truth = [0.0, 0.5, 1.0] + + assert np.allclose(result, truth) + assert np.array_equal(result["Year"].values, np.arange(2000, 2003)) + + +def test_piecewise_smoothstep_scale_raises_for_mismatched_lengths(): + with pytest.raises(ValueError, match="Length of 'years' and 'values' must be the same"): + piecewise_smoothstep_scale([2000, 2005], [0.0]) + + +def test_piecewise_smoothstep_scale_raises_for_decreasing_years(): + with pytest.raises(ValueError, match="Years must be in ascending order"): + piecewise_smoothstep_scale([2000, 2005, 2004], [0.0, 0.5, 1.0]) + + +def test_smoothstep_scale_wrapper_matches_piecewise(): + result = smoothstep_scale(2000, 2002, 2004, 2006, 1.0, 3.0) + expected = piecewise_smoothstep_scale([2000, 2002, 2004, 2006], [1.0, 1.0, 3.0, 3.0]) + + assert np.allclose(result, expected) + assert np.array_equal(result["Year"].values, np.arange(2000, 2007)) + + +def test_smoothstep_scale_wrapper_matches_piecewise_smoother(): + result = smoothstep_scale(2000, 2002, 2004, 2006, 1.0, 3.0, smoother=True) + expected = piecewise_smoothstep_scale( + [2000, 2002, 2004, 2006], + [1.0, 1.0, 3.0, 3.0], + smoother=True, + ) + + assert np.allclose(result, expected) + assert np.array_equal(result["Year"].values, np.arange(2000, 2007)) + diff --git a/docs/config_file.rst b/docs/config_file.rst index 7e09384..55c915b 100644 --- a/docs/config_file.rst +++ b/docs/config_file.rst @@ -51,3 +51,53 @@ and ``function`` is the name of the function to be executed. The parameters are specified as a dictionary of key-value pairs, where the keys are the parameter names. +YAML constructors +----------------- + +Configuration files can also use YAML tags to construct values during parsing, +which can be useful for creating complex objects without the need to read the +pipeline in a python environment and create the objects there. + +The agrifoodpy library provides access to arbitrary numpy and xarray functions +through the following YAML constructors: + +* ``!numpy.`` +* ``!xarray.`` + +Additionally, there are specific constructors for utility functions defined in +agrifoodpy: + +* ``!scale.linear`` +* ``!scale.logistic`` +* ``!scale.step`` +* ``!scale.pulse`` +* ``!scale.smoothstep`` +* ``!scale.piecewise_linear`` +* ``!scale.piecewise_constant`` +* ``!scale.piecewise_smoothstep`` + +The ``!scale`` constructors call scaling utilities from +``agrifoodpy.utils.scaling`` and return an ``xarray.DataArray``. + +Example with positional arguments: + +.. code-block:: yaml + + nodes: + - function: agrifoodpy.utils.nodes.write_to_datablock + name: Linear scale + params: + key: my_scale + value: !scale.linear [2020, 2022, 2024, 2026, 1.0, 3.0] + +Example with keyword arguments: + +.. code-block:: yaml + + nodes: + - function: agrifoodpy.utils.nodes.write_to_datablock + name: Logistic scale + params: + key: my_scale + value: !scale.logistic {y0: 2020, y1: 2022, y2: 2024, y3: 2026, c_init: 1.0, c_end: 3.0} +