Skip to content

Clamping function #1011

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vgro merged 1 commit into from
Sep 1, 2025
Merged

Clamping function #1011

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 45 additions & 27 deletions tests/models/abiotic/test_wind.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""Test module for abiotic.wind.py."""

import numpy as np
from numpy.testing import assert_allclose


def test_calculate_zero_plane_displacement(dummy_climate_data_varying_canopy):
Expand All @@ -16,7 +17,7 @@ def test_calculate_zero_plane_displacement(dummy_climate_data_varying_canopy):
zero_plane_scaling_parameter=7.5,
)

np.testing.assert_allclose(result, np.array([0.0, 0.0, 25.86256, 25.86256]))
assert_allclose(result, np.array([0.0, 0.0, 25.86256, 25.86256]))


def test_calculate_roughness_length_momentum(dummy_climate_data_varying_canopy):
Expand All @@ -38,7 +39,7 @@ def test_calculate_roughness_length_momentum(dummy_climate_data_varying_canopy):
min_roughness_length=0.01,
)

np.testing.assert_allclose(
assert_allclose(
result, np.array([0.01, 0.01666, 0.524479, 0.524479]), rtol=1e-3, atol=1e-3
)

Expand Down Expand Up @@ -72,7 +73,7 @@ def test_calculate_wind_profile(
]
)

np.testing.assert_allclose(result, exp_wind, rtol=1e-3, atol=1e-3)
assert_allclose(result, exp_wind, rtol=1e-3, atol=1e-3)


def test_calculate_friction_velocity(dummy_climate_data_varying_canopy):
Expand All @@ -92,7 +93,7 @@ def test_calculate_friction_velocity(dummy_climate_data_varying_canopy):
von_karman_constant=0.4,
)
exp_friction_velocity = np.array([0.080945, 0.085658, 0.099079, 0.099079])
np.testing.assert_allclose(result, exp_friction_velocity, rtol=1e-3, atol=1e-3)
assert_allclose(result, exp_friction_velocity, rtol=1e-3, atol=1e-3)


def test_calculate_ventilation_rate_scalar():
Expand Down Expand Up @@ -122,7 +123,7 @@ def test_calculate_ventilation_rate_array():
expected = np.array([5.0e-02, 1.0e-03, 1.0e03])

result = calculate_ventilation_rate(ra, h)
np.testing.assert_allclose(result, expected)
assert_allclose(result, expected)


def test_calculate_ventilation_rate_zero_denominator():
Expand Down Expand Up @@ -158,33 +159,50 @@ def test_calculate_mixing_coefficients():

assert result.shape == layer_midpoints.shape
assert np.all(result >= 0)
np.testing.assert_allclose(result, expected, rtol=1e-6)
assert_allclose(result, expected, rtol=1e-6)


def test_compute_excess_all_cases():
"""Test calculate excess."""
def test_clamp_variable_within_limits():
"""Test clamping of variable within limits."""

from virtual_ecosystem.models.abiotic.wind import (
calculate_excess,
)
from virtual_ecosystem.models.abiotic.wind import clamp_variable_within_limits

# Set limits
limits = (4, 6)

limits = (0.0, 5.0)
# Cells with mix of issues: undershoot and overshoot, can be absorbed by a single
# cell vs need to spread further up into canopy, empty layers and complete, and mix
# of under and overshoot. Two cases explicitly test the edge case where undershoot
# and overshoot cannot be absorbed within the canopy, leading to out of limit values
# in the top layer.

n = np.nan
variable = np.array(
[
[5, 5, 5, 5, 5, 4, 5, 5, 5, 5, 5, 5],
[5, 5, 5, 5, 5, 4, 5, 5, 5, 5, 5, 8],
[5, 5, 5, 5, 3, 4, 5, 5, 5, 7, 5, 2],
[5, n, 5, n, n, 4, 4, 5, n, n, 7, 7],
[5, 5, 3, 3, 5, 0, 2, 7, 7, 7, 9, 3],
]
)

# mixed array: NaN, undershoot, inside, edge, overshoot
vals = np.array([np.nan, -2.0, 2.5, 0.0, 5.0, 7.0, 10.0])
excess, valid = calculate_excess(vals, limits)
variable_expected = np.array(
[
[5, 5, 5, 5, 5, 0, 5, 5, 5, 6, 7, 5],
[5, 5, 5, 5, 4, 4, 4, 5, 5, 6, 6, 6],
[5, 5, 5, 4, 4, 4, 4, 5, 6, 6, 6, 4],
[5, n, 4, n, n, 4, 4, 6, n, n, 6, 6],
[5, 5, 4, 4, 5, 4, 4, 6, 6, 6, 6, 4],
]
)

# expected excess: [0, -2, 0, 0, 0, 2, 5]
expected_excess = np.array([0.0, -2.0, 0.0, 0.0, 0.0, 2.0, 5.0])
expected_valid = np.array([False, True, True, True, True, True, True])
clamped_variable = clamp_variable_within_limits(variable=variable, limits=limits)

assert np.allclose(excess, expected_excess)
assert np.array_equal(valid, expected_valid)
assert_allclose(clamped_variable, variable_expected)

# sanity check: for all valid values, (value - excess) lies within limits
corrected = vals.copy()
corrected[valid] = corrected[valid] - excess[valid]
assert np.all((corrected[valid] >= limits[0]) & (corrected[valid] <= limits[1]))
# sanity checks: the column sums should be maintained.
assert_allclose(variable.sum(axis=0), clamped_variable.sum(axis=0))


def test_mix_and_ventilate(dummy_climate_data_varying_canopy, fixture_core_components):
Expand Down Expand Up @@ -245,7 +263,7 @@ def test_mix_and_ventilate(dummy_climate_data_varying_canopy, fixture_core_compo
limits=(0, 100),
time_interval=3600.0,
)
np.testing.assert_allclose(result, exp_result, rtol=1e-6, atol=1e-6)
assert_allclose(result, exp_result, rtol=1e-6, atol=1e-6)


def test_advect_from_toplayer():
Expand All @@ -272,7 +290,7 @@ def test_advect_from_toplayer():
time_interval=time_interval,
)

np.testing.assert_allclose(result, expected_specific_humidity)
assert_allclose(result, expected_specific_humidity)


def test_calculate_aerodynamic_resistance(
Expand Down Expand Up @@ -302,4 +320,4 @@ def test_calculate_aerodynamic_resistance(
wind_speed=np.array([1.0, 2.0, 0.5, 0.01]),
von_karman_constant=0.4,
)
np.testing.assert_allclose(result, exp_ra, rtol=1e-3, atol=1e-3)
assert_allclose(result, exp_ra, rtol=1e-3, atol=1e-3)
81 changes: 44 additions & 37 deletions virtual_ecosystem/models/abiotic/wind.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -279,24 +279,52 @@ def calculate_mixing_coefficients_canopy(
return mixing_coefficients


def calculate_excess(
value: NDArray[np.floating], limits: tuple[float, float]
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
"""Calculate excess in layers during mixing in valid layers.
def clamp_variable_within_limits(
variable: NDArray[np.floating], limits: tuple[float, float]
) -> NDArray[np.floating]:
"""Clamp an array of canopy data within limits.

Args:
value: Value of variable that is mixed
limits: Upper and lower limit for input variable
This function iterates from the bottom of the canopy, clamping the values of the
input array within the limits. When a value is altered by clamping, the residual is
added to the layer above to maintain the variable total within cells. Residual
values may be redistributed across multiple layers and empty values (representing
unoccupied canopy layers) are skipped.

Returns:
net excess and valid array
Note:
If the vertical layers cannot absorb all of the accumulated residuals without
themselves being clamped, then the values in the top layer can still fall
outside the clamping limits.

Args:
variable: A numpy array containing canopy data.
limits: A tuple giving the upper and lower bounds within which to clamp the data
"""
low, high = limits
valid = ~np.isnan(value)
over = np.where(valid & (value > high), value - high, 0)
under = np.where(valid & (value < low), value - low, 0)

return over + under, valid
# Get a map of nan values and initialise the out_of_limits array
out_of_limits = np.zeros_like(variable[0])
nan_map = np.isnan(variable)
n_layers = variable.shape[0]

# Loop up from the row index of lowest layer, stopping before the top layer
for layer in np.arange(n_layers - 1, 0, -1):
# Calculate the clamped values for the current layer
in_limits = np.clip(variable[layer], *limits)

# Add under and overshoots to the out_of_limits array, trapping cells that
# contain no vegetation in the layer (np.nan)
out_of_limits += np.where(nan_map[layer], 0, variable[layer] - in_limits)

# Set the clamped data in the current layer
variable[layer] = in_limits

# Add out of limits to the layer above
variable[layer - 1] += out_of_limits
# Update out_of_limits
# - np.nan cells carry over the current out_of_limits total
# - otherwise the out_of_limits has been set into the layer above, so is zeroed
out_of_limits = np.where(nan_map[layer - 1], out_of_limits, 0)

return variable


def mix_and_ventilate(
Expand Down Expand Up @@ -386,30 +414,9 @@ def mix_and_ventilate(
input_variable_mixed[1] -= vent_change

# Redistribute overshoot/undershoot
valid = ~np.isnan(input_variable_mixed)
layer_indices = np.arange(n_layers)[:, None]
layer_indices_valid = np.where(valid, layer_indices, -1)
last_layer_above_valid = np.maximum.accumulate(layer_indices_valid, axis=0)
nearest_above = np.vstack(
[np.full((1, n_cells), -1, dtype=int), last_layer_above_valid[:-1, :]]
input_variable_mixed = clamp_variable_within_limits(
variable=input_variable_mixed, limits=limits
)
cols = np.arange(n_cells)

for layer in range(n_layers - 1, 0, -1):
value = input_variable_mixed[layer, :]
excess, valid_mask = calculate_excess(value=value, limits=limits)

input_variable_mixed[layer, valid_mask] -= excess[valid_mask]

# Redistribute upward
target_layer_index = nearest_above[layer, :]
mask = (excess != 0) & (target_layer_index >= 0)
if np.any(mask):
np.add.at(
input_variable_mixed,
(target_layer_index[mask], cols[mask]),
excess[mask],
)

return input_variable_mixed

Expand Down
Loading