diff --git a/nsvb/estimators.py b/nsvb/estimators.py index 77f5a39..983ddbf 100644 --- a/nsvb/estimators.py +++ b/nsvb/estimators.py @@ -1,3 +1,8 @@ +from typing import Union + +import numpy as np +from numpy.typing import ArrayLike, NDArray + from nsvb.models import MODEL_MAP from nsvb.tables import REF_SPECIES, TABLES @@ -5,10 +10,15 @@ def _run_model_form( - table_name: str, spcd: int, dia: float, ht: float, division: str = "" -) -> float: + table_name: str, + spcd: Union[int, ArrayLike], + dia: Union[float, ArrayLike], + ht: Union[float, ArrayLike], + division: Union[str, ArrayLike] = "", +) -> Union[float, NDArray]: """ Run the model form for the given table. + Works with scalar or array inputs. Parameters: table_name (str): Table name. @@ -20,24 +30,41 @@ def _run_model_form( Returns: float: Model form result. """ - try: - table_name_spcd = f"{table_name}a" - table_data = TABLES[table_name_spcd] - data = table_data.get((spcd, division), table_data[(spcd, "")]) - except KeyError: - spgrp = int(REF_SPECIES[spcd]["JENKINS_SPGRPCD"]) - table_name_spgrp = f"{table_name}b" - table_data = TABLES[table_name_spgrp] - data = table_data.get(spgrp) - wdsg = float(REF_SPECIES[spcd]["WOOD_SPGR_GREENVOL_DRYWT"]) - data["wdsg"] = wdsg - model_function = MODEL_MAP[data["model"]] - return model_function(dia, ht, **data) + # Scalar implementation (unchanged) + def _scalar_lookup(spcd_val, dia_val, ht_val, div_val): + try: + table_name_spcd = f"{table_name}a" + table_data = TABLES[table_name_spcd] + data = table_data.get((spcd_val, div_val), table_data[(spcd_val, "")]) + except KeyError: + spgrp = int(REF_SPECIES[spcd_val]["JENKINS_SPGRPCD"]) + table_name_spgrp = f"{table_name}b" + table_data = TABLES[table_name_spgrp] + data = table_data.get(spgrp) + wdsg = float(REF_SPECIES[spcd_val]["WOOD_SPGR_GREENVOL_DRYWT"]) + data = data.copy() + data["wdsg"] = wdsg + model_function = MODEL_MAP[data["model"]] + return model_function(dia_val, ht_val, **data) + + # Check if inputs are arrays + is_array = isinstance(spcd, np.ndarray) or isinstance(dia, np.ndarray) or isinstance(ht, np.ndarray) + + if is_array: + # Vectorize the scalar function + vectorized_fn = np.vectorize(_scalar_lookup) + return vectorized_fn(spcd, dia, ht, division) + else: + # Use scalar path directly + return _scalar_lookup(int(spcd), float(dia), float(ht), str(division)) def total_inside_bark_wood_volume( - spcd: int, dia: float, ht: float, division: str = "" -) -> float: + spcd: Union[int, ArrayLike], + dia: Union[float, ArrayLike], + ht: Union[float, ArrayLike], + division: Union[str, ArrayLike] = "", +) -> Union[float, NDArray]: """ Predict gross total stem wood volume as a function of diameter at breast height (D) and @@ -59,8 +86,11 @@ def total_inside_bark_wood_volume( def total_bark_wood_volume( - spcd: int, dia: float, ht: float, division: str = "" -) -> float: + spcd: Union[int, ArrayLike], + dia: Union[float, ArrayLike], + ht: Union[float, ArrayLike], + division: Union[str, ArrayLike] = "", +) -> Union[float, NDArray]: """ Predict gross total stem bark volume as a function of D and H. Uses the appropriate model form and coefficients from table S2. @@ -80,8 +110,11 @@ def total_bark_wood_volume( def total_outside_bark_volume( - spcd: int, dia: float, ht: float, division: str = "" -) -> float: + spcd: Union[int, ArrayLike], + dia: Union[float, ArrayLike], + ht: Union[float, ArrayLike], + division: Union[str, ArrayLike] = "", +) -> Union[float, NDArray]: """ Obtain gross total stem outside-bark volume as the sum of wood and bark gross volumes. @@ -103,8 +136,12 @@ def total_outside_bark_volume( def total_stem_wood_dry_weight( - spcd: int, dia: float, ht: float, division: str = "", cull: float = 0 -) -> float: + spcd: Union[int, ArrayLike], + dia: Union[float, ArrayLike], + ht: Union[float, ArrayLike], + division: Union[str, ArrayLike] = "", + cull: Union[float, ArrayLike] = 0, +) -> Union[float, NDArray]: """ Convert total stem wood gross volume to biomass weight using published wood density @@ -113,6 +150,8 @@ def total_stem_wood_dry_weight( (accounting for nonzero weight of cull), and dead tree wood density reduction. + Works with scalar or array inputs. + Corresponds to step 7 of "Examples of Tree-Level Calculations" in the GTR. Parameters: @@ -125,31 +164,54 @@ def total_stem_wood_dry_weight( Returns: float: Total stem wood dry weight in pounds (lb). """ - wdsg = float(REF_SPECIES[spcd]["WOOD_SPGR_GREENVOL_DRYWT"]) - v_tot_ib = total_inside_bark_wood_volume(spcd, dia, ht, division) - - if cull > 0: - # It is considered that most cull will be rotten wood, which would - # still contribute to the stem weight. As such, it is assumed the - # density of cull wood is reduced by the proportion for DECAYCD = 3 - # (see table 1; wood density proportion (DensProp) is 0.54 for - # hardwood species and 0.92 for softwood species) - dens_prop = 0.54 if REF_SPECIES[spcd]["SFTWD_HRDWD"] == "H" else 0.92 + # Check if array input + is_array = isinstance(spcd, np.ndarray) or isinstance(dia, np.ndarray) - return ( - v_tot_ib - * (1 - cull / 100 * (1 - dens_prop)) - * wdsg - * WEIGHT_CUBIC_FOOT_WATER - ) - - # No cull - return v_tot_ib * wdsg * WEIGHT_CUBIC_FOOT_WATER + # Scalar path + if not is_array: + wdsg = float(REF_SPECIES[spcd]["WOOD_SPGR_GREENVOL_DRYWT"]) + v_tot_ib = total_inside_bark_wood_volume(spcd, dia, ht, division) + + if cull > 0: + dens_prop = 0.54 if REF_SPECIES[spcd]["SFTWD_HRDWD"] == "H" else 0.92 + return ( + v_tot_ib + * (1 - cull / 100 * (1 - dens_prop)) + * wdsg + * WEIGHT_CUBIC_FOOT_WATER + ) + return v_tot_ib * wdsg * WEIGHT_CUBIC_FOOT_WATER + + # Array path + spcd_arr = np.atleast_1d(spcd) + dia_arr = np.atleast_1d(dia) + ht_arr = np.atleast_1d(ht) + div_arr = np.atleast_1d(division) + cull_arr = np.atleast_1d(cull) + + # Get volume (will be vectorized through _run_model_form) + v_tot_ib = total_inside_bark_wood_volume(spcd_arr, dia_arr, ht_arr, div_arr) + + # Vectorize lookups from REF_SPECIES + wdsg_arr = np.array([float(REF_SPECIES[int(s)]["WOOD_SPGR_GREENVOL_DRYWT"]) for s in spcd_arr]) + dens_prop_arr = np.array([0.54 if REF_SPECIES[int(s)]["SFTWD_HRDWD"] == "H" else 0.92 for s in spcd_arr]) + + # Vectorized calculation + weight = np.where( + cull_arr > 0, + v_tot_ib * (1 - cull_arr / 100 * (1 - dens_prop_arr)) * wdsg_arr * WEIGHT_CUBIC_FOOT_WATER, + v_tot_ib * wdsg_arr * WEIGHT_CUBIC_FOOT_WATER + ) + + return weight def total_stem_bark_weight( - spcd: int, dia: float, ht: float, division: str = "" -) -> float: + spcd: Union[int, ArrayLike], + dia: Union[float, ArrayLike], + ht: Union[float, ArrayLike], + division: Union[str, ArrayLike] = "", +) -> Union[float, NDArray]: """ Predict total stem bark biomass as a function of D and H. Reduce the prediction if necessary for @@ -172,7 +234,12 @@ def total_stem_bark_weight( return _run_model_form("s6", spcd, dia, ht, division) -def total_branch_weight(spcd: int, dia: float, ht: float, division: str = "") -> float: +def total_branch_weight( + spcd: Union[int, ArrayLike], + dia: Union[float, ArrayLike], + ht: Union[float, ArrayLike], + division: Union[str, ArrayLike] = "", +) -> Union[float, NDArray]: """ Predict total branch biomass as a function of D and H. Reduce the prediction if necessary for @@ -196,8 +263,11 @@ def total_branch_weight(spcd: int, dia: float, ht: float, division: str = "") -> def total_aboveground_biomass( - spcd: int, dia: float, ht: float, division: str = "" -) -> float: + spcd: Union[int, ArrayLike], + dia: Union[float, ArrayLike], + ht: Union[float, ArrayLike], + division: Union[str, ArrayLike] = "", +) -> Union[float, NDArray]: """ Predict total aboveground biomass as a function of D and H. Reduce the prediction if necessary @@ -223,8 +293,11 @@ def total_aboveground_biomass( def total_foliage_dry_weight( - spcd: int, dia: float, ht: float, division: str = "" -) -> float: + spcd: Union[int, ArrayLike], + dia: Union[float, ArrayLike], + ht: Union[float, ArrayLike], + division: Union[str, ArrayLike] = "", +) -> Union[float, NDArray]: """ Directly predict total foliage dry weight as a function of D and H. Use the appropriate model diff --git a/nsvb/models.py b/nsvb/models.py index 736cde2..11ffabb 100644 --- a/nsvb/models.py +++ b/nsvb/models.py @@ -1,7 +1,12 @@ -from math import exp +from typing import Union +import numpy as np +from numpy.typing import ArrayLike, NDArray -def schumacher_hall_method(dia: float, ht: float, **kwargs) -> float: + +def schumacher_hall_method( + dia: Union[float, ArrayLike], ht: Union[float, ArrayLike], **kwargs +) -> Union[float, NDArray]: """ Schumacher-Hall Method. @@ -22,7 +27,9 @@ def schumacher_hall_method(dia: float, ht: float, **kwargs) -> float: return a * (dia**b) * (ht**c) + e -def segmented_model(dia: float, ht: float, **kwargs) -> float: +def segmented_model( + dia: Union[float, ArrayLike], ht: Union[float, ArrayLike], **kwargs +) -> Union[float, NDArray]: """ Segmented Model. @@ -44,12 +51,16 @@ def segmented_model(dia: float, ht: float, **kwargs) -> float: c = kwargs.get("c") k = kwargs.get("k") e = kwargs.get("e", 0) - if dia < k: - return a * (dia**b) * (ht**c) + e - return a * (k ** (b - b1)) * (dia**b1) * (ht**c) + e + return np.where( + dia < k, + a * (dia**b) * (ht**c) + e, + a * (k ** (b - b1)) * (dia**b1) * (ht**c) + e + ) -def continuously_variable_model(dia: float, ht: float, **kwargs) -> float: +def continuously_variable_model( + dia: Union[float, ArrayLike], ht: Union[float, ArrayLike], **kwargs +) -> Union[float, NDArray]: """ Continuously Variable Model. @@ -71,10 +82,12 @@ def continuously_variable_model(dia: float, ht: float, **kwargs) -> float: c = kwargs.get("c") c1 = kwargs.get("c1") e = kwargs.get("e", 0) - return a * (a1 * ((1 - exp(-b * dia)) ** c1)) * (ht**c) + e + return a * (a1 * ((1 - np.exp(-b * dia)) ** c1)) * (ht**c) + e -def modifed_wiley_model(dia: float, ht: float, **kwargs) -> float: +def modifed_wiley_model( + dia: Union[float, ArrayLike], ht: Union[float, ArrayLike], **kwargs +) -> Union[float, NDArray]: """ Modified Wiley Model. @@ -94,10 +107,12 @@ def modifed_wiley_model(dia: float, ht: float, **kwargs) -> float: b1 = kwargs.get("b1") c = kwargs.get("c") e = kwargs.get("e", 0) - return a * (dia**b) * (ht**c) * exp(-(b1 * dia)) + e + return a * (dia**b) * (ht**c) * np.exp(-(b1 * dia)) + e -def modified_schumaker_hall(dia: float, ht: float, **kwargs) -> float: +def modified_schumaker_hall( + dia: Union[float, ArrayLike], ht: Union[float, ArrayLike], **kwargs +) -> Union[float, NDArray]: """ Modified Schumacher-Hall Method. diff --git a/setup.py b/setup.py index 7354fd0..e1c3c20 100644 --- a/setup.py +++ b/setup.py @@ -52,5 +52,5 @@ def get_version(): package_data={"nsvb": ["data/*"]}, include_package_data=True, python_requires=">=3.9", - install_requires=[], + install_requires=["numpy>=1.20.0"], ) diff --git a/tests/test_examples.py b/tests/test_examples.py index 5644eb9..8cccfd4 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1,3 +1,5 @@ +import numpy as np + from nsvb.estimators import ( total_inside_bark_wood_volume, total_bark_wood_volume, @@ -684,3 +686,125 @@ def test_total_foliage_dry_weight(self): total_foliage_dry_weight(self.spcd, self.dia, self.ht, self.division) == 47.82328163632339 ) + + +class TestVectorized: + """ + Test vectorization by combining all 4 examples into arrays. + + Uses the same trees from TestExample1-4: + - Example 1: Douglas-fir (spcd=202, dia=20.0, ht=110, division="240") + - Example 2: Red maple (spcd=316, dia=11.1, ht=38, division="M210", cull=3) + - Example 3: Tanoak (spcd=631, dia=11.3, ht=28, division="M240") + - Example 4: White oak (spcd=802, dia=18.1, ht=65, division="M220", cull=2) + """ + + # Arrays of all 4 examples + spcd = np.array([202, 316, 631, 802]) + dia = np.array([20.0, 11.1, 11.3, 18.1]) + ht = np.array([110, 38, 28, 65]) + division = np.array(["240", "M210", "M240", "M220"]) + cull = np.array([0, 3, 0, 2]) # Examples 1 and 3 have no cull + + def test_inside_bark_wood_volume(self): + """Vectorized inside bark volume matches individual examples.""" + result = total_inside_bark_wood_volume(self.spcd, self.dia, self.ht, self.division) + + # Should return array + assert isinstance(result, np.ndarray) + assert len(result) == 4 + + # Should match each example's expected value + assert result[0] == 88.45229093648126 # Example 1 + assert result[1] == 9.42711333158677 # Example 2 + assert result[2] == 7.283116395242574 # Example 3 + assert result[3] == 42.27783673140729 # Example 4 + + def test_total_bark_wood_volume(self): + """Vectorized bark volume matches individual examples.""" + result = total_bark_wood_volume(self.spcd, self.dia, self.ht, self.division) + + assert isinstance(result, np.ndarray) + assert len(result) == 4 + + assert result[0] == 13.197130062388565 # Example 1 + assert result[1] == 2.1551061436670853 # Example 2 + assert result[2] == 1.9071364767677488 # Example 3 + assert result[3] == 8.361568897350095 # Example 4 + + def test_total_stem_wood_dry_weight(self): + """Vectorized stem wood weight matches individual examples.""" + # Test without cull first + result_no_cull = total_stem_wood_dry_weight( + self.spcd, self.dia, self.ht, self.division + ) + + assert isinstance(result_no_cull, np.ndarray) + assert len(result_no_cull) == 4 + + assert result_no_cull[0] == 2483.7403294963938 # Example 1 + assert result_no_cull[1] == 288.2434172265971 # Example 2 (no cull) + assert result_no_cull[2] == 263.59054857661926 # Example 3 + assert result_no_cull[3] == 1582.8822072238888 # Example 4 (no cull) + + # Test with cull + result_with_cull = total_stem_wood_dry_weight( + self.spcd, self.dia, self.ht, self.division, self.cull + ) + + assert isinstance(result_with_cull, np.ndarray) + assert len(result_with_cull) == 4 + + assert result_with_cull[0] == 2483.7403294963938 # Example 1 (cull=0) + assert result_with_cull[1] == 284.26565806887004 # Example 2 (cull=3) + assert result_with_cull[2] == 263.59054857661926 # Example 3 (cull=0) + # Example 4 with cull=2 commented out in original tests + # assert result_with_cull[3] == 1564.617593936140 + + def test_total_stem_bark_weight(self): + """Vectorized stem bark weight matches individual examples.""" + result = total_stem_bark_weight(self.spcd, self.dia, self.ht, self.division) + + assert isinstance(result, np.ndarray) + assert len(result) == 4 + + assert result[0] == 361.7824889136451 # Example 1 + assert result[1] == 52.94546582033252 # Example 2 + assert result[2] == 46.81666440280295 # Example 3 + assert result[3] == 237.1544176737046 # Example 4 + + def test_total_branch_weight(self): + """Vectorized branch weight matches individual examples.""" + result = total_branch_weight(self.spcd, self.dia, self.ht, self.division) + + assert isinstance(result, np.ndarray) + assert len(result) == 4 + + assert result[0] == 277.4877562341372 # Example 1 + assert result[1] == 135.00192318003036 # Example 2 + assert result[2] == 226.78800239146196 # Example 3 + assert result[3] == 770.2515898127575 # Example 4 + + def test_total_aboveground_biomass(self): + """Vectorized total aboveground biomass matches individual examples.""" + result = total_aboveground_biomass(self.spcd, self.dia, self.ht, self.division) + + assert isinstance(result, np.ndarray) + assert len(result) == 4 + + assert result[0] == 3154.553996629238 # Example 1 + assert result[1] == 532.5847996695031 # Example 2 + assert result[2] == 492.6214580952344 # Example 3 + # Example 4 doesn't have test for total_aboveground_biomass + + def test_total_foliage_dry_weight(self): + """Vectorized foliage weight matches individual examples.""" + result = total_foliage_dry_weight(self.spcd, self.dia, self.ht, self.division) + + assert isinstance(result, np.ndarray) + assert len(result) == 4 + + assert result[0] == 83.63478892024017 # Example 1 + assert result[1] == 22.807960628763336 # Example 2 + # Example 3 is dead tree, foliage = 0 (not tested in original) + assert result[3] == 47.82328163632339 # Example 4