test_constraints.py

# © Crown copyright, Met Office (2022-2024) and CSET contributors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Test constraint operators."""

from datetime import datetime

import iris
import numpy as np
import pytest

from CSET.operators import constraints


def test_generate_stash_constraint():
    """Generate iris cube constraint for UM STASH code."""
    stash_constraint = constraints.generate_stash_constraint("m01s03i236")
    expected_stash_constraint = "AttributeConstraint({'STASH': 'm01s03i236'})"
    assert repr(stash_constraint) == expected_stash_constraint


def test_generate_var_constraint():
    """Generate iris cube constraint for str variable name."""
    var_constraint = constraints.generate_var_constraint("test")
    expected_var_constraint = "Constraint(name='test')"
    assert repr(var_constraint) == expected_var_constraint


def test_generate_var_constraint_stash():
    """Generate iris cube constraint for UM STASH code with var constraint."""
    var_constraint = constraints.generate_var_constraint("m01s03i236")
    expected_stash_constraint = "AttributeConstraint({'STASH': 'm01s03i236'})"
    assert repr(var_constraint) == expected_stash_constraint


def test_generate_cell_methods_constraint():
    """Generate iris cube constraint for cell methods."""
    cell_methods_constraint = constraints.generate_cell_methods_constraint(["mean"])
    expected_cell_methods_constraint = "Constraint(cube_func=<function generate_cell_methods_constraint.<locals>.check_cell_methods at"
    assert expected_cell_methods_constraint in repr(cell_methods_constraint)


def test_generate_cell_methods_constraint_sum():
    """Generate aggregate iris cube constraint for cell methods."""
    cell_methods_constraint = constraints.generate_cell_methods_constraint(["sum"])
    expected_cell_methods_constraint = "Constraint(cube_func=<function generate_cell_methods_constraint.<locals>.check_cell_methods at"
    assert expected_cell_methods_constraint in repr(cell_methods_constraint)


def test_generate_cell_methods_constraint_no_aggregation():
    """Generate iris cube constraint for no aggregation cell methods."""
    cell_methods_constraint = constraints.generate_cell_methods_constraint([])
    expected_cell_methods_constraint = "Constraint(cube_func=<function generate_cell_methods_constraint.<locals>.check_no_aggregation at"
    assert expected_cell_methods_constraint in repr(cell_methods_constraint)


def test_generate_cell_methods_constraint_varname():
    """Generate variable-dependent iris cube constraint for cell methods."""
    cell_methods_constraint = constraints.generate_cell_methods_constraint(
        [], "number_of_lightning_flashes"
    )
    expected_cell_methods_constraint = "Constraint(cube_func=<function generate_cell_methods_constraint.<locals>.check_cell_sum at"
    assert expected_cell_methods_constraint in repr(cell_methods_constraint)


def test_generate_cell_methods_mean_constraint_varname():
    """Generate variable-dependent iris cube constrain for mean cell methods."""
    cell_methods_constraint = constraints.generate_cell_methods_constraint([], "albedo")
    expected_cell_methods_constraint = "Constraint(cube_func=<function generate_cell_methods_constraint.<locals>.check_cell_mean at"
    assert expected_cell_methods_constraint in repr(cell_methods_constraint)


def test_generate_time_constraint():
    """Generate iris cube constraint for dates."""
    # Try with str dates
    time_constraint = constraints.generate_time_constraint(
        "2023-03-24T00:00", "2023-03-24T06:00"
    )
    expected_time_constraint = "Constraint(coord_values={'time': <function generate_time_constraint.<locals>.<lambda> at "
    assert expected_time_constraint in repr(time_constraint)
    # Try with datetime.datetime dates
    time_constraint = constraints.generate_time_constraint(
        datetime.fromisoformat("2023-03-24T00:00:00+00:00"),
        datetime.fromisoformat("2023-03-24T06:00:00+00:00"),
    )
    assert expected_time_constraint in repr(time_constraint)
    # Try with implicit end
    time_constraint = constraints.generate_time_constraint("2023-03-24T00:00:00+00:00")
    assert expected_time_constraint in repr(time_constraint)


def test_generate_level_constraint_single_level():
    """Generate constraint for a single level."""
    pressure_constraint = constraints.generate_level_constraint(
        coordinate="pressure", levels=1000
    )
    expected_pressure_constraint = "Constraint(coord_values={'pressure': [1000]})"
    assert expected_pressure_constraint in repr(pressure_constraint)


def test_generate_level_constraint_multi_level():
    """Generate constraint for multiple pressure levels."""
    pressure_constraint = constraints.generate_level_constraint(
        coordinate="pressure", levels=[200, 800]
    )
    expected_pressure_constraint = "Constraint(coord_values={'pressure': [200, 800]})"
    assert expected_pressure_constraint in repr(pressure_constraint)


def test_generate_level_constraint_all_level():
    """Generate constraint for all levels."""
    pressure_constraint = constraints.generate_level_constraint(
        coordinate="pressure", levels="*"
    )
    expected_pressure_constraint = "Constraint(coord_values={'pressure': <function generate_level_constraint.<locals>.<lambda> at"
    assert expected_pressure_constraint in repr(pressure_constraint)


def test_generate_level_constraint_no_pressure():
    """Generate constraint for not having pressure levels."""
    pressure_constraint = constraints.generate_level_constraint(
        coordinate="pressure", levels=[]
    )
    expected_pressure_constraint = (
        "Constraint(cube_func=<function generate_level_constraint.<locals>.no_levels at"
    )
    assert expected_pressure_constraint in repr(pressure_constraint)


def test_generate_area_constraint():
    """Generate area constraint with lat-lon limits."""
    area_constraint = constraints.generate_area_constraint(0.0, 0.0, 0.1, 0.1)
    actual = repr(area_constraint)
    assert "Constraint(coord_values={" in actual
    assert (
        "'grid_latitude': <function generate_area_constraint.<locals>.bound_lat at 0x"
        in actual
    )
    assert (
        "'grid_longitude': <function generate_area_constraint.<locals>.bound_lon at 0x"
        in actual
    )


def test_generate_area_constraint_no_limits():
    """Generate area constraint with no limits."""
    area_constraint = constraints.generate_area_constraint(None, None, None, None)
    expected_area_constraint = "Constraint()"
    assert expected_area_constraint in repr(area_constraint)


def test_generate_area_constraint_invalid_arguments():
    """Generate area constraint raises exception with invalid arguments."""
    # Non-numbers are rejected.
    with pytest.raises(TypeError):
        constraints.generate_area_constraint(1, 2, 3, "four")

    # Mixed numbers and Nones are rejected.
    with pytest.raises(TypeError):
        constraints.generate_area_constraint(None, None, None, 0)


def test_generate_remove_single_ensemble_member_constraint():
    """Generate a constraint to remove a single ensemble member using default value."""
    single_member_constraint = (
        constraints.generate_remove_single_ensemble_member_constraint()
    )
    assert (
        "Constraint(coord_values={'realization': <function generate_remove_single_ensemble_member_constraint.<locals>.<lambda> at 0x"
        in repr(single_member_constraint)
    )


def test_generate_remove_single_ensemble_member_constraint_any_value():
    """Generate a constraint to remove a single ensemble member using chosen value."""
    single_member_constraint = (
        constraints.generate_remove_single_ensemble_member_constraint(ensemble_member=2)
    )
    assert (
        "Constraint(coord_values={'realization': <function generate_remove_single_ensemble_member_constraint.<locals>.<lambda> at 0x"
        in repr(single_member_constraint)
    )


def test_generate_realization_constraint():
    """Generate a constraint for a single realization."""
    single_member_constraint = constraints.generate_realization_constraint(
        ensemble_members=2
    )
    assert "Constraint(coord_values={'realization': (2,)})" in repr(
        single_member_constraint
    )


def test_generate_realization_constraint_multiple_realizations():
    """Generate a constraint for multiple realizations."""
    multi_member_constraint = constraints.generate_realization_constraint(
        ensemble_members=[2, 4, 6, 8]
    )
    assert "Constraint(coord_values={'realization': [2, 4, 6, 8]})" in repr(
        multi_member_constraint
    )


def test_generate_hour_constraint():
    """Generate hour constraint with hour_start."""
    hour_constraint = constraints.generate_hour_constraint(hour_start=12)
    expected_hour_constraint = "Constraint(coord_values={'hour': <function generate_hour_constraint.<locals>.<lambda> at"
    assert expected_hour_constraint in repr(hour_constraint)


def test_generate_hour_constraint_both_limits():
    """Generate hour constraint with hour_start and hour_end."""
    hour_constraint = constraints.generate_hour_constraint(hour_start=12, hour_end=15)
    expected_hour_constraint = "Constraint(coord_values={'hour': <function generate_hour_constraint.<locals>.<lambda> at"
    assert expected_hour_constraint in repr(hour_constraint)


def test_generate_hour_constraint_negative_values():
    """Generate hour constraint raises exception when arguments are negative."""
    with pytest.raises(ValueError):
        constraints.generate_hour_constraint(hour_start=-1)

    with pytest.raises(ValueError):
        constraints.generate_hour_constraint(hour_start=0, hour_end=-1)


def test_generate_hour_constraint_too_large_values():
    """Generate hour constraint raises exception when arguments are too big."""
    with pytest.raises(ValueError):
        constraints.generate_hour_constraint(hour_start=24)

    with pytest.raises(ValueError):
        constraints.generate_hour_constraint(hour_start=22, hour_end=24)


def test_combine_constraints():
    """Combine constraint."""
    stash_constraint = constraints.generate_stash_constraint("m01s03i236")
    var_constraint = constraints.generate_var_constraint("test")
    combined_constraint = constraints.combine_constraints(
        stash_constraint,
        a=var_constraint,
    )
    expected_combined_constraint = "ConstraintCombination(AttributeConstraint({'STASH': 'm01s03i236'}), Constraint(name='test'), <built-in function and_>)"
    assert repr(combined_constraint) == expected_combined_constraint
    var_constraint_2 = constraints.generate_var_constraint("test_2")
    combined_constraint = constraints.combine_constraints(
        stash_constraint,
        a=var_constraint,
        b=var_constraint_2,
    )
    expected_combined_constraint = "ConstraintCombination(ConstraintCombination(AttributeConstraint({'STASH': 'm01s03i236'}), Constraint(name='test'), <built-in function and_>), Constraint(name='test_2'), <built-in function and_>)"
    assert repr(combined_constraint) == expected_combined_constraint


def test_generate_attribute_constraint_value_none():
    """Attribute constraint with no value."""
    attr_constraint = constraints.generate_attribute_constraint(attribute="test")
    expected_attr_constraint = "Constraint(cube_func=<function generate_attribute_constraint.<locals>.<lambda> at "
    assert expected_attr_constraint in repr(attr_constraint)


def test_generate_attribute_constraint_with_value():
    """Attribute constraint with value."""
    attr_constraint = constraints.generate_attribute_constraint(
        attribute="test", value="2"
    )
    expected_attr_constraint = "AttributeConstraint({'test': '2'})"
    assert expected_attr_constraint in repr(attr_constraint)


def test_generate_var_constraint_multiple_names():
    """Test constraint works for multiple variable names."""
    # Create two cubes with different names
    cube1 = iris.cube.Cube(np.arange(5), long_name="temperature_long")
    cube1.var_name = "var_temperature"
    cube2 = iris.cube.Cube(np.arange(5), standard_name="wind_speed")
    # Third cube that should NOT match
    cube3 = iris.cube.Cube(np.arange(5), long_name="surface_pressure")
    # Generate constraint with multiple names
    constraint = constraints.generate_var_constraint(["var_temperature", "wind_speed"])
    # Apply constraint
    cubes = iris.cube.CubeList([cube1, cube2, cube3])
    result = cubes.extract(constraint)
    # Check correct cubes are selected
    result_names = [c.name() for c in result]

    assert cube1 in result
    assert cube2 in result
    assert cube3 not in result

    assert "temperature_long" in result_names
    assert "wind_speed" in result_names
    assert "surface_pressure" not in result_names
    # Should only return 2 cubes
    assert len(result) == 2


def test_generate_remove_single_level_constraint():
    """Tests constraint to remove default model_level_number of zero."""
    remove_level_constraint = constraints.generate_remove_single_level_constraint(
        coord="model_level_number"
    )
    expected_constraint = "Constraint(coord_values={'model_level_number': <function generate_remove_single_level_constraint.<locals>.<lambda> at"
    assert expected_constraint in repr(remove_level_constraint)


def test_generate_remove_single_level_constraint_non_default():
    """Tests constraint to remove default model_level_number of one."""
    remove_level_constraint = constraints.generate_remove_single_level_constraint(
        coord="model_level_number", level=1
    )
    expected_constraint = "Constraint(coord_values={'model_level_number': <function generate_remove_single_level_constraint.<locals>.<lambda> at"
    assert expected_constraint in repr(remove_level_constraint)