From 46ce0e2342919776c544a0b6aae964f6f7645a4b Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Mon, 22 Jan 2024 09:36:55 +0100 Subject: [PATCH 01/53] [optional modules] Add CLI option and liskov directives --- tools/src/icon4pytools/liskov/cli.py | 9 ++- .../liskov/codegen/integration/deserialise.py | 63 +++++++++++++++++++ .../liskov/codegen/integration/interface.py | 10 +++ .../src/icon4pytools/liskov/external/gt4py.py | 1 + .../src/icon4pytools/liskov/parsing/parse.py | 10 +++ .../icon4pytools/liskov/parsing/transform.py | 31 ++++++++- .../icon4pytools/liskov/parsing/validation.py | 31 +++++---- .../liskov/pipeline/collection.py | 4 +- 8 files changed, 144 insertions(+), 15 deletions(-) diff --git a/tools/src/icon4pytools/liskov/cli.py b/tools/src/icon4pytools/liskov/cli.py index dd6eb7876b..099695b42f 100644 --- a/tools/src/icon4pytools/liskov/cli.py +++ b/tools/src/icon4pytools/liskov/cli.py @@ -56,6 +56,12 @@ def main(ctx: click.Context) -> None: default=True, help="Adds fused or unfused stencils.", ) +@click.option( + "--enable-dsl-optional", + type=click.Choice(['advection', 'graupel', 'no']), + default='no', + help="Specify the DSL optional module to enable.", +) @click.argument( "input_path", type=click.Path(exists=True, dir_okay=False, resolve_path=True, path_type=pathlib.Path), @@ -70,10 +76,11 @@ def integrate( fused: bool, profile: bool, metadatagen: bool, + enable_dsl_optional: str, ) -> None: mode = "integration" iface = parse_fortran_file(input_path, output_path, mode) - iface_gt4py = process_stencils(iface, fused) + iface_gt4py = process_stencils(iface, fused, optional_modules_to_enable=enable_dsl_optional) run_code_generation( input_path, output_path, diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index fd3c9aecb5..09ffb92a4c 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -22,6 +22,7 @@ EndCreateData, EndDeleteData, EndFusedStencilData, + EndOptionalStencilData, EndIfData, EndProfileData, EndStencilData, @@ -32,6 +33,7 @@ StartCreateData, StartDeleteData, StartFusedStencilData, + StartOptionalStencilData, StartProfileData, StartStencilData, UnusedDirective, @@ -260,6 +262,25 @@ def __call__(self, parsed: ts.ParsedDict) -> list[EndFusedStencilData]: return deserialised +class EndOptionalStencilDataFactory(DataFactoryBase): + directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.EndOptionalStencil + dtype: Type[EndOptionalStencilData] = EndOptionalStencilData + + def __call__(self, parsed: ts.ParsedDict) -> list[EndOptionalStencilData]: + deserialised = [] + extracted = extract_directive(parsed["directives"], self.directive_cls) + for i, directive in enumerate(extracted): + named_args = parsed["content"]["EndOptionalStencil"][i] + stencil_name = _extract_stencil_name(named_args, directive) + deserialised.append( + self.dtype( + name=stencil_name, + startln=directive.startln, + ) + ) + return deserialised + + class StartStencilDataFactoryBase(DataFactoryBase): directive_cls: Type[ts.ParsedDirective] dtype: Type[StartStencilData | StartFusedStencilData] @@ -407,6 +428,46 @@ class StartFusedStencilDataFactory(StartStencilDataFactoryBase): dtype: Type[StartFusedStencilData] = StartFusedStencilData +class StartOptionalStencilDataFactory(StartStencilDataFactoryBase): + directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartOptionalStencil + dtype: Type[StartOptionalStencilData] = StartOptionalStencilData + + def create_stencil_data( + self, + parsed: ts.ParsedDict, + field_dimensions: dict[str, Any], + directives: Sequence[ts.ParsedDirective], + directive_cls: Type[ts.ParsedDirective], + dtype: Type[StartOptionalStencilData], + ) -> list[StartOptionalStencilData]: + """Create and return a list of StartOptionalStencilData objects from parsed directives.""" + deserialised = [] + for i, directive in enumerate(directives): + named_args = parsed["content"][directive_cls.__name__][i] + additional_attrs = self._pop_additional_attributes(dtype, named_args) + acc_present = string_to_bool(pop_item_from_dict(named_args, "accpresent", "true")) + stencil_name = _extract_stencil_name(named_args, directive) + bounds = self._make_bounds(named_args) + fields = self._make_fields(named_args, field_dimensions) + fields_w_tolerance = self._update_tolerances(named_args, fields) + + # Add optional_module to additional attributes + additional_attrs["optional_module"] = named_args.get("optional_module", "") + + deserialised.append( + dtype( + name=stencil_name, + fields=fields_w_tolerance, + bounds=bounds, + startln=directive.startln, + acc_present=acc_present, + **additional_attrs, + ) + ) + return deserialised + + + class StartStencilDataFactory(StartStencilDataFactoryBase): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartStencil dtype: Type[StartStencilData] = StartStencilData @@ -435,6 +496,8 @@ class IntegrationCodeDeserialiser(Deserialiser): "EndStencil": EndStencilDataFactory(), "StartFusedStencil": StartFusedStencilDataFactory(), "EndFusedStencil": EndFusedStencilDataFactory(), + "StartOptionalStencil": StartOptionalStencilDataFactory(), + "EndOptionalStencil": EndOptionalStencilDataFactory(), "StartDelete": StartDeleteDataFactory(), "EndDelete": EndDeleteDataFactory(), "EndIf": EndIfDataFactory(), diff --git a/tools/src/icon4pytools/liskov/codegen/integration/interface.py b/tools/src/icon4pytools/liskov/codegen/integration/interface.py index 93903cab78..ee5c60d133 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/interface.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/interface.py @@ -95,6 +95,10 @@ class StartStencilData(BaseStartStencilData): class StartFusedStencilData(BaseStartStencilData): ... +@dataclass +class StartOptionalStencilData(BaseStartStencilData): + optional_module: str + @dataclass class BaseEndStencilData(CodeGenInput): @@ -112,6 +116,10 @@ class EndStencilData(BaseEndStencilData): class EndFusedStencilData(BaseEndStencilData): ... +@dataclass +class EndOptionalStencilData(BaseEndStencilData): + ... + @dataclass class StartDeleteData(CodeGenInput): @@ -134,6 +142,8 @@ class IntegrationCodeInterface: EndStencil: Sequence[EndStencilData] StartFusedStencil: Sequence[StartFusedStencilData] EndFusedStencil: Sequence[EndFusedStencilData] + StartOptionalStencil: Sequence[StartOptionalStencilData] + EndOptionalStencil: Sequence[EndOptionalStencilData] StartDelete: Sequence[StartDeleteData] | UnusedDirective EndDelete: Sequence[EndDeleteData] | UnusedDirective Declare: Sequence[DeclareData] diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index 42b6363616..c3a24577a0 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -47,6 +47,7 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: self._set_in_out_field(self.parsed.StartStencil) self._set_in_out_field(self.parsed.StartFusedStencil) + self._set_in_out_field(self.parsed.StartOptionalStencil) return self.parsed diff --git a/tools/src/icon4pytools/liskov/parsing/parse.py b/tools/src/icon4pytools/liskov/parsing/parse.py index c4052fa56c..9425f3a1e3 100644 --- a/tools/src/icon4pytools/liskov/parsing/parse.py +++ b/tools/src/icon4pytools/liskov/parsing/parse.py @@ -192,6 +192,14 @@ class EndFusedStencil(WithArguments): pattern = "END FUSED STENCIL" +class StartOptionalStencil(WithArguments): + pattern = "START OPTIONAL STENCIL" + + +class EndOptionalStencil(WithArguments): + pattern = "END OPTIONAL STENCIL" + + class Declare(WithArguments): pattern = "DECLARE" @@ -237,6 +245,8 @@ class Insert(FreeForm): EndStencil, StartFusedStencil, EndFusedStencil, + StartOptionalStencil, + EndOptionalStencil, StartDelete, EndDelete, Imports, diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index f5ace9add4..a7a8add2ad 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -31,9 +31,10 @@ class StencilTransformer(Step): - def __init__(self, parsed: IntegrationCodeInterface, fused: bool) -> None: + def __init__(self, parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: str) -> None: self.parsed = parsed self.fused = fused + self.optional_modules_to_enable = optional_modules_to_enable def __call__(self, data: Any = None) -> IntegrationCodeInterface: """Transform stencils in the parse tree based on the 'fused' flag, transforming or removing as necessary. @@ -56,6 +57,12 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: self._remove_fused_stencils() self._remove_delete() + # Conditionally process optional stencils based on self.optional_modules_to_enable + if self.optional_modules_to_enable.lower() != "no": + self.process_optional_stencils_to_enable() + else: + self._remove_optional_stencils() + return self.parsed def _process_stencils_for_deletion(self) -> None: @@ -118,3 +125,25 @@ def _remove_fused_stencils(self) -> None: def _remove_delete(self) -> None: self.parsed.StartDelete = [] self.parsed.EndDelete = [] + + def _process_optional_stencils_to_enable(self): + # Check if optional_module_to_enable matches the optional_module in StartOptionalStencil + if self.parsed.StartOptionalStencil.optional_module == self.optional_modules_to_enable: + # Remove optional_module attribute from StartOptionalStencil + del self.parsed.StartOptionalStencil.optional_module + + # Add elements of StartOptionalStencil to StartStencil + self.parsed.StartStencil.__dict__.update(self.parsed.StartOptionalStencil.__dict__) + + # Optionally, you can remove the StartOptionalStencil object itself if needed + del self.parsed.StartOptionalStencil + + # Add any other logic you might need after processing optional stencils + + def _remove_optional_stencils(self): + pass + # Implement the logic to remove optional stencils when not enabled + # This method will be called when self.optional_modules_to_enable is "no" + self.parsed.StartOptionalStencil = [] + self.parsed.EndOptionalStencil = [] + # delete other liskov statements? diff --git a/tools/src/icon4pytools/liskov/parsing/validation.py b/tools/src/icon4pytools/liskov/parsing/validation.py index 434d29fe10..6adad4a9ac 100644 --- a/tools/src/icon4pytools/liskov/parsing/validation.py +++ b/tools/src/icon4pytools/liskov/parsing/validation.py @@ -122,6 +122,8 @@ def _validate_directive_uniqueness(self, directives: Sequence[ts.ParsedDirective parse.EndCreate, parse.StartStencil, parse.EndStencil, + parse.StartOptionalStencil, + parse.EndOptionalStencil, parse.EndIf, parse.EndProfile, parse.StartProfile, @@ -136,17 +138,24 @@ def _validate_directive_uniqueness(self, directives: Sequence[ts.ParsedDirective def _validate_required_directives(self, directives: Sequence[ts.ParsedDirective]) -> None: """Check that all required directives are used at least once.""" - expected = [ - parse.Declare, - parse.Imports, - parse.StartStencil, - parse.EndStencil, - ] - for expected_type in expected: - if not any([isinstance(d, expected_type) for d in directives]): - raise RequiredDirectivesError( - f"Error in {self.filepath}.\n Missing required directive of type {expected_type.pattern} in source." - ) + group_1 = (parse.StartStencil, parse.EndStencil) + group_2 = (parse.StartOptionalStencil, parse.EndOptionalStencil) + + if not ( + (any(isinstance(d, group_1[0]) for d in directives) and + any(isinstance(d, group_1[1]) for d in directives)) or + (any(isinstance(d, group_2[0]) for d in directives) and + any(isinstance(d, group_2[1]) for d in directives)) + ) or not ( + any(isinstance(d, parse.Declare) for d in directives) and + any(isinstance(d, parse.Imports) for d in directives) + ): + raise RequiredDirectivesError( + f"Error in {self.filepath}.\n Missing required directive(s) in source. " + f"At least one of {group_1[0].pattern}, {group_1[1].pattern} or " + f"{group_2[0].pattern}, {group_2[1].pattern} must be present, and " + f"{parse.Declare.pattern} and {parse.Imports.pattern} are required." + ) def _validate_stencil_directives(self, directives: Sequence[ts.ParsedDirective]) -> None: """Validate that the number of start and end stencil directives match in the input `directives`. diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index 4d458ad1cc..21d5b8d3b9 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -70,7 +70,7 @@ def parse_fortran_file( @linear_pipeline -def process_stencils(parsed: IntegrationCodeInterface, fused: bool) -> list[Step]: +def process_stencils(parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: str) -> list[Step]: """Execute a linear pipeline to transform stencils and produce either fused or unfused execution. This function takes an input `parsed` object of type `IntegrationCodeInterface` and a `fused` boolean flag. @@ -84,7 +84,7 @@ def process_stencils(parsed: IntegrationCodeInterface, fused: bool) -> list[Step Returns: The updated and transformed object with fields containing information from GT4Py stencils. """ - return [StencilTransformer(parsed, fused), UpdateFieldsWithGt4PyStencils(parsed)] + return [StencilTransformer(parsed, fused, optional_modules_to_enable), UpdateFieldsWithGt4PyStencils(parsed)] @linear_pipeline From 8ba115ef22aaed29937b57d007ef864034899501 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Mon, 22 Jan 2024 09:59:13 +0100 Subject: [PATCH 02/53] Update mo_advection_btraj_compute_o1_dsl --- .../mo_advection_traj_btraj_compute_o1_dsl.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py index 941cb99661..cacfdda7bf 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py @@ -33,6 +33,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl( dual_normal_cell_2: Field[[ECDim], wpfloat], p_dthalf: wpfloat, ) -> tuple[ + Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], vpfloat], @@ -41,6 +42,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl( lvn_pos = where(p_vn > 0.0, True, False) p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1])) + p_cell_rel_idx_dsl = where(lvn_pos, int32(0), int32(1)) p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) z_ntdistv_bary_1 = -( @@ -67,7 +69,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl( + z_ntdistv_bary_2 * dual_normal_cell_2(E2EC[1]), ) - return p_cell_idx, p_cell_blk, astype(p_distv_bary_1, vpfloat), astype(p_distv_bary_2, vpfloat) + return p_cell_idx, p_cell_rel_idx_dsl, p_cell_blk, astype(p_distv_bary_1, vpfloat), astype(p_distv_bary_2, vpfloat) @program(grid_type=GridType.UNSTRUCTURED) @@ -83,6 +85,7 @@ def mo_advection_traj_btraj_compute_o1_dsl( primal_normal_cell_2: Field[[ECDim], wpfloat], dual_normal_cell_2: Field[[ECDim], wpfloat], p_cell_idx: Field[[EdgeDim, KDim], int32], + p_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], p_cell_blk: Field[[EdgeDim, KDim], int32], p_distv_bary_1: Field[[EdgeDim, KDim], vpfloat], p_distv_bary_2: Field[[EdgeDim, KDim], vpfloat], @@ -104,7 +107,13 @@ def mo_advection_traj_btraj_compute_o1_dsl( primal_normal_cell_2, dual_normal_cell_2, p_dthalf, - out=(p_cell_idx, p_cell_blk, p_distv_bary_1, p_distv_bary_2), + out=( + p_cell_idx, + p_cell_rel_idx_dsl, + p_cell_blk, + p_distv_bary_1, + p_distv_bary_2, + ), domain={ EdgeDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end), From 8c240d1b237431be7464fde08131ff5c299e6c05 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Mon, 22 Jan 2024 11:08:22 +0100 Subject: [PATCH 03/53] Pre-commit run --- tools/src/icon4pytools/liskov/cli.py | 4 ++-- .../liskov/codegen/integration/deserialise.py | 3 +-- .../liskov/codegen/integration/interface.py | 2 ++ .../src/icon4pytools/liskov/parsing/transform.py | 9 ++++++--- .../icon4pytools/liskov/parsing/validation.py | 16 ++++++++++------ .../icon4pytools/liskov/pipeline/collection.py | 9 +++++++-- 6 files changed, 28 insertions(+), 15 deletions(-) diff --git a/tools/src/icon4pytools/liskov/cli.py b/tools/src/icon4pytools/liskov/cli.py index 099695b42f..03d9cf0a55 100644 --- a/tools/src/icon4pytools/liskov/cli.py +++ b/tools/src/icon4pytools/liskov/cli.py @@ -58,8 +58,8 @@ def main(ctx: click.Context) -> None: ) @click.option( "--enable-dsl-optional", - type=click.Choice(['advection', 'graupel', 'no']), - default='no', + type=click.Choice(["advection", "graupel", "no"]), + default="no", help="Specify the DSL optional module to enable.", ) @click.argument( diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index 09ffb92a4c..4e069bf99e 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -22,8 +22,8 @@ EndCreateData, EndDeleteData, EndFusedStencilData, - EndOptionalStencilData, EndIfData, + EndOptionalStencilData, EndProfileData, EndStencilData, FieldAssociationData, @@ -467,7 +467,6 @@ def create_stencil_data( return deserialised - class StartStencilDataFactory(StartStencilDataFactoryBase): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartStencil dtype: Type[StartStencilData] = StartStencilData diff --git a/tools/src/icon4pytools/liskov/codegen/integration/interface.py b/tools/src/icon4pytools/liskov/codegen/integration/interface.py index ee5c60d133..40511578c5 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/interface.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/interface.py @@ -95,6 +95,7 @@ class StartStencilData(BaseStartStencilData): class StartFusedStencilData(BaseStartStencilData): ... + @dataclass class StartOptionalStencilData(BaseStartStencilData): optional_module: str @@ -116,6 +117,7 @@ class EndStencilData(BaseEndStencilData): class EndFusedStencilData(BaseEndStencilData): ... + @dataclass class EndOptionalStencilData(BaseEndStencilData): ... diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index a7a8add2ad..b2ad7eb716 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -31,7 +31,9 @@ class StencilTransformer(Step): - def __init__(self, parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: str) -> None: + def __init__( + self, parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: str + ) -> None: self.parsed = parsed self.fused = fused self.optional_modules_to_enable = optional_modules_to_enable @@ -139,11 +141,12 @@ def _process_optional_stencils_to_enable(self): del self.parsed.StartOptionalStencil # Add any other logic you might need after processing optional stencils - + def _remove_optional_stencils(self): pass # Implement the logic to remove optional stencils when not enabled # This method will be called when self.optional_modules_to_enable is "no" self.parsed.StartOptionalStencil = [] self.parsed.EndOptionalStencil = [] - # delete other liskov statements? + + # delete other liskov statements? diff --git a/tools/src/icon4pytools/liskov/parsing/validation.py b/tools/src/icon4pytools/liskov/parsing/validation.py index 6adad4a9ac..cb86ab3b69 100644 --- a/tools/src/icon4pytools/liskov/parsing/validation.py +++ b/tools/src/icon4pytools/liskov/parsing/validation.py @@ -142,13 +142,17 @@ def _validate_required_directives(self, directives: Sequence[ts.ParsedDirective] group_2 = (parse.StartOptionalStencil, parse.EndOptionalStencil) if not ( - (any(isinstance(d, group_1[0]) for d in directives) and - any(isinstance(d, group_1[1]) for d in directives)) or - (any(isinstance(d, group_2[0]) for d in directives) and - any(isinstance(d, group_2[1]) for d in directives)) + ( + any(isinstance(d, group_1[0]) for d in directives) + and any(isinstance(d, group_1[1]) for d in directives) + ) + or ( + any(isinstance(d, group_2[0]) for d in directives) + and any(isinstance(d, group_2[1]) for d in directives) + ) ) or not ( - any(isinstance(d, parse.Declare) for d in directives) and - any(isinstance(d, parse.Imports) for d in directives) + any(isinstance(d, parse.Declare) for d in directives) + and any(isinstance(d, parse.Imports) for d in directives) ): raise RequiredDirectivesError( f"Error in {self.filepath}.\n Missing required directive(s) in source. " diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index 21d5b8d3b9..370690dc25 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -70,7 +70,9 @@ def parse_fortran_file( @linear_pipeline -def process_stencils(parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: str) -> list[Step]: +def process_stencils( + parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: str +) -> list[Step]: """Execute a linear pipeline to transform stencils and produce either fused or unfused execution. This function takes an input `parsed` object of type `IntegrationCodeInterface` and a `fused` boolean flag. @@ -84,7 +86,10 @@ def process_stencils(parsed: IntegrationCodeInterface, fused: bool, optional_mod Returns: The updated and transformed object with fields containing information from GT4Py stencils. """ - return [StencilTransformer(parsed, fused, optional_modules_to_enable), UpdateFieldsWithGt4PyStencils(parsed)] + return [ + StencilTransformer(parsed, fused, optional_modules_to_enable), + UpdateFieldsWithGt4PyStencils(parsed), + ] @linear_pipeline From ebad3c6dbe54e74a7de938cdaf58be5b6c95def9 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Wed, 24 Jan 2024 10:44:30 +0100 Subject: [PATCH 04/53] Remove START OPTIONAL STENCIL directives; only optional arg. --- .../liskov/codegen/integration/deserialise.py | 66 ++----------------- .../liskov/codegen/integration/interface.py | 13 +--- .../src/icon4pytools/liskov/external/gt4py.py | 1 - .../src/icon4pytools/liskov/parsing/parse.py | 10 --- .../icon4pytools/liskov/parsing/transform.py | 20 ++---- .../icon4pytools/liskov/parsing/validation.py | 35 ++++------ 6 files changed, 24 insertions(+), 121 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index 4e069bf99e..2807319203 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -23,7 +23,6 @@ EndDeleteData, EndFusedStencilData, EndIfData, - EndOptionalStencilData, EndProfileData, EndStencilData, FieldAssociationData, @@ -33,7 +32,6 @@ StartCreateData, StartDeleteData, StartFusedStencilData, - StartOptionalStencilData, StartProfileData, StartStencilData, UnusedDirective, @@ -262,25 +260,6 @@ def __call__(self, parsed: ts.ParsedDict) -> list[EndFusedStencilData]: return deserialised -class EndOptionalStencilDataFactory(DataFactoryBase): - directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.EndOptionalStencil - dtype: Type[EndOptionalStencilData] = EndOptionalStencilData - - def __call__(self, parsed: ts.ParsedDict) -> list[EndOptionalStencilData]: - deserialised = [] - extracted = extract_directive(parsed["directives"], self.directive_cls) - for i, directive in enumerate(extracted): - named_args = parsed["content"]["EndOptionalStencil"][i] - stencil_name = _extract_stencil_name(named_args, directive) - deserialised.append( - self.dtype( - name=stencil_name, - startln=directive.startln, - ) - ) - return deserialised - - class StartStencilDataFactoryBase(DataFactoryBase): directive_cls: Type[ts.ParsedDirective] dtype: Type[StartStencilData | StartFusedStencilData] @@ -313,6 +292,8 @@ def create_stencil_data( fields = self._make_fields(named_args, field_dimensions) fields_w_tolerance = self._update_tolerances(named_args, fields) + additional_attrs["optional_module"] = named_args.get("optional_module", "") + deserialised.append( dtype( name=stencil_name, @@ -326,7 +307,9 @@ def create_stencil_data( return deserialised def _pop_additional_attributes( - self, dtype: Type[StartStencilData | StartFusedStencilData], named_args: dict[str, Any] + self, + dtype: Type[StartStencilData | StartFusedStencilData], + named_args: dict[str, Any], ) -> dict: """Pop and return additional attributes specific to StartStencilData.""" additional_attrs = {} @@ -428,43 +411,6 @@ class StartFusedStencilDataFactory(StartStencilDataFactoryBase): dtype: Type[StartFusedStencilData] = StartFusedStencilData -class StartOptionalStencilDataFactory(StartStencilDataFactoryBase): - directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartOptionalStencil - dtype: Type[StartOptionalStencilData] = StartOptionalStencilData - - def create_stencil_data( - self, - parsed: ts.ParsedDict, - field_dimensions: dict[str, Any], - directives: Sequence[ts.ParsedDirective], - directive_cls: Type[ts.ParsedDirective], - dtype: Type[StartOptionalStencilData], - ) -> list[StartOptionalStencilData]: - """Create and return a list of StartOptionalStencilData objects from parsed directives.""" - deserialised = [] - for i, directive in enumerate(directives): - named_args = parsed["content"][directive_cls.__name__][i] - additional_attrs = self._pop_additional_attributes(dtype, named_args) - acc_present = string_to_bool(pop_item_from_dict(named_args, "accpresent", "true")) - stencil_name = _extract_stencil_name(named_args, directive) - bounds = self._make_bounds(named_args) - fields = self._make_fields(named_args, field_dimensions) - fields_w_tolerance = self._update_tolerances(named_args, fields) - - # Add optional_module to additional attributes - additional_attrs["optional_module"] = named_args.get("optional_module", "") - - deserialised.append( - dtype( - name=stencil_name, - fields=fields_w_tolerance, - bounds=bounds, - startln=directive.startln, - acc_present=acc_present, - **additional_attrs, - ) - ) - return deserialised class StartStencilDataFactory(StartStencilDataFactoryBase): @@ -495,8 +441,6 @@ class IntegrationCodeDeserialiser(Deserialiser): "EndStencil": EndStencilDataFactory(), "StartFusedStencil": StartFusedStencilDataFactory(), "EndFusedStencil": EndFusedStencilDataFactory(), - "StartOptionalStencil": StartOptionalStencilDataFactory(), - "EndOptionalStencil": EndOptionalStencilDataFactory(), "StartDelete": StartDeleteDataFactory(), "EndDelete": EndDeleteDataFactory(), "EndIf": EndIfDataFactory(), diff --git a/tools/src/icon4pytools/liskov/codegen/integration/interface.py b/tools/src/icon4pytools/liskov/codegen/integration/interface.py index 40511578c5..07cb43872e 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/interface.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/interface.py @@ -89,6 +89,7 @@ class BaseStartStencilData(CodeGenInput): class StartStencilData(BaseStartStencilData): mergecopy: Optional[bool] copies: Optional[bool] + optional_module: Optional[str] @dataclass @@ -96,11 +97,6 @@ class StartFusedStencilData(BaseStartStencilData): ... -@dataclass -class StartOptionalStencilData(BaseStartStencilData): - optional_module: str - - @dataclass class BaseEndStencilData(CodeGenInput): name: str @@ -118,11 +114,6 @@ class EndFusedStencilData(BaseEndStencilData): ... -@dataclass -class EndOptionalStencilData(BaseEndStencilData): - ... - - @dataclass class StartDeleteData(CodeGenInput): startln: int @@ -144,8 +135,6 @@ class IntegrationCodeInterface: EndStencil: Sequence[EndStencilData] StartFusedStencil: Sequence[StartFusedStencilData] EndFusedStencil: Sequence[EndFusedStencilData] - StartOptionalStencil: Sequence[StartOptionalStencilData] - EndOptionalStencil: Sequence[EndOptionalStencilData] StartDelete: Sequence[StartDeleteData] | UnusedDirective EndDelete: Sequence[EndDeleteData] | UnusedDirective Declare: Sequence[DeclareData] diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index c3a24577a0..42b6363616 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -47,7 +47,6 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: self._set_in_out_field(self.parsed.StartStencil) self._set_in_out_field(self.parsed.StartFusedStencil) - self._set_in_out_field(self.parsed.StartOptionalStencil) return self.parsed diff --git a/tools/src/icon4pytools/liskov/parsing/parse.py b/tools/src/icon4pytools/liskov/parsing/parse.py index 9425f3a1e3..c4052fa56c 100644 --- a/tools/src/icon4pytools/liskov/parsing/parse.py +++ b/tools/src/icon4pytools/liskov/parsing/parse.py @@ -192,14 +192,6 @@ class EndFusedStencil(WithArguments): pattern = "END FUSED STENCIL" -class StartOptionalStencil(WithArguments): - pattern = "START OPTIONAL STENCIL" - - -class EndOptionalStencil(WithArguments): - pattern = "END OPTIONAL STENCIL" - - class Declare(WithArguments): pattern = "DECLARE" @@ -245,8 +237,6 @@ class Insert(FreeForm): EndStencil, StartFusedStencil, EndFusedStencil, - StartOptionalStencil, - EndOptionalStencil, StartDelete, EndDelete, Imports, diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index b2ad7eb716..ccaf7f38e9 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -61,7 +61,7 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: # Conditionally process optional stencils based on self.optional_modules_to_enable if self.optional_modules_to_enable.lower() != "no": - self.process_optional_stencils_to_enable() + self._process_optional_stencils_to_enable() else: self._remove_optional_stencils() @@ -129,24 +129,18 @@ def _remove_delete(self) -> None: self.parsed.EndDelete = [] def _process_optional_stencils_to_enable(self): - # Check if optional_module_to_enable matches the optional_module in StartOptionalStencil - if self.parsed.StartOptionalStencil.optional_module == self.optional_modules_to_enable: - # Remove optional_module attribute from StartOptionalStencil - del self.parsed.StartOptionalStencil.optional_module + if self.parsed.StartStencil.optional_module == self.optional_modules_to_enable: + del self.parsed.StartStencil.optional_module - # Add elements of StartOptionalStencil to StartStencil - self.parsed.StartStencil.__dict__.update(self.parsed.StartOptionalStencil.__dict__) + self.parsed.StartStencil.__dict__.update(self.parsed.StartStencil.__dict__) - # Optionally, you can remove the StartOptionalStencil object itself if needed - del self.parsed.StartOptionalStencil - - # Add any other logic you might need after processing optional stencils + del self.parsed.StartStencil def _remove_optional_stencils(self): pass # Implement the logic to remove optional stencils when not enabled # This method will be called when self.optional_modules_to_enable is "no" - self.parsed.StartOptionalStencil = [] - self.parsed.EndOptionalStencil = [] + self.parsed.StartStencil = [] + self.parsed.EndStencil = [] # delete other liskov statements? diff --git a/tools/src/icon4pytools/liskov/parsing/validation.py b/tools/src/icon4pytools/liskov/parsing/validation.py index cb86ab3b69..434d29fe10 100644 --- a/tools/src/icon4pytools/liskov/parsing/validation.py +++ b/tools/src/icon4pytools/liskov/parsing/validation.py @@ -122,8 +122,6 @@ def _validate_directive_uniqueness(self, directives: Sequence[ts.ParsedDirective parse.EndCreate, parse.StartStencil, parse.EndStencil, - parse.StartOptionalStencil, - parse.EndOptionalStencil, parse.EndIf, parse.EndProfile, parse.StartProfile, @@ -138,28 +136,17 @@ def _validate_directive_uniqueness(self, directives: Sequence[ts.ParsedDirective def _validate_required_directives(self, directives: Sequence[ts.ParsedDirective]) -> None: """Check that all required directives are used at least once.""" - group_1 = (parse.StartStencil, parse.EndStencil) - group_2 = (parse.StartOptionalStencil, parse.EndOptionalStencil) - - if not ( - ( - any(isinstance(d, group_1[0]) for d in directives) - and any(isinstance(d, group_1[1]) for d in directives) - ) - or ( - any(isinstance(d, group_2[0]) for d in directives) - and any(isinstance(d, group_2[1]) for d in directives) - ) - ) or not ( - any(isinstance(d, parse.Declare) for d in directives) - and any(isinstance(d, parse.Imports) for d in directives) - ): - raise RequiredDirectivesError( - f"Error in {self.filepath}.\n Missing required directive(s) in source. " - f"At least one of {group_1[0].pattern}, {group_1[1].pattern} or " - f"{group_2[0].pattern}, {group_2[1].pattern} must be present, and " - f"{parse.Declare.pattern} and {parse.Imports.pattern} are required." - ) + expected = [ + parse.Declare, + parse.Imports, + parse.StartStencil, + parse.EndStencil, + ] + for expected_type in expected: + if not any([isinstance(d, expected_type) for d in directives]): + raise RequiredDirectivesError( + f"Error in {self.filepath}.\n Missing required directive of type {expected_type.pattern} in source." + ) def _validate_stencil_directives(self, directives: Sequence[ts.ParsedDirective]) -> None: """Validate that the number of start and end stencil directives match in the input `directives`. From e40db7f8375dd6275cd32918193a1104ef273f71 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Wed, 24 Jan 2024 10:46:43 +0100 Subject: [PATCH 05/53] Pre-commit fix --- .../advection/mo_advection_traj_btraj_compute_o1_dsl.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py index cacfdda7bf..f35f0d4564 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py @@ -69,7 +69,13 @@ def _mo_advection_traj_btraj_compute_o1_dsl( + z_ntdistv_bary_2 * dual_normal_cell_2(E2EC[1]), ) - return p_cell_idx, p_cell_rel_idx_dsl, p_cell_blk, astype(p_distv_bary_1, vpfloat), astype(p_distv_bary_2, vpfloat) + return ( + p_cell_idx, + p_cell_rel_idx_dsl, + p_cell_blk, + astype(p_distv_bary_1, vpfloat), + astype(p_distv_bary_2, vpfloat), + ) @program(grid_type=GridType.UNSTRUCTURED) From 4bd6faa1d2f347cb60394bef2017277b88c834c4 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Tue, 30 Jan 2024 13:54:28 +0100 Subject: [PATCH 06/53] Add advection unit tests in tox --- model/tox.ini | 2 ++ 1 file changed, 2 insertions(+) diff --git a/model/tox.ini b/model/tox.ini index 7aed4fb051..992b99c586 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -30,11 +30,13 @@ allowlist_externals = commands = pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/diffusion/tests/diffusion_stencil_tests --benchmark-skip -n auto {posargs} pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/dycore/tests/dycore_stencil_tests --benchmark-skip -n auto {posargs} + pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/advection/tests/stencil_tests --benchmark-skip -n auto {posargs} [testenv:run_benchmarks] commands = pytest -s -m "not slow_tests" atmosphere/diffusion/tests/diffusion_stencil_tests --benchmark-only {posargs} pytest -s -m "not slow_tests" atmosphere/dycore/tests/dycore_stencil_tests --benchmark-only {posargs} + pytest -s -m "not slow_tests" atmosphere/advection/tests/stencil_tests --benchmark-only {posargs} [testenv:dev] From 59ea47c83695e611e87ab99261d6e49e7033aab2 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Tue, 30 Jan 2024 13:58:40 +0100 Subject: [PATCH 07/53] Fix optional_module attribute in liskov --- .../liskov/codegen/integration/deserialise.py | 6 ++-- .../icon4pytools/liskov/parsing/transform.py | 29 +++++++------------ 2 files changed, 13 insertions(+), 22 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index 2807319203..289142a6a7 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -291,9 +291,6 @@ def create_stencil_data( bounds = self._make_bounds(named_args) fields = self._make_fields(named_args, field_dimensions) fields_w_tolerance = self._update_tolerances(named_args, fields) - - additional_attrs["optional_module"] = named_args.get("optional_module", "") - deserialised.append( dtype( name=stencil_name, @@ -316,7 +313,8 @@ def _pop_additional_attributes( if dtype == StartStencilData: mergecopy = string_to_bool(pop_item_from_dict(named_args, "mergecopy", "false")) copies = string_to_bool(pop_item_from_dict(named_args, "copies", "true")) - additional_attrs = {"mergecopy": mergecopy, "copies": copies} + optional_module = pop_item_from_dict(named_args, "optional_module", "no") + additional_attrs = {"mergecopy": mergecopy, "copies": copies, "optional_module": optional_module} return additional_attrs @staticmethod diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index ccaf7f38e9..bc380c3313 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -59,11 +59,7 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: self._remove_fused_stencils() self._remove_delete() - # Conditionally process optional stencils based on self.optional_modules_to_enable - if self.optional_modules_to_enable.lower() != "no": - self._process_optional_stencils_to_enable() - else: - self._remove_optional_stencils() + self._process_optional_stencils_to_enable() return self.parsed @@ -129,18 +125,15 @@ def _remove_delete(self) -> None: self.parsed.EndDelete = [] def _process_optional_stencils_to_enable(self): - if self.parsed.StartStencil.optional_module == self.optional_modules_to_enable: - del self.parsed.StartStencil.optional_module - - self.parsed.StartStencil.__dict__.update(self.parsed.StartStencil.__dict__) - - del self.parsed.StartStencil + stencils_to_remove = [] + modified_stencil_list = [] - def _remove_optional_stencils(self): - pass - # Implement the logic to remove optional stencils when not enabled - # This method will be called when self.optional_modules_to_enable is "no" - self.parsed.StartStencil = [] - self.parsed.EndStencil = [] + for start_single, end_single in zip( + self.parsed.StartStencil, self.parsed.EndStencil, strict=True + ): + optional_module_value = getattr(start_single, 'optional_module', None) - # delete other liskov statements? + if optional_module_value not in self.optional_modules_to_enable: + stencils_to_remove += [start_single, end_single] + self._remove_stencils(stencils_to_remove) + \ No newline at end of file From d85eaf05012880f28d0798aa1942a1c3076b9755 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Wed, 31 Jan 2024 00:03:02 +0100 Subject: [PATCH 08/53] Change field name in set_zero_c_k --- .../src/icon4py/model/atmosphere/advection/set_zero_c_k.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py index a6b50d09ef..74a88179d4 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py @@ -24,5 +24,5 @@ def _set_zero_c_k() -> Field[[CellDim, KDim], float]: @program(grid_type=GridType.UNSTRUCTURED) -def set_zero_c_k(field: Field[[CellDim, KDim], float]): - _set_zero_c_k(out=field) +def set_zero_c_k(z_cfl: Field[[CellDim, KDim], float]): + _set_zero_c_k(out=z_cfl) From bd4c1b7a338dcdfc0e66afa8c465400b979f8e15 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Wed, 31 Jan 2024 00:03:58 +0100 Subject: [PATCH 09/53] Adapt optional feature in liskov --- .../liskov/codegen/integration/deserialise.py | 4 +--- .../liskov/codegen/integration/generate.py | 1 + tools/src/icon4pytools/liskov/parsing/transform.py | 13 ++++--------- 3 files changed, 6 insertions(+), 12 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index 289142a6a7..c042041b49 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -313,7 +313,7 @@ def _pop_additional_attributes( if dtype == StartStencilData: mergecopy = string_to_bool(pop_item_from_dict(named_args, "mergecopy", "false")) copies = string_to_bool(pop_item_from_dict(named_args, "copies", "true")) - optional_module = pop_item_from_dict(named_args, "optional_module", "no") + optional_module = pop_item_from_dict(named_args, "optional_module", "None") additional_attrs = {"mergecopy": mergecopy, "copies": copies, "optional_module": optional_module} return additional_attrs @@ -409,8 +409,6 @@ class StartFusedStencilDataFactory(StartStencilDataFactoryBase): dtype: Type[StartFusedStencilData] = StartFusedStencilData - - class StartStencilDataFactory(StartStencilDataFactoryBase): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartStencil dtype: Type[StartStencilData] = StartStencilData diff --git a/tools/src/icon4pytools/liskov/codegen/integration/generate.py b/tools/src/icon4pytools/liskov/codegen/integration/generate.py index 0603e226f4..be0a44d448 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/generate.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/generate.py @@ -143,6 +143,7 @@ def _generate_start_stencil(self) -> None: acc_present=stencil.acc_present, mergecopy=stencil.mergecopy, copies=stencil.copies, + optional_module=stencil.optional_module, ) i += 2 diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index bc380c3313..cc85c9c059 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -126,14 +126,9 @@ def _remove_delete(self) -> None: def _process_optional_stencils_to_enable(self): stencils_to_remove = [] - modified_stencil_list = [] - for start_single, end_single in zip( - self.parsed.StartStencil, self.parsed.EndStencil, strict=True - ): - optional_module_value = getattr(start_single, 'optional_module', None) + for start_stencil, end_stencil in zip(self.parsed.StartStencil, self.parsed.EndStencil): + if start_stencil.optional_module not in self.optional_modules_to_enable and start_stencil.optional_module != "None": + stencils_to_remove += [start_stencil, end_stencil] - if optional_module_value not in self.optional_modules_to_enable: - stencils_to_remove += [start_single, end_single] - self._remove_stencils(stencils_to_remove) - \ No newline at end of file + self._remove_stencils(stencils_to_remove) From 9fb954d6ee3c70d04ec22a877503e610c3fa7a16 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Wed, 31 Jan 2024 00:05:41 +0100 Subject: [PATCH 10/53] Pre-commit fixes --- .../icon4pytools/liskov/codegen/integration/deserialise.py | 6 +++++- tools/src/icon4pytools/liskov/parsing/transform.py | 5 ++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index c042041b49..4b19efcd0b 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -314,7 +314,11 @@ def _pop_additional_attributes( mergecopy = string_to_bool(pop_item_from_dict(named_args, "mergecopy", "false")) copies = string_to_bool(pop_item_from_dict(named_args, "copies", "true")) optional_module = pop_item_from_dict(named_args, "optional_module", "None") - additional_attrs = {"mergecopy": mergecopy, "copies": copies, "optional_module": optional_module} + additional_attrs = { + "mergecopy": mergecopy, + "copies": copies, + "optional_module": optional_module, + } return additional_attrs @staticmethod diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index cc85c9c059..1e858734bd 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -128,7 +128,10 @@ def _process_optional_stencils_to_enable(self): stencils_to_remove = [] for start_stencil, end_stencil in zip(self.parsed.StartStencil, self.parsed.EndStencil): - if start_stencil.optional_module not in self.optional_modules_to_enable and start_stencil.optional_module != "None": + if ( + start_stencil.optional_module not in self.optional_modules_to_enable + and start_stencil.optional_module != "None" + ): stencils_to_remove += [start_stencil, end_stencil] self._remove_stencils(stencils_to_remove) From bf03135ea445daa61ffd736ca20ae9ab4ceff488 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 7 Feb 2024 10:54:18 +0100 Subject: [PATCH 11/53] Use list comprehension again --- tools/src/icon4pytools/liskov/parsing/parse.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tools/src/icon4pytools/liskov/parsing/parse.py b/tools/src/icon4pytools/liskov/parsing/parse.py index c4052fa56c..aef87af306 100644 --- a/tools/src/icon4pytools/liskov/parsing/parse.py +++ b/tools/src/icon4pytools/liskov/parsing/parse.py @@ -83,7 +83,10 @@ def _determine_type( def _preprocess(self, directives: Sequence[ts.ParsedDirective]) -> Sequence[ts.ParsedDirective]: """Preprocess the directives by removing unnecessary characters and formatting the directive strings.""" - return [d.__class__(self._clean_string(d.string), d.startln, d.endln) for d in directives] + return [ + d.__class__(self._clean_string(d.string, d.type_name), d.startln, d.endln) + for d in directives + ] def _run_validation_passes(self, preprocessed: Sequence[ts.ParsedDirective]) -> None: """Run validation passes on the directives.""" From c13d3015ee28657df3abee70775d2ce1bf52bce0 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 7 Feb 2024 10:44:47 +0100 Subject: [PATCH 12/53] Special case string cleaning for DSL INSERT --- .../src/icon4pytools/liskov/parsing/parse.py | 24 ++++++++++++------- tools/tests/liskov/fortran_samples.py | 2 +- tools/tests/liskov/test_parser.py | 2 +- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/tools/src/icon4pytools/liskov/parsing/parse.py b/tools/src/icon4pytools/liskov/parsing/parse.py index aef87af306..be6716f5c2 100644 --- a/tools/src/icon4pytools/liskov/parsing/parse.py +++ b/tools/src/icon4pytools/liskov/parsing/parse.py @@ -25,8 +25,6 @@ from icon4pytools.liskov.pipeline.definition import Step -REPLACE_CHARS = [ts.DIRECTIVE_IDENT, "&", "\n"] - logger = setup_logger(__name__) @@ -83,10 +81,13 @@ def _determine_type( def _preprocess(self, directives: Sequence[ts.ParsedDirective]) -> Sequence[ts.ParsedDirective]: """Preprocess the directives by removing unnecessary characters and formatting the directive strings.""" - return [ - d.__class__(self._clean_string(d.string, d.type_name), d.startln, d.endln) - for d in directives - ] + preprocessed = [] + for d in directives: + new_directive = d.__class__( + self._clean_string(d.string, d.type_name), d.startln, d.endln + ) + preprocessed.append(new_directive) + return preprocessed def _run_validation_passes(self, preprocessed: Sequence[ts.ParsedDirective]) -> None: """Run validation passes on the directives.""" @@ -94,9 +95,16 @@ def _run_validation_passes(self, preprocessed: Sequence[ts.ParsedDirective]) -> validator(self.input_filepath).validate(preprocessed) @staticmethod - def _clean_string(string: str) -> str: + def _clean_string(string: str, type_name: str) -> str: """Remove leading or trailing whitespaces, and words from the REPLACE_CHARS list.""" - return " ".join([c for c in string.strip().split() if c not in REPLACE_CHARS]) + replace_chars = [ts.DIRECTIVE_IDENT] + + # DSL INSERT Statements should be inserted verbatim meaning no string cleaning + # other than the directive identifier. + if type_name != "Insert": + replace_chars += ["&", "\n"] + + return " ".join([c for c in string.strip().split() if c not in replace_chars]) @staticmethod def _parse(directives: Sequence[ts.ParsedDirective]) -> ts.ParsedContent: diff --git a/tools/tests/liskov/fortran_samples.py b/tools/tests/liskov/fortran_samples.py index da2b056f3b..0f00e48a79 100644 --- a/tools/tests/liskov/fortran_samples.py +++ b/tools/tests/liskov/fortran_samples.py @@ -28,7 +28,7 @@ SINGLE_STENCIL = """\ !$DSL IMPORTS() - + !$DSL INSERT( & p_cell_rel_idx_dsl, & ) !$DSL START CREATE() !$DSL DECLARE(vn=nproma,p_patch%nlev,p_patch%nblks_e; suffix=dsl) diff --git a/tools/tests/liskov/test_parser.py b/tools/tests/liskov/test_parser.py index 0b7094d56c..8828e48e5a 100644 --- a/tools/tests/liskov/test_parser.py +++ b/tools/tests/liskov/test_parser.py @@ -83,7 +83,7 @@ def test_parse_single_directive(directive, string, startln, endln, expected_cont @mark.parametrize( "stencil, num_directives, num_content", [ - (SINGLE_STENCIL, 9, 8), + (SINGLE_STENCIL, 10, 9), (SINGLE_STENCIL_WITH_COMMENTS, 9, 8), (MULTIPLE_STENCILS, 11, 7), (SINGLE_FUSED, 9, 7), From 1d416f979c1a1b4d75556782bacbec6a1829d54c Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 7 Feb 2024 16:48:09 +0100 Subject: [PATCH 13/53] improvements --- tools/README.md | 5 ++++- tools/docs/ICON_Liskov_integration_style_guide.md | 2 +- tools/src/icon4pytools/liskov/codegen/shared/generate.py | 7 ++++++- 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/tools/README.md b/tools/README.md index c77c90dcb9..417ba43a49 100644 --- a/tools/README.md +++ b/tools/README.md @@ -254,7 +254,10 @@ call wrap_run_calculate_diagnostic_quantities_for_turbulence( & #### `!$DSL INSERT()` -This directive allows the user to generate any text that is placed between the parentheses. This is useful for situations where custom code generation is necessary. +This directive allows the user to generate any text that is placed between the parentheses. +This is useful for situations where custom code generation is necessary. +Note that, the `INSERT`` statement is verbatim, such that there is no filtering or fortran formatting. +Also, that line continatuion with `&` is not provided for `INSERT` statements. #### `!$DSL START PROFILE()` diff --git a/tools/docs/ICON_Liskov_integration_style_guide.md b/tools/docs/ICON_Liskov_integration_style_guide.md index cd5f573563..1e4406d095 100644 --- a/tools/docs/ICON_Liskov_integration_style_guide.md +++ b/tools/docs/ICON_Liskov_integration_style_guide.md @@ -16,7 +16,7 @@ The code should look as clean and concise as possible. Also it should be similar - `!$DSL END STENCIL` as close as possible to the end of the original stencil section, no empty line before, one empty line after. - `!$DSL START FUSED STENCIL` should be placed before the `!$DSL START STENCIL` of the first stencil in the fused batch, one empty line before, one empty line after. - `!$DSL END FUSED STENCIL` should be placed after the `!$DSL END STENCIL` of the last stencil in the fused batch, one empty line before, one empty line after. -- `!$DSL INSERT` one empty line before and after, unless the inserted code is part of an `ACC` data region. +- `!$DSL INSERT` one empty line before and after, unless the inserted code is part of an `ACC` data region, or a function call. - `!$DSL START CREATE` after the `!$ACC CREATE` block, no empty line before and one empty line after. ## Example diff --git a/tools/src/icon4pytools/liskov/codegen/shared/generate.py b/tools/src/icon4pytools/liskov/codegen/shared/generate.py index 9b6f14d3d8..ac6f28176e 100644 --- a/tools/src/icon4pytools/liskov/codegen/shared/generate.py +++ b/tools/src/icon4pytools/liskov/codegen/shared/generate.py @@ -20,6 +20,7 @@ from icon4pytools.icon4pygen.bindings.utils import format_fortran_code from icon4pytools.liskov.codegen.shared.types import CodeGenInput, GeneratedCode from icon4pytools.liskov.pipeline.definition import Step +from icon4pytools.liskov.codegen.integration.template import InsertStatement class CodeGenerator(Step): @@ -53,7 +54,11 @@ def _generate_fortran_code( """ parent = parent_node(**kwargs) source = code_generator.apply(parent) - formatted_source = format_fortran_code(source) + # DSL INSERT Statements should be inserted verbatim meaning no fortran formatting + if parent_node is InsertStatement: + formatted_source = source + else: + formatted_source = format_fortran_code(source) return formatted_source def _generate( From df96eeb7495bf16a9a34ce39bb7a3e0e610fd0cd Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 15 Feb 2024 10:22:02 +0100 Subject: [PATCH 14/53] WIP --- .../advection/vert_adv_stencil_01.py | 42 ++++++++++++++++--- 1 file changed, 36 insertions(+), 6 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py index 44f933b18f..c74f486977 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py @@ -12,11 +12,32 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where + from icon4py.model.common.dimension import CellDim, KDim, Koff +@field_operator +def _vert_adv_stencil_01a( + tracer_now: Field[[CellDim, KDim], float], + rhodz_now: Field[[CellDim, KDim], float], + p_mflx_tracer_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + rhodz_new: Field[[CellDim, KDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + # tracer_new = ( + # tracer_now * rhodz_now + # + p_dtime * (p_mflx_tracer_v(Koff[1]) * deepatmo_divzl - p_mflx_tracer_v * deepatmo_divzu) + # ) / rhodz_new + # tracer_new = tracer_now * rhodz_now / rhodz_now + tracer_new = tracer_now + + return tracer_new + + @field_operator def _vert_adv_stencil_01( tracer_now: Field[[CellDim, KDim], float], @@ -25,12 +46,15 @@ def _vert_adv_stencil_01( deepatmo_divzl: Field[[KDim], float], deepatmo_divzu: Field[[KDim], float], rhodz_new: Field[[CellDim, KDim], float], + k: Field[[KDim], int32], p_dtime: float, + ivadv_tracer: int32, + iadv_slev_jt: int32, ) -> Field[[CellDim, KDim], float]: - tracer_new = ( - tracer_now * rhodz_now - + p_dtime * (p_mflx_tracer_v(Koff[1]) * deepatmo_divzl - p_mflx_tracer_v * deepatmo_divzu) - ) / rhodz_new + k = broadcast(k, (CellDim, KDim)) + + # tracer_new = (where((iadv_slev_jt<=k), _vert_adv_stencil_01a(tracer_now, rhodz_now, p_mflx_tracer_v, deepatmo_divzl, deepatmo_divzu, rhodz_new, p_dtime, ), tracer_now*2.) if (ivadv_tracer!=int32(0)) else tracer_now ) + tracer_new = (where((iadv_slev_jt<=k), tracer_now*1.4, tracer_now) if (ivadv_tracer!=int32(0)) else tracer_now ) return tracer_new @@ -43,8 +67,11 @@ def vert_adv_stencil_01( deepatmo_divzl: Field[[KDim], float], deepatmo_divzu: Field[[KDim], float], rhodz_new: Field[[CellDim, KDim], float], - tracer_new: Field[[CellDim, KDim], float], + k: Field[[KDim], int32], p_dtime: float, + ivadv_tracer: int32, + iadv_slev_jt: int32, + tracer_new: Field[[CellDim, KDim], float], ): _vert_adv_stencil_01( tracer_now, @@ -53,6 +80,9 @@ def vert_adv_stencil_01( deepatmo_divzl, deepatmo_divzu, rhodz_new, + k, p_dtime, + ivadv_tracer, + iadv_slev_jt, out=tracer_new, ) From 93df4cd96ee17ff468725c88c19bb27495e259df Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 15 Feb 2024 16:02:42 +0100 Subject: [PATCH 15/53] clean up --- .../atmosphere/advection/vert_adv_stencil_01.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py index c74f486977..bc5229d09a 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py @@ -28,12 +28,10 @@ def _vert_adv_stencil_01a( rhodz_new: Field[[CellDim, KDim], float], p_dtime: float, ) -> Field[[CellDim, KDim], float]: - # tracer_new = ( - # tracer_now * rhodz_now - # + p_dtime * (p_mflx_tracer_v(Koff[1]) * deepatmo_divzl - p_mflx_tracer_v * deepatmo_divzu) - # ) / rhodz_new - # tracer_new = tracer_now * rhodz_now / rhodz_now - tracer_new = tracer_now + tracer_new = ( + tracer_now * rhodz_now + + p_dtime * (p_mflx_tracer_v(Koff[1]) * deepatmo_divzl - p_mflx_tracer_v * deepatmo_divzu) + ) / rhodz_new return tracer_new @@ -53,8 +51,7 @@ def _vert_adv_stencil_01( ) -> Field[[CellDim, KDim], float]: k = broadcast(k, (CellDim, KDim)) - # tracer_new = (where((iadv_slev_jt<=k), _vert_adv_stencil_01a(tracer_now, rhodz_now, p_mflx_tracer_v, deepatmo_divzl, deepatmo_divzu, rhodz_new, p_dtime, ), tracer_now*2.) if (ivadv_tracer!=int32(0)) else tracer_now ) - tracer_new = (where((iadv_slev_jt<=k), tracer_now*1.4, tracer_now) if (ivadv_tracer!=int32(0)) else tracer_now ) + tracer_new = (where((iadv_slev_jt<=k), _vert_adv_stencil_01a(tracer_now, rhodz_now, p_mflx_tracer_v, deepatmo_divzl, deepatmo_divzu, rhodz_new, p_dtime, ), tracer_now*2.) if (ivadv_tracer!=int32(0)) else tracer_now ) return tracer_new From b5b7cc29ce85edbf37268038cf636223e48e10e7 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 15 Feb 2024 16:50:20 +0100 Subject: [PATCH 16/53] fix output name --- .../model/atmosphere/advection/hor_adv_stencil_01.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py index 59ad061a9d..cd8069e55e 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py @@ -27,14 +27,14 @@ def _hor_adv_stencil_01( geofac_div: Field[[CEDim], float], p_dtime: float, ) -> Field[[CellDim, KDim], float]: - tracer_new = ( + tracer_new_hor = ( tracer_now * rhodz_now - p_dtime * deepatmo_divh * neighbor_sum(p_mflx_tracer_h(C2E) * geofac_div(C2CE), axis=C2EDim) ) / rhodz_new - return tracer_new + return tracer_new_hor @program @@ -45,7 +45,7 @@ def hor_adv_stencil_01( rhodz_now: Field[[CellDim, KDim], float], rhodz_new: Field[[CellDim, KDim], float], geofac_div: Field[[CEDim], float], - tracer_new: Field[[CellDim, KDim], float], + tracer_new_hor: Field[[CellDim, KDim], float], p_dtime: float, ): _hor_adv_stencil_01( @@ -56,5 +56,5 @@ def hor_adv_stencil_01( rhodz_new, geofac_div, p_dtime, - out=tracer_new, + out=tracer_new_hor, ) From 6703bdc1c3f6e8e16736b6793e5982a83dbc3a13 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Fri, 16 Feb 2024 15:08:41 +0100 Subject: [PATCH 17/53] fix + unittest --- .../advection/vert_adv_stencil_01.py | 2 +- .../stencil_tests/test_vert_adv_stencil_01.py | 29 ++++++++++++++----- 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py index bc5229d09a..42db3e110a 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py @@ -51,7 +51,7 @@ def _vert_adv_stencil_01( ) -> Field[[CellDim, KDim], float]: k = broadcast(k, (CellDim, KDim)) - tracer_new = (where((iadv_slev_jt<=k), _vert_adv_stencil_01a(tracer_now, rhodz_now, p_mflx_tracer_v, deepatmo_divzl, deepatmo_divzu, rhodz_new, p_dtime, ), tracer_now*2.) if (ivadv_tracer!=int32(0)) else tracer_now ) + tracer_new = (where((iadv_slev_jt<=k), _vert_adv_stencil_01a(tracer_now, rhodz_now, p_mflx_tracer_v, deepatmo_divzl, deepatmo_divzu, rhodz_new, p_dtime, ), tracer_now) if (ivadv_tracer!=int32(0)) else tracer_now ) return tracer_new diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py index 598d496e3c..5857e496bf 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py @@ -13,10 +13,12 @@ import numpy as np import pytest +from gt4py.next import as_field +from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.atmosphere.advection.vert_adv_stencil_01 import vert_adv_stencil_01 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field +from icon4py.model.common.test_utils.helpers import StencilTest, _shape, random_field, zero_field class TestVertAdvStencil01(StencilTest): @@ -32,14 +34,21 @@ def reference( deepatmo_divzl: np.array, deepatmo_divzu: np.array, rhodz_new: np.array, - p_dtime, + k: np.array, + ivadv_tracer: int32, + iadv_slev_jt: int32, + p_dtime: np.float64, **kwargs, ) -> np.array: - tracer_new = ( - tracer_now * rhodz_now - + p_dtime - * (p_mflx_tracer_v[:, 1:] * deepatmo_divzl - p_mflx_tracer_v[:, :-1] * deepatmo_divzu) - ) / rhodz_new + if (ivadv_tracer!=int32(0)): + tracer_new = np.where( (iadv_slev_jt<=k), ( + tracer_now * rhodz_now + + p_dtime + * (p_mflx_tracer_v[:, 1:] * deepatmo_divzl - p_mflx_tracer_v[:, :-1] * deepatmo_divzu) + ) / rhodz_new, tracer_now) + else: + tracer_new = tracer_now + return dict(tracer_new=tracer_new) @pytest.fixture @@ -50,7 +59,10 @@ def input_data(self, grid): deepatmo_divzl = random_field(grid, KDim) deepatmo_divzu = random_field(grid, KDim) rhodz_new = random_field(grid, CellDim, KDim) + k = as_field((KDim,), np.arange(0, _shape(grid, KDim)[0], dtype=int32)) p_dtime = np.float64(5.0) + ivadv_tracer= int32(1) + iadv_slev_jt= int32(4) tracer_new = zero_field(grid, CellDim, KDim) return dict( tracer_now=tracer_now, @@ -59,6 +71,9 @@ def input_data(self, grid): deepatmo_divzl=deepatmo_divzl, deepatmo_divzu=deepatmo_divzu, rhodz_new=rhodz_new, + k=k, p_dtime=p_dtime, + ivadv_tracer= ivadv_tracer, + iadv_slev_jt= iadv_slev_jt, tracer_new=tracer_new, ) From 3733fb9df9e3a761520f90f8ce248da46865b478 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Fri, 16 Feb 2024 15:11:14 +0100 Subject: [PATCH 18/53] fix style --- .../advection/vert_adv_stencil_01.py | 21 +++++++++++-- .../stencil_tests/test_vert_adv_stencil_01.py | 30 ++++++++++++------- .../liskov/codegen/shared/generate.py | 1 - 3 files changed, 37 insertions(+), 15 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py index 42db3e110a..dfe82b7064 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py @@ -14,7 +14,6 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where - from icon4py.model.common.dimension import CellDim, KDim, Koff @@ -50,8 +49,24 @@ def _vert_adv_stencil_01( iadv_slev_jt: int32, ) -> Field[[CellDim, KDim], float]: k = broadcast(k, (CellDim, KDim)) - - tracer_new = (where((iadv_slev_jt<=k), _vert_adv_stencil_01a(tracer_now, rhodz_now, p_mflx_tracer_v, deepatmo_divzl, deepatmo_divzu, rhodz_new, p_dtime, ), tracer_now) if (ivadv_tracer!=int32(0)) else tracer_now ) + + tracer_new = ( + where( + (iadv_slev_jt <= k), + _vert_adv_stencil_01a( + tracer_now, + rhodz_now, + p_mflx_tracer_v, + deepatmo_divzl, + deepatmo_divzu, + rhodz_new, + p_dtime, + ), + tracer_now, + ) + if (ivadv_tracer != int32(0)) + else tracer_now + ) return tracer_new diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py index 5857e496bf..14e64e37cd 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py @@ -40,14 +40,22 @@ def reference( p_dtime: np.float64, **kwargs, ) -> np.array: - if (ivadv_tracer!=int32(0)): - tracer_new = np.where( (iadv_slev_jt<=k), ( - tracer_now * rhodz_now - + p_dtime - * (p_mflx_tracer_v[:, 1:] * deepatmo_divzl - p_mflx_tracer_v[:, :-1] * deepatmo_divzu) - ) / rhodz_new, tracer_now) + if ivadv_tracer != int32(0): + tracer_new = np.where( + (iadv_slev_jt <= k), + ( + tracer_now * rhodz_now + + p_dtime + * ( + p_mflx_tracer_v[:, 1:] * deepatmo_divzl + - p_mflx_tracer_v[:, :-1] * deepatmo_divzu + ) + ) + / rhodz_new, + tracer_now, + ) else: - tracer_new = tracer_now + tracer_new = tracer_now return dict(tracer_new=tracer_new) @@ -61,8 +69,8 @@ def input_data(self, grid): rhodz_new = random_field(grid, CellDim, KDim) k = as_field((KDim,), np.arange(0, _shape(grid, KDim)[0], dtype=int32)) p_dtime = np.float64(5.0) - ivadv_tracer= int32(1) - iadv_slev_jt= int32(4) + ivadv_tracer = int32(1) + iadv_slev_jt = int32(4) tracer_new = zero_field(grid, CellDim, KDim) return dict( tracer_now=tracer_now, @@ -73,7 +81,7 @@ def input_data(self, grid): rhodz_new=rhodz_new, k=k, p_dtime=p_dtime, - ivadv_tracer= ivadv_tracer, - iadv_slev_jt= iadv_slev_jt, + ivadv_tracer=ivadv_tracer, + iadv_slev_jt=iadv_slev_jt, tracer_new=tracer_new, ) diff --git a/tools/src/icon4pytools/liskov/codegen/shared/generate.py b/tools/src/icon4pytools/liskov/codegen/shared/generate.py index 2417376fa5..84c5d42907 100644 --- a/tools/src/icon4pytools/liskov/codegen/shared/generate.py +++ b/tools/src/icon4pytools/liskov/codegen/shared/generate.py @@ -21,7 +21,6 @@ from icon4pytools.liskov.codegen.integration.template import InsertStatement from icon4pytools.liskov.codegen.shared.types import CodeGenInput, GeneratedCode from icon4pytools.liskov.pipeline.definition import Step -from icon4pytools.liskov.codegen.integration.template import InsertStatement class CodeGenerator(Step): From 395ebcd1c48fae5970ff4ffbc6a3f48491bde485 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 19 Feb 2024 13:28:12 +0100 Subject: [PATCH 19/53] polish liskov --- .../icon4pytools/liskov/parsing/transform.py | 51 ++++++++++++------ .../liskov/pipeline/collection.py | 5 +- tools/tests/liskov/test_transform.py | 53 +++++++++++++++---- 3 files changed, 83 insertions(+), 26 deletions(-) diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index 1e858734bd..0f9b3c8e9c 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -29,14 +29,20 @@ logger = setup_logger(__name__) +def _remove_stencils(parsed: IntegrationCodeInterface, stencils_to_remove: list[CodeGenInput]) -> None: + attributes_to_modify = ["StartStencil", "EndStencil"] -class StencilTransformer(Step): + for attr_name in attributes_to_modify: + current_stencil_list = getattr(parsed, attr_name) + modified_stencil_list = [_ for _ in current_stencil_list if _ not in stencils_to_remove] + setattr(parsed, attr_name, modified_stencil_list) + +class FusedStencilTransformer(Step): def __init__( - self, parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: str + self, parsed: IntegrationCodeInterface, fused: bool ) -> None: self.parsed = parsed self.fused = fused - self.optional_modules_to_enable = optional_modules_to_enable def __call__(self, data: Any = None) -> IntegrationCodeInterface: """Transform stencils in the parse tree based on the 'fused' flag, transforming or removing as necessary. @@ -59,8 +65,6 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: self._remove_fused_stencils() self._remove_delete() - self._process_optional_stencils_to_enable() - return self.parsed def _process_stencils_for_deletion(self) -> None: @@ -76,7 +80,7 @@ def _process_stencils_for_deletion(self) -> None: self._create_delete_directives(start_single, end_single) stencils_to_remove += [start_single, end_single] - self._remove_stencils(stencils_to_remove) + _remove_stencils(self.parsed, stencils_to_remove) def _stencil_is_removable( self, @@ -108,14 +112,6 @@ def _create_delete_directives( directive.append(cls(startln=param.startln)) setattr(self.parsed, attr, directive) - def _remove_stencils(self, stencils_to_remove: list[CodeGenInput]) -> None: - attributes_to_modify = ["StartStencil", "EndStencil"] - - for attr_name in attributes_to_modify: - current_stencil_list = getattr(self.parsed, attr_name) - modified_stencil_list = [_ for _ in current_stencil_list if _ not in stencils_to_remove] - setattr(self.parsed, attr_name, modified_stencil_list) - def _remove_fused_stencils(self) -> None: self.parsed.StartFusedStencil = [] self.parsed.EndFusedStencil = [] @@ -124,6 +120,31 @@ def _remove_delete(self) -> None: self.parsed.StartDelete = [] self.parsed.EndDelete = [] + +class OptionalModulesTransformer(Step): + def __init__( + self, parsed: IntegrationCodeInterface, optional_modules_to_enable: str + ) -> None: + self.parsed = parsed + self.optional_modules_to_enable = optional_modules_to_enable + + def __call__(self, data: Any = None) -> IntegrationCodeInterface: + """Transform stencils in the parse tree based on the 'optional_modules_to_enable' flag, transforming or removing as necessary. + + This method processes stencils present in the 'parsed' object according to the 'optional_modules_to_enable' + flag. If 'optional_modules_to_enable' is a str, it identifies and processes stencils that are eligible for + deletion. If stencil is enabled with. + + Args: + data (Any): Optional data to be passed. Default is None. + + Returns: + IntegrationCodeInterface: The interface object along with any transformations applied. + """ + self._process_optional_stencils_to_enable() + + return self.parsed + def _process_optional_stencils_to_enable(self): stencils_to_remove = [] @@ -134,4 +155,4 @@ def _process_optional_stencils_to_enable(self): ): stencils_to_remove += [start_stencil, end_stencil] - self._remove_stencils(stencils_to_remove) + _remove_stencils(self.parsed, stencils_to_remove) \ No newline at end of file diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index 370690dc25..7a95b885d9 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -22,7 +22,7 @@ from icon4pytools.liskov.external.gt4py import UpdateFieldsWithGt4PyStencils from icon4pytools.liskov.parsing.parse import DirectivesParser from icon4pytools.liskov.parsing.scan import DirectivesScanner -from icon4pytools.liskov.parsing.transform import StencilTransformer +from icon4pytools.liskov.parsing.transform import FusedStencilTransformer, OptionalModulesTransformer from icon4pytools.liskov.pipeline.definition import Step, linear_pipeline @@ -87,7 +87,8 @@ def process_stencils( The updated and transformed object with fields containing information from GT4Py stencils. """ return [ - StencilTransformer(parsed, fused, optional_modules_to_enable), + FusedStencilTransformer(parsed, fused), + OptionalModulesTransformer(parsed, optional_modules_to_enable), UpdateFieldsWithGt4PyStencils(parsed), ] diff --git a/tools/tests/liskov/test_transform.py b/tools/tests/liskov/test_transform.py index 4c2a60454c..f45333a40e 100644 --- a/tools/tests/liskov/test_transform.py +++ b/tools/tests/liskov/test_transform.py @@ -33,7 +33,7 @@ StartProfileData, StartStencilData, ) -from icon4pytools.liskov.parsing.transform import StencilTransformer +from icon4pytools.liskov.parsing.transform import FusedStencilTransformer, OptionalModulesTransformer @pytest.fixture @@ -94,6 +94,7 @@ def integration_code_interface(): acc_present=False, mergecopy=False, copies=True, + optional_module="None", ) end_stencil_data1 = EndStencilData( name="stencil1", startln=3, noendif=False, noprofile=False, noaccenddata=False @@ -126,6 +127,7 @@ def integration_code_interface(): acc_present=False, mergecopy=False, copies=True, + optional_module='advection', ) end_stencil_data2 = EndStencilData( name="stencil2", startln=6, noendif=False, noprofile=False, noaccenddata=False @@ -165,20 +167,27 @@ def integration_code_interface(): @pytest.fixture -def stencil_transform_fused(integration_code_interface): - return StencilTransformer(integration_code_interface, fused=True) +def fused_stencil_transform_fused(integration_code_interface): + return FusedStencilTransformer(integration_code_interface, fused=True) +@pytest.fixture +def fused_stencil_transform_unfused(integration_code_interface): + return FusedStencilTransformer(integration_code_interface, fused=False) + +@pytest.fixture +def optional_modules_transform_enabled(integration_code_interface): + return OptionalModulesTransformer(integration_code_interface, optional_modules_to_enable='advection') @pytest.fixture -def stencil_transform_unfused(integration_code_interface): - return StencilTransformer(integration_code_interface, fused=False) +def optional_modules_transform_disabled(integration_code_interface): + return OptionalModulesTransformer(integration_code_interface, optional_modules_to_enable='no') def test_transform_fused( - stencil_transform_fused, + fused_stencil_transform_fused, ): # Check that the transformed interface is as expected - transformed = stencil_transform_fused() + transformed = fused_stencil_transform_fused() assert len(transformed.StartFusedStencil) == 1 assert len(transformed.EndFusedStencil) == 1 assert len(transformed.StartStencil) == 1 @@ -188,10 +197,10 @@ def test_transform_fused( def test_transform_unfused( - stencil_transform_unfused, + fused_stencil_transform_unfused, ): # Check that the transformed interface is as expected - transformed = stencil_transform_unfused() + transformed = fused_stencil_transform_unfused() assert not transformed.StartFusedStencil assert not transformed.EndFusedStencil @@ -199,3 +208,29 @@ def test_transform_unfused( assert len(transformed.EndStencil) == 2 assert not transformed.StartDelete assert not transformed.EndDelete + +def test_transform_optional_enabled( + optional_modules_transform_enabled, +): + # Check that the transformed interface is as expected + transformed = optional_modules_transform_enabled() + assert len(transformed.StartFusedStencil) == 1 + assert len(transformed.EndFusedStencil) == 1 + assert len(transformed.StartStencil) == 2 + assert len(transformed.EndStencil) == 2 + assert len(transformed.StartDelete) == 1 + assert len(transformed.EndDelete) == 1 + + +def test_transform_optional_disabled( + optional_modules_transform_disabled, +): + # Check that the transformed interface is as expected + transformed = optional_modules_transform_disabled() + + assert len(transformed.StartFusedStencil) == 1 + assert len(transformed.EndFusedStencil) == 1 + assert len(transformed.StartStencil) == 1 + assert len(transformed.EndStencil) == 1 + assert len(transformed.StartDelete) == 1 + assert len(transformed.EndDelete) == 1 From 9f20290ba9898e3300cad6f2698c5e8483ba55b9 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 19 Feb 2024 13:46:18 +0100 Subject: [PATCH 20/53] fix style --- .../icon4pytools/liskov/parsing/transform.py | 16 ++++++++-------- .../icon4pytools/liskov/pipeline/collection.py | 5 ++++- tools/tests/liskov/test_transform.py | 17 +++++++++++++---- 3 files changed, 25 insertions(+), 13 deletions(-) diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index 0f9b3c8e9c..234010cac8 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -29,7 +29,10 @@ logger = setup_logger(__name__) -def _remove_stencils(parsed: IntegrationCodeInterface, stencils_to_remove: list[CodeGenInput]) -> None: + +def _remove_stencils( + parsed: IntegrationCodeInterface, stencils_to_remove: list[CodeGenInput] +) -> None: attributes_to_modify = ["StartStencil", "EndStencil"] for attr_name in attributes_to_modify: @@ -37,10 +40,9 @@ def _remove_stencils(parsed: IntegrationCodeInterface, stencils_to_remove: list[ modified_stencil_list = [_ for _ in current_stencil_list if _ not in stencils_to_remove] setattr(parsed, attr_name, modified_stencil_list) + class FusedStencilTransformer(Step): - def __init__( - self, parsed: IntegrationCodeInterface, fused: bool - ) -> None: + def __init__(self, parsed: IntegrationCodeInterface, fused: bool) -> None: self.parsed = parsed self.fused = fused @@ -122,9 +124,7 @@ def _remove_delete(self) -> None: class OptionalModulesTransformer(Step): - def __init__( - self, parsed: IntegrationCodeInterface, optional_modules_to_enable: str - ) -> None: + def __init__(self, parsed: IntegrationCodeInterface, optional_modules_to_enable: str) -> None: self.parsed = parsed self.optional_modules_to_enable = optional_modules_to_enable @@ -155,4 +155,4 @@ def _process_optional_stencils_to_enable(self): ): stencils_to_remove += [start_stencil, end_stencil] - _remove_stencils(self.parsed, stencils_to_remove) \ No newline at end of file + _remove_stencils(self.parsed, stencils_to_remove) diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index 7a95b885d9..719575ed1f 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -22,7 +22,10 @@ from icon4pytools.liskov.external.gt4py import UpdateFieldsWithGt4PyStencils from icon4pytools.liskov.parsing.parse import DirectivesParser from icon4pytools.liskov.parsing.scan import DirectivesScanner -from icon4pytools.liskov.parsing.transform import FusedStencilTransformer, OptionalModulesTransformer +from icon4pytools.liskov.parsing.transform import ( + FusedStencilTransformer, + OptionalModulesTransformer, +) from icon4pytools.liskov.pipeline.definition import Step, linear_pipeline diff --git a/tools/tests/liskov/test_transform.py b/tools/tests/liskov/test_transform.py index f45333a40e..6a1898f3eb 100644 --- a/tools/tests/liskov/test_transform.py +++ b/tools/tests/liskov/test_transform.py @@ -33,7 +33,10 @@ StartProfileData, StartStencilData, ) -from icon4pytools.liskov.parsing.transform import FusedStencilTransformer, OptionalModulesTransformer +from icon4pytools.liskov.parsing.transform import ( + FusedStencilTransformer, + OptionalModulesTransformer, +) @pytest.fixture @@ -127,7 +130,7 @@ def integration_code_interface(): acc_present=False, mergecopy=False, copies=True, - optional_module='advection', + optional_module="advection", ) end_stencil_data2 = EndStencilData( name="stencil2", startln=6, noendif=False, noprofile=False, noaccenddata=False @@ -170,17 +173,22 @@ def integration_code_interface(): def fused_stencil_transform_fused(integration_code_interface): return FusedStencilTransformer(integration_code_interface, fused=True) + @pytest.fixture def fused_stencil_transform_unfused(integration_code_interface): return FusedStencilTransformer(integration_code_interface, fused=False) + @pytest.fixture def optional_modules_transform_enabled(integration_code_interface): - return OptionalModulesTransformer(integration_code_interface, optional_modules_to_enable='advection') + return OptionalModulesTransformer( + integration_code_interface, optional_modules_to_enable="advection" + ) + @pytest.fixture def optional_modules_transform_disabled(integration_code_interface): - return OptionalModulesTransformer(integration_code_interface, optional_modules_to_enable='no') + return OptionalModulesTransformer(integration_code_interface, optional_modules_to_enable="no") def test_transform_fused( @@ -209,6 +217,7 @@ def test_transform_unfused( assert not transformed.StartDelete assert not transformed.EndDelete + def test_transform_optional_enabled( optional_modules_transform_enabled, ): From 7791db76389ed3fad01a06f520cee1be89668385 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 19 Feb 2024 13:48:28 +0100 Subject: [PATCH 21/53] add to Readme --- tools/README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tools/README.md b/tools/README.md index b774421c0b..d50085c6c7 100644 --- a/tools/README.md +++ b/tools/README.md @@ -148,6 +148,8 @@ In addition, other optional keyword arguments are the following: - `copies`: Takes a boolean string input, and controls whether before field copies should be made or not. If set to False only the `#ifdef __DSL_VERIFY` directive is generated. Defaults to true.

+- `optional_module`: Takes a boolean string input, and controls whether stencils is part of an optional module. Defaults to "None".

+ #### `!$DSL END STENCIL()` This directive denotes the end of a stencil. The required argument is `name`, which must match the name of the preceding `START STENCIL` directive. From b7760a290a51a40ea7da11e953916df0d420beaf Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 19 Feb 2024 16:27:47 +0100 Subject: [PATCH 22/53] fix test --- tools/tests/liskov/test_external.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tools/tests/liskov/test_external.py b/tools/tests/liskov/test_external.py index 55f2e5f3b2..5b38206f56 100644 --- a/tools/tests/liskov/test_external.py +++ b/tools/tests/liskov/test_external.py @@ -77,6 +77,7 @@ def test_stencil_collector_invalid_member(): acc_present=False, mergecopy=False, copies=True, + optional_module="None" ) ], Imports=None, From d989c668d0019d9853ddba8a697803e6814a4672 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 19 Feb 2024 16:33:56 +0100 Subject: [PATCH 23/53] fix test --- tools/tests/liskov/test_generation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tools/tests/liskov/test_generation.py b/tools/tests/liskov/test_generation.py index 846aff6edb..66bc71b2b3 100644 --- a/tools/tests/liskov/test_generation.py +++ b/tools/tests/liskov/test_generation.py @@ -76,6 +76,7 @@ def integration_code_interface(): acc_present=False, mergecopy=False, copies=True, + optional_module="None" ) end_stencil_data = EndStencilData( name="stencil1", startln=3, noendif=False, noprofile=False, noaccenddata=False From c69afa45b68e69ad865dca48a7d9c95402e55605 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 19 Feb 2024 16:36:43 +0100 Subject: [PATCH 24/53] fix style --- tools/tests/liskov/test_external.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/tests/liskov/test_external.py b/tools/tests/liskov/test_external.py index 5b38206f56..454d01e428 100644 --- a/tools/tests/liskov/test_external.py +++ b/tools/tests/liskov/test_external.py @@ -77,7 +77,7 @@ def test_stencil_collector_invalid_member(): acc_present=False, mergecopy=False, copies=True, - optional_module="None" + optional_module="None", ) ], Imports=None, From e273bbed0ef261f454fe78ac0d11eb9fb72333f2 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 19 Feb 2024 16:39:28 +0100 Subject: [PATCH 25/53] fix style --- tools/tests/liskov/test_generation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/tests/liskov/test_generation.py b/tools/tests/liskov/test_generation.py index 66bc71b2b3..73ccc977f7 100644 --- a/tools/tests/liskov/test_generation.py +++ b/tools/tests/liskov/test_generation.py @@ -76,7 +76,7 @@ def integration_code_interface(): acc_present=False, mergecopy=False, copies=True, - optional_module="None" + optional_module="None", ) end_stencil_data = EndStencilData( name="stencil1", startln=3, noendif=False, noprofile=False, noaccenddata=False From c85bd8a808d5496ffb6b49cd251b3d38e6cc0b9f Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 20 Feb 2024 11:25:00 +0100 Subject: [PATCH 26/53] fix test --- .../src/icon4py/model/atmosphere/advection/set_zero_c_k.py | 4 ++-- .../test_mo_advection_traj_btraj_compute_o1_dsl.py | 5 +++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py index 74a88179d4..a6b50d09ef 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py @@ -24,5 +24,5 @@ def _set_zero_c_k() -> Field[[CellDim, KDim], float]: @program(grid_type=GridType.UNSTRUCTURED) -def set_zero_c_k(z_cfl: Field[[CellDim, KDim], float]): - _set_zero_c_k(out=z_cfl) +def set_zero_c_k(field: Field[[CellDim, KDim], float]): + _set_zero_c_k(out=field) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py index 05ca5d197a..cab4d4f9ca 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py @@ -69,8 +69,10 @@ def reference( dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) + p_cell_rel_idx_dsl = np.where(lvn_pos, int32(0), int32(1)) p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) + z_ntdistv_bary_1 = -( p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) @@ -96,6 +98,7 @@ def reference( return dict( p_cell_idx=p_cell_idx, + p_cell_rel_idx_dsl=p_cell_rel_idx_dsl, p_cell_blk=p_cell_blk, p_distv_bary_1=p_distv_bary_1, p_distv_bary_2=p_distv_bary_2, @@ -122,6 +125,7 @@ def input_data(self, grid): dual_normal_cell_2 = random_field(grid, EdgeDim, E2CDim) dual_normal_cell_2_new = as_1D_sparse_field(dual_normal_cell_2, ECDim) p_cell_idx = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) + p_cell_rel_idx_dsl = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) p_cell_blk = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) p_distv_bary_1 = random_field(grid, EdgeDim, KDim) p_distv_bary_2 = random_field(grid, EdgeDim, KDim) @@ -139,6 +143,7 @@ def input_data(self, grid): primal_normal_cell_2=primal_normal_cell_2_new, dual_normal_cell_2=dual_normal_cell_2_new, p_cell_idx=p_cell_idx, + p_cell_rel_idx_dsl=p_cell_rel_idx_dsl, p_cell_blk=p_cell_blk, p_distv_bary_1=p_distv_bary_1, p_distv_bary_2=p_distv_bary_2, From 967c38d031c6afb2391d90dfd0bc58ec66d5bb9a Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 20 Feb 2024 11:29:37 +0100 Subject: [PATCH 27/53] fix test --- .../tests/stencil_tests/test_hor_adv_stencil_01.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py index 51fb24ffb9..b517a76c3c 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py @@ -27,7 +27,7 @@ class TestHorAdvStencil01(StencilTest): PROGRAM = hor_adv_stencil_01 - OUTPUTS = ("tracer_new",) + OUTPUTS = ("tracer_new_hor",) @staticmethod def reference( @@ -43,13 +43,13 @@ def reference( ) -> np.array: geofac_div = reshape(geofac_div, grid.connectivities[C2EDim].shape) geofac_div = np.expand_dims(geofac_div, axis=-1) - tracer_new = ( + tracer_new_hor = ( tracer_now * rhodz_now - p_dtime * deepatmo_divh * np.sum(p_mflx_tracer_h[grid.connectivities[C2EDim]] * geofac_div, axis=1) ) / rhodz_new - return dict(tracer_new=tracer_new) + return dict(tracer_new_hor=tracer_new_hor) @pytest.fixture def input_data(self, grid): @@ -61,7 +61,7 @@ def input_data(self, grid): geofac_div = random_field(grid, CellDim, C2EDim) geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) p_dtime = np.float64(5.0) - tracer_new = zero_field(grid, CellDim, KDim) + tracer_new_hor = zero_field(grid, CellDim, KDim) return dict( p_mflx_tracer_h=p_mflx_tracer_h, deepatmo_divh=deepatmo_divh, @@ -70,5 +70,5 @@ def input_data(self, grid): rhodz_new=rhodz_new, geofac_div=geofac_div_new, p_dtime=p_dtime, - tracer_new=tracer_new, + tracer_new_hor=tracer_new_hor, ) From eae09f14d5a3f58ff36a8b316da397ddd8fb7776 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 20 Feb 2024 11:30:32 +0100 Subject: [PATCH 28/53] fix style --- .../stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py | 1 - 1 file changed, 1 deletion(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py index cab4d4f9ca..0be7c47b80 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py @@ -72,7 +72,6 @@ def reference( p_cell_rel_idx_dsl = np.where(lvn_pos, int32(0), int32(1)) p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) - z_ntdistv_bary_1 = -( p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) From 45ff5a31574936ac84000b13118262c3b804bf45 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 20 Feb 2024 14:02:00 +0100 Subject: [PATCH 29/53] addresse comments --- tools/src/icon4pytools/liskov/cli.py | 7 ++++--- .../liskov/codegen/integration/deserialise.py | 12 ++++++++---- .../codegen/serialisation/deserialise.py | 1 + .../icon4pytools/liskov/parsing/transform.py | 15 +++++++-------- .../icon4pytools/liskov/pipeline/collection.py | 18 ++++++++++++------ 5 files changed, 32 insertions(+), 21 deletions(-) diff --git a/tools/src/icon4pytools/liskov/cli.py b/tools/src/icon4pytools/liskov/cli.py index 03d9cf0a55..4eae9a2b42 100644 --- a/tools/src/icon4pytools/liskov/cli.py +++ b/tools/src/icon4pytools/liskov/cli.py @@ -58,8 +58,9 @@ def main(ctx: click.Context) -> None: ) @click.option( "--enable-dsl-optional", - type=click.Choice(["advection", "graupel", "no"]), - default="no", + type=click.Choice(["advection", "graupel", False]), + default=[False], + multiple=True, help="Specify the DSL optional module to enable.", ) @click.argument( @@ -76,7 +77,7 @@ def integrate( fused: bool, profile: bool, metadatagen: bool, - enable_dsl_optional: str, + enable_dsl_optional: list, ) -> None: mode = "integration" iface = parse_fortran_file(input_path, output_path, mode) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index 4b19efcd0b..bba6c1fe8d 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -53,6 +53,10 @@ TOLERANCE_ARGS = ["abs_tol", "rel_tol"] DEFAULT_DECLARE_IDENT_TYPE = "REAL(wp)" DEFAULT_DECLARE_SUFFIX = "before" +DEFAULT_STARTSTENCIL_ACC_PRESENT = "true" +DEFAULT_STARTSTENCIL_MERGECOPY = "false" +DEFAULT_STARTSTENCIL_COPIES = "true" +DEFAULT_STARTSTENCIL_OPTIONAL_MODULE = "None" logger = setup_logger(__name__) @@ -286,7 +290,7 @@ def create_stencil_data( for i, directive in enumerate(directives): named_args = parsed["content"][directive_cls.__name__][i] additional_attrs = self._pop_additional_attributes(dtype, named_args) - acc_present = string_to_bool(pop_item_from_dict(named_args, "accpresent", "true")) + acc_present = string_to_bool(pop_item_from_dict(named_args, "accpresent", DEFAULT_STARTSTENCIL_ACC_PRESENT)) stencil_name = _extract_stencil_name(named_args, directive) bounds = self._make_bounds(named_args) fields = self._make_fields(named_args, field_dimensions) @@ -311,9 +315,9 @@ def _pop_additional_attributes( """Pop and return additional attributes specific to StartStencilData.""" additional_attrs = {} if dtype == StartStencilData: - mergecopy = string_to_bool(pop_item_from_dict(named_args, "mergecopy", "false")) - copies = string_to_bool(pop_item_from_dict(named_args, "copies", "true")) - optional_module = pop_item_from_dict(named_args, "optional_module", "None") + mergecopy = string_to_bool(pop_item_from_dict(named_args, "mergecopy", DEFAULT_STARTSTENCIL_MERGECOPY)) + copies = string_to_bool(pop_item_from_dict(named_args, "copies", DEFAULT_STARTSTENCIL_COPIES)) + optional_module = pop_item_from_dict(named_args, "optional_module", DEFAULT_STARTSTENCIL_OPTIONAL_MODULE) additional_attrs = { "mergecopy": mergecopy, "copies": copies, diff --git a/tools/src/icon4pytools/liskov/codegen/serialisation/deserialise.py b/tools/src/icon4pytools/liskov/codegen/serialisation/deserialise.py index 3e711ccaf5..43b5547909 100644 --- a/tools/src/icon4pytools/liskov/codegen/serialisation/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/serialisation/deserialise.py @@ -39,6 +39,7 @@ "accpresent", "mergecopy", "copies", + "optional_module", "horizontal_lower", "horizontal_upper", "vertical_lower", diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index 234010cac8..ef46952bde 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -124,7 +124,7 @@ def _remove_delete(self) -> None: class OptionalModulesTransformer(Step): - def __init__(self, parsed: IntegrationCodeInterface, optional_modules_to_enable: str) -> None: + def __init__(self, parsed: IntegrationCodeInterface, optional_modules_to_enable: list) -> None: self.parsed = parsed self.optional_modules_to_enable = optional_modules_to_enable @@ -141,18 +141,17 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: Returns: IntegrationCodeInterface: The interface object along with any transformations applied. """ - self._process_optional_stencils_to_enable() + self._enable_optional_stencils() return self.parsed - def _process_optional_stencils_to_enable(self): + def _enable_optional_stencils(self): stencils_to_remove = [] for start_stencil, end_stencil in zip(self.parsed.StartStencil, self.parsed.EndStencil): - if ( - start_stencil.optional_module not in self.optional_modules_to_enable - and start_stencil.optional_module != "None" - ): - stencils_to_remove += [start_stencil, end_stencil] + if start_stencil.optional_module == "None": + pass + elif start_stencil.optional_module not in self.optional_modules_to_enable: + stencils_to_remove += [start_stencil, end_stencil] _remove_stencils(self.parsed, stencils_to_remove) diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index 719575ed1f..dba0936487 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -74,7 +74,7 @@ def parse_fortran_file( @linear_pipeline def process_stencils( - parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: str + parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: list ) -> list[Step]: """Execute a linear pipeline to transform stencils and produce either fused or unfused execution. @@ -89,11 +89,17 @@ def process_stencils( Returns: The updated and transformed object with fields containing information from GT4Py stencils. """ - return [ - FusedStencilTransformer(parsed, fused), - OptionalModulesTransformer(parsed, optional_modules_to_enable), - UpdateFieldsWithGt4PyStencils(parsed), - ] + if optional_modules_to_enable==[False]: + return [ + FusedStencilTransformer(parsed, fused), + UpdateFieldsWithGt4PyStencils(parsed), + ] + else: + return [ + FusedStencilTransformer(parsed, fused), + OptionalModulesTransformer(parsed, optional_modules_to_enable), + UpdateFieldsWithGt4PyStencils(parsed), + ] @linear_pipeline From 9b63734eff3c3cd674396c937c87fba47640b8fd Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 20 Feb 2024 14:05:58 +0100 Subject: [PATCH 30/53] fix style --- .../liskov/codegen/integration/deserialise.py | 16 +++++++++++---- .../icon4pytools/liskov/parsing/transform.py | 8 ++++---- .../liskov/pipeline/collection.py | 20 +++++++++---------- 3 files changed, 26 insertions(+), 18 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index bba6c1fe8d..d4c1ce1280 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -290,7 +290,9 @@ def create_stencil_data( for i, directive in enumerate(directives): named_args = parsed["content"][directive_cls.__name__][i] additional_attrs = self._pop_additional_attributes(dtype, named_args) - acc_present = string_to_bool(pop_item_from_dict(named_args, "accpresent", DEFAULT_STARTSTENCIL_ACC_PRESENT)) + acc_present = string_to_bool( + pop_item_from_dict(named_args, "accpresent", DEFAULT_STARTSTENCIL_ACC_PRESENT) + ) stencil_name = _extract_stencil_name(named_args, directive) bounds = self._make_bounds(named_args) fields = self._make_fields(named_args, field_dimensions) @@ -315,9 +317,15 @@ def _pop_additional_attributes( """Pop and return additional attributes specific to StartStencilData.""" additional_attrs = {} if dtype == StartStencilData: - mergecopy = string_to_bool(pop_item_from_dict(named_args, "mergecopy", DEFAULT_STARTSTENCIL_MERGECOPY)) - copies = string_to_bool(pop_item_from_dict(named_args, "copies", DEFAULT_STARTSTENCIL_COPIES)) - optional_module = pop_item_from_dict(named_args, "optional_module", DEFAULT_STARTSTENCIL_OPTIONAL_MODULE) + mergecopy = string_to_bool( + pop_item_from_dict(named_args, "mergecopy", DEFAULT_STARTSTENCIL_MERGECOPY) + ) + copies = string_to_bool( + pop_item_from_dict(named_args, "copies", DEFAULT_STARTSTENCIL_COPIES) + ) + optional_module = pop_item_from_dict( + named_args, "optional_module", DEFAULT_STARTSTENCIL_OPTIONAL_MODULE + ) additional_attrs = { "mergecopy": mergecopy, "copies": copies, diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index ef46952bde..bc55044466 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -149,9 +149,9 @@ def _enable_optional_stencils(self): stencils_to_remove = [] for start_stencil, end_stencil in zip(self.parsed.StartStencil, self.parsed.EndStencil): - if start_stencil.optional_module == "None": - pass - elif start_stencil.optional_module not in self.optional_modules_to_enable: - stencils_to_remove += [start_stencil, end_stencil] + if start_stencil.optional_module == "None": + pass + elif start_stencil.optional_module not in self.optional_modules_to_enable: + stencils_to_remove += [start_stencil, end_stencil] _remove_stencils(self.parsed, stencils_to_remove) diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index dba0936487..bfe92c788f 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -89,17 +89,17 @@ def process_stencils( Returns: The updated and transformed object with fields containing information from GT4Py stencils. """ - if optional_modules_to_enable==[False]: - return [ - FusedStencilTransformer(parsed, fused), - UpdateFieldsWithGt4PyStencils(parsed), - ] + if optional_modules_to_enable == [False]: + return [ + FusedStencilTransformer(parsed, fused), + UpdateFieldsWithGt4PyStencils(parsed), + ] else: - return [ - FusedStencilTransformer(parsed, fused), - OptionalModulesTransformer(parsed, optional_modules_to_enable), - UpdateFieldsWithGt4PyStencils(parsed), - ] + return [ + FusedStencilTransformer(parsed, fused), + OptionalModulesTransformer(parsed, optional_modules_to_enable), + UpdateFieldsWithGt4PyStencils(parsed), + ] @linear_pipeline From 412588813199479134b1bd43bff78cca6d337cea Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 20 Feb 2024 14:23:03 +0100 Subject: [PATCH 31/53] False doesnot work :( --- tools/src/icon4pytools/liskov/cli.py | 4 ++-- tools/src/icon4pytools/liskov/pipeline/collection.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tools/src/icon4pytools/liskov/cli.py b/tools/src/icon4pytools/liskov/cli.py index 4eae9a2b42..89a6034802 100644 --- a/tools/src/icon4pytools/liskov/cli.py +++ b/tools/src/icon4pytools/liskov/cli.py @@ -58,8 +58,8 @@ def main(ctx: click.Context) -> None: ) @click.option( "--enable-dsl-optional", - type=click.Choice(["advection", "graupel", False]), - default=[False], + type=click.Choice(["advection", "graupel", "None"]), + default=[None], multiple=True, help="Specify the DSL optional module to enable.", ) diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index bfe92c788f..f9e611b9d9 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -89,7 +89,7 @@ def process_stencils( Returns: The updated and transformed object with fields containing information from GT4Py stencils. """ - if optional_modules_to_enable == [False]: + if optional_modules_to_enable == ["None"]: return [ FusedStencilTransformer(parsed, fused), UpdateFieldsWithGt4PyStencils(parsed), From fb37d59b6e1ad9e1795b3174093e2ab523ed3bbe Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 20 Feb 2024 14:39:36 +0100 Subject: [PATCH 32/53] fix ty[es --- tools/src/icon4pytools/liskov/parsing/transform.py | 2 +- tools/src/icon4pytools/liskov/pipeline/collection.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index bc55044466..c13fbd3394 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -124,7 +124,7 @@ def _remove_delete(self) -> None: class OptionalModulesTransformer(Step): - def __init__(self, parsed: IntegrationCodeInterface, optional_modules_to_enable: list) -> None: + def __init__(self, parsed: IntegrationCodeInterface, optional_modules_to_enable: list[str]) -> None: self.parsed = parsed self.optional_modules_to_enable = optional_modules_to_enable diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index f9e611b9d9..6d8c38f222 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -74,7 +74,7 @@ def parse_fortran_file( @linear_pipeline def process_stencils( - parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: list + parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: list[str] ) -> list[Step]: """Execute a linear pipeline to transform stencils and produce either fused or unfused execution. From b9515e7eb94558121110151082e27e32d0a13348 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 20 Feb 2024 14:40:19 +0100 Subject: [PATCH 33/53] fix style --- tools/src/icon4pytools/liskov/parsing/transform.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index c13fbd3394..b85aaa5e1d 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -124,7 +124,9 @@ def _remove_delete(self) -> None: class OptionalModulesTransformer(Step): - def __init__(self, parsed: IntegrationCodeInterface, optional_modules_to_enable: list[str]) -> None: + def __init__( + self, parsed: IntegrationCodeInterface, optional_modules_to_enable: list[str] + ) -> None: self.parsed = parsed self.optional_modules_to_enable = optional_modules_to_enable From 5173126668bb0f73da5541b8fd83acdaa43fa5bf Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 20 Feb 2024 15:16:43 +0100 Subject: [PATCH 34/53] Add comma separated argument list for optional modules --- .../model/common/test_utils/helpers.py | 2 +- tools/src/icon4pytools/liskov/cli.py | 19 ++++++++----- .../icon4pytools/liskov/parsing/transform.py | 21 ++++++++++----- .../liskov/pipeline/collection.py | 27 ++++++++----------- 4 files changed, 38 insertions(+), 31 deletions(-) diff --git a/model/common/src/icon4py/model/common/test_utils/helpers.py b/model/common/src/icon4py/model/common/test_utils/helpers.py index 2858ee09ba..c5a79d5b0c 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -227,7 +227,7 @@ class StencilTest: """ PROGRAM: ClassVar[Program] - OUTPUTS: ClassVar[tuple[str, ...]] + OUTPUTS: ClassVar[tuple[str | Output, ...]] def __init_subclass__(cls, **kwargs): # Add two methods for verification and benchmarking. In order to have names that diff --git a/tools/src/icon4pytools/liskov/cli.py b/tools/src/icon4pytools/liskov/cli.py index 89a6034802..c57323992f 100644 --- a/tools/src/icon4pytools/liskov/cli.py +++ b/tools/src/icon4pytools/liskov/cli.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pathlib +from typing import Optional import click @@ -27,6 +28,10 @@ logger = setup_logger(__name__) +def split_comma(ctx, param, value) -> Optional[tuple[str]]: + return tuple(v.strip() for v in value.split(",")) if value else None + + @click.group(invoke_without_command=True) @click.pass_context def main(ctx: click.Context) -> None: @@ -57,11 +62,9 @@ def main(ctx: click.Context) -> None: help="Adds fused or unfused stencils.", ) @click.option( - "--enable-dsl-optional", - type=click.Choice(["advection", "graupel", "None"]), - default=[None], - multiple=True, - help="Specify the DSL optional module to enable.", + "--optional-modules-to-enable", + callback=split_comma, + help="Specify a list of comma-separated optional DSL modules to enable.", ) @click.argument( "input_path", @@ -77,11 +80,13 @@ def integrate( fused: bool, profile: bool, metadatagen: bool, - enable_dsl_optional: list, + optional_modules_to_enable: Optional[tuple[str]], ) -> None: mode = "integration" iface = parse_fortran_file(input_path, output_path, mode) - iface_gt4py = process_stencils(iface, fused, optional_modules_to_enable=enable_dsl_optional) + iface_gt4py = process_stencils( + iface, fused, optional_modules_to_enable=optional_modules_to_enable + ) run_code_generation( input_path, output_path, diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index bc55044466..3c743bf0bb 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -10,9 +10,10 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Any +from typing import Any, Optional from icon4pytools.common.logger import setup_logger +from icon4pytools.liskov.codegen.integration.deserialise import DEFAULT_STARTSTENCIL_OPTIONAL_MODULE from icon4pytools.liskov.codegen.integration.interface import ( EndDeleteData, EndFusedStencilData, @@ -124,7 +125,9 @@ def _remove_delete(self) -> None: class OptionalModulesTransformer(Step): - def __init__(self, parsed: IntegrationCodeInterface, optional_modules_to_enable: list) -> None: + def __init__( + self, parsed: IntegrationCodeInterface, optional_modules_to_enable: Optional[tuple[str]] + ) -> None: self.parsed = parsed self.optional_modules_to_enable = optional_modules_to_enable @@ -141,15 +144,19 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: Returns: IntegrationCodeInterface: The interface object along with any transformations applied. """ - self._enable_optional_stencils() - + if self.optional_modules_to_enable is not None: + logger.info( + f"Transforming stencils to enable optional modules: {', '.join(self.optional_modules_to_enable)}" + ) + self._enable_optional_stencils() return self.parsed - def _enable_optional_stencils(self): + def _enable_optional_stencils(self) -> None: + """Identify and remove stencils from the parse tree based on their optional module status.""" stencils_to_remove = [] - + assert self.optional_modules_to_enable is not None for start_stencil, end_stencil in zip(self.parsed.StartStencil, self.parsed.EndStencil): - if start_stencil.optional_module == "None": + if start_stencil.optional_module == DEFAULT_STARTSTENCIL_OPTIONAL_MODULE: pass elif start_stencil.optional_module not in self.optional_modules_to_enable: stencils_to_remove += [start_stencil, end_stencil] diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index f9e611b9d9..8b924926b3 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later from pathlib import Path -from typing import Any +from typing import Any, Optional from icon4pytools.liskov.codegen.integration.deserialise import IntegrationCodeDeserialiser from icon4pytools.liskov.codegen.integration.generate import IntegrationCodeGenerator @@ -74,32 +74,27 @@ def parse_fortran_file( @linear_pipeline def process_stencils( - parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: list + parsed: IntegrationCodeInterface, fused: bool, optional_modules_to_enable: Optional[tuple[str]] ) -> list[Step]: """Execute a linear pipeline to transform stencils and produce either fused or unfused execution. This function takes an input `parsed` object of type `IntegrationCodeInterface` and a `fused` boolean flag. - It then executes a linear pipeline, consisting of two steps: transformation of stencils for fusion or unfusion, - and updating fields with information from GT4Py stencils. + It then executes a linear pipeline, consisting of three steps: transformation of stencils for fusion or unfusion, + enabling optional modules, and updating fields with information from GT4Py stencils. Args: parsed (IntegrationCodeInterface): The input object containing parsed integration code. fused (bool): A boolean flag indicating whether to produce fused (True) or unfused (False) execution. + optional_modules_to_enable (Optional[tuple[str]]): A tuple of optional modules to enable. Returns: - The updated and transformed object with fields containing information from GT4Py stencils. + The updated and transformed IntegrationCodeInterface object. """ - if optional_modules_to_enable == ["None"]: - return [ - FusedStencilTransformer(parsed, fused), - UpdateFieldsWithGt4PyStencils(parsed), - ] - else: - return [ - FusedStencilTransformer(parsed, fused), - OptionalModulesTransformer(parsed, optional_modules_to_enable), - UpdateFieldsWithGt4PyStencils(parsed), - ] + return [ + FusedStencilTransformer(parsed, fused), + OptionalModulesTransformer(parsed, optional_modules_to_enable), + UpdateFieldsWithGt4PyStencils(parsed), + ] @linear_pipeline From dcab426b4e998e816e31244a711969e0ac6373e7 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 20 Feb 2024 15:19:12 +0100 Subject: [PATCH 35/53] Add comma separated argument list for optional modules --- tools/src/icon4pytools/liskov/parsing/transform.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index 3c743bf0bb..9b228102b3 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -154,11 +154,10 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: def _enable_optional_stencils(self) -> None: """Identify and remove stencils from the parse tree based on their optional module status.""" stencils_to_remove = [] - assert self.optional_modules_to_enable is not None for start_stencil, end_stencil in zip(self.parsed.StartStencil, self.parsed.EndStencil): if start_stencil.optional_module == DEFAULT_STARTSTENCIL_OPTIONAL_MODULE: pass - elif start_stencil.optional_module not in self.optional_modules_to_enable: + elif start_stencil.optional_module not in self.optional_modules_to_enable: # type: ignore stencils_to_remove += [start_stencil, end_stencil] _remove_stencils(self.parsed, stencils_to_remove) From 435ca1d51e499e35eb72db85d576f60e71afd218 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 20 Feb 2024 15:42:59 +0100 Subject: [PATCH 36/53] Test that optional module is not part of serialised savepoints --- tools/tests/liskov/test_serialisation_deserialiser.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tools/tests/liskov/test_serialisation_deserialiser.py b/tools/tests/liskov/test_serialisation_deserialiser.py index 0431086beb..518029af3a 100644 --- a/tools/tests/liskov/test_serialisation_deserialiser.py +++ b/tools/tests/liskov/test_serialisation_deserialiser.py @@ -105,6 +105,7 @@ def parsed_dict(): "horizontal_lower": "i_startidx", "horizontal_upper": "i_endidx", "accpresent": "True", + "optional_module": "advection", } ], "StartProfile": [{"name": "apply_nabla2_to_vn_in_lateral_boundary"}], @@ -127,5 +128,7 @@ def test_savepoint_data_factory(parsed_dict): savepoints = SavepointDataFactory()(parsed_dict) assert len(savepoints) == 2 assert any([isinstance(sp, SavepointData) for sp in savepoints]) + # check that unnecessary keys have been removed + assert not any(f.variable == "optional_module" for sp in savepoints for f in sp.fields) assert any([isinstance(f, FieldSerialisationData) for f in savepoints[0].fields]) assert any([isinstance(m, Metadata) for m in savepoints[0].metadata]) From 9da0eb71c589fec3f4513b5efcfa801987c01657 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 20 Feb 2024 16:16:18 +0100 Subject: [PATCH 37/53] Refactor OptionalModulesTransformer --- .../icon4pytools/liskov/parsing/transform.py | 42 +++++++++++-------- tools/tests/liskov/test_transform.py | 4 +- 2 files changed, 27 insertions(+), 19 deletions(-) diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index 9b228102b3..a723ab2f4b 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -132,32 +132,40 @@ def __init__( self.optional_modules_to_enable = optional_modules_to_enable def __call__(self, data: Any = None) -> IntegrationCodeInterface: - """Transform stencils in the parse tree based on the 'optional_modules_to_enable' flag, transforming or removing as necessary. - - This method processes stencils present in the 'parsed' object according to the 'optional_modules_to_enable' - flag. If 'optional_modules_to_enable' is a str, it identifies and processes stencils that are eligible for - deletion. If stencil is enabled with. + """Transform stencils in the parse tree based on 'optional_modules_to_enable', either enabling specific modules or removing them. Args: - data (Any): Optional data to be passed. Default is None. + data (Any): Optional data to be passed. Defaults to None. Returns: - IntegrationCodeInterface: The interface object along with any transformations applied. + IntegrationCodeInterface: The modified interface object. """ if self.optional_modules_to_enable is not None: - logger.info( - f"Transforming stencils to enable optional modules: {', '.join(self.optional_modules_to_enable)}" - ) - self._enable_optional_stencils() + action = "enabling" + else: + action = "removing" + logger.info(f"Transforming stencils by {action} optional modules.") + self._transform_stencils() + return self.parsed - def _enable_optional_stencils(self) -> None: - """Identify and remove stencils from the parse tree based on their optional module status.""" + def _transform_stencils(self) -> None: + """Identify stencils to transform based on 'optional_modules_to_enable' and applies necessary changes.""" stencils_to_remove = [] for start_stencil, end_stencil in zip(self.parsed.StartStencil, self.parsed.EndStencil): - if start_stencil.optional_module == DEFAULT_STARTSTENCIL_OPTIONAL_MODULE: - pass - elif start_stencil.optional_module not in self.optional_modules_to_enable: # type: ignore - stencils_to_remove += [start_stencil, end_stencil] + if self._should_remove_stencil(start_stencil): + stencils_to_remove.extend([start_stencil, end_stencil]) _remove_stencils(self.parsed, stencils_to_remove) + + def _should_remove_stencil(self, stencil: StartStencilData) -> bool: + """Determine if a stencil should be removed based on 'optional_modules_to_enable'. + + Returns: + bool: True if the stencil should be removed, False otherwise. + """ + if stencil.optional_module == DEFAULT_STARTSTENCIL_OPTIONAL_MODULE: + return False + if self.optional_modules_to_enable is None: + return True + return stencil.optional_module not in self.optional_modules_to_enable diff --git a/tools/tests/liskov/test_transform.py b/tools/tests/liskov/test_transform.py index 6a1898f3eb..1ac5a2bcbb 100644 --- a/tools/tests/liskov/test_transform.py +++ b/tools/tests/liskov/test_transform.py @@ -182,13 +182,13 @@ def fused_stencil_transform_unfused(integration_code_interface): @pytest.fixture def optional_modules_transform_enabled(integration_code_interface): return OptionalModulesTransformer( - integration_code_interface, optional_modules_to_enable="advection" + integration_code_interface, optional_modules_to_enable=["advection"] ) @pytest.fixture def optional_modules_transform_disabled(integration_code_interface): - return OptionalModulesTransformer(integration_code_interface, optional_modules_to_enable="no") + return OptionalModulesTransformer(integration_code_interface, optional_modules_to_enable=None) def test_transform_fused( From 7d2321d0fe05da25ffa2c2e0216f01d9aee3c071 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 20 Feb 2024 18:12:54 +0100 Subject: [PATCH 38/53] Skip slow stencils (move markers to correct place) --- .../{stencil_tests => advection_stencil_tests}/__init__.py | 0 .../test_btraj_dreg_stencil_01.py | 0 .../test_btraj_dreg_stencil_02.py | 0 .../test_btraj_dreg_stencil_03.py | 0 .../test_divide_flux_area_list_stencil_01.py | 2 +- .../test_divide_flux_area_list_stencil_02.py | 0 .../test_face_val_ppm_stencil_01.py | 0 .../test_face_val_ppm_stencil_02.py | 0 .../test_face_val_ppm_stencil_02a.py | 0 .../test_face_val_ppm_stencil_02b.py | 0 .../test_face_val_ppm_stencil_02c.py | 0 .../test_face_val_ppm_stencil_05.py | 0 .../test_hflux_ffsl_hybrid_stencil_01a.py | 0 .../test_hflux_ffsl_hybrid_stencil_02.py | 0 .../test_hflx_limiter_mo_stencil_01a.py | 0 .../test_hflx_limiter_mo_stencil_01b.py | 0 .../test_hflx_limiter_mo_stencil_02.py | 0 .../test_hflx_limiter_mo_stencil_03.py | 0 .../test_hflx_limiter_mo_stencil_04.py | 0 .../test_hflx_limiter_pd_stencil_01.py | 0 .../test_hflx_limiter_pd_stencil_02.py | 0 .../test_hor_adv_stencil_01.py | 0 .../test_mo_advection_traj_btraj_compute_o1_dsl.py | 0 .../test_prep_gauss_quadrature_c_list_stencil.py | 2 +- .../test_prep_gauss_quadrature_c_stencil.py | 2 +- .../test_rbf_intp_edge_stencil_01.py | 0 .../test_recon_lsq_cell_c_svd_stencil.py | 2 +- .../test_recon_lsq_cell_l_svd_stencil.py | 0 .../test_set_zero_c.py | 0 .../test_set_zero_c_k.py | 0 .../test_step_advection_stencil_01.py | 0 .../test_step_advection_stencil_02.py | 0 .../test_step_advection_stencil_03.py | 0 .../test_step_advection_stencil_04.py | 0 .../test_upwind_hflux_miura3_stencil_01.py | 0 .../test_upwind_hflux_miura_cycl_stencil_01.py | 0 .../test_upwind_hflux_miura_cycl_stencil_02.py | 0 .../test_upwind_hflux_miura_cycl_stencil_03a.py | 0 .../test_upwind_hflux_miura_cycl_stencil_03b.py | 0 .../test_upwind_hflux_miura_stencil_01.py | 0 .../test_upwind_vflux_ppm_stencil_01.py | 0 .../test_vert_adv_stencil_01.py | 0 .../test_vlimit_prbl_sm_stencil_01.py | 0 .../test_vlimit_prbl_sm_stencil_02.py | 0 model/tox.ini | 3 ++- 45 files changed, 6 insertions(+), 5 deletions(-) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/__init__.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_btraj_dreg_stencil_01.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_btraj_dreg_stencil_02.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_btraj_dreg_stencil_03.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_divide_flux_area_list_stencil_01.py (99%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_divide_flux_area_list_stencil_02.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_face_val_ppm_stencil_01.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_face_val_ppm_stencil_02.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_face_val_ppm_stencil_02a.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_face_val_ppm_stencil_02b.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_face_val_ppm_stencil_02c.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_face_val_ppm_stencil_05.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_hflux_ffsl_hybrid_stencil_01a.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_hflux_ffsl_hybrid_stencil_02.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_hflx_limiter_mo_stencil_01a.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_hflx_limiter_mo_stencil_01b.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_hflx_limiter_mo_stencil_02.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_hflx_limiter_mo_stencil_03.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_hflx_limiter_mo_stencil_04.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_hflx_limiter_pd_stencil_01.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_hflx_limiter_pd_stencil_02.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_hor_adv_stencil_01.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_mo_advection_traj_btraj_compute_o1_dsl.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_prep_gauss_quadrature_c_list_stencil.py (99%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_prep_gauss_quadrature_c_stencil.py (99%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_rbf_intp_edge_stencil_01.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_recon_lsq_cell_c_svd_stencil.py (99%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_recon_lsq_cell_l_svd_stencil.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_set_zero_c.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_set_zero_c_k.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_step_advection_stencil_01.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_step_advection_stencil_02.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_step_advection_stencil_03.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_step_advection_stencil_04.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_upwind_hflux_miura3_stencil_01.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_upwind_hflux_miura_cycl_stencil_01.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_upwind_hflux_miura_cycl_stencil_02.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_upwind_hflux_miura_cycl_stencil_03a.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_upwind_hflux_miura_cycl_stencil_03b.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_upwind_hflux_miura_stencil_01.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_upwind_vflux_ppm_stencil_01.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_vert_adv_stencil_01.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_vlimit_prbl_sm_stencil_01.py (100%) rename model/atmosphere/advection/tests/{stencil_tests => advection_stencil_tests}/test_vlimit_prbl_sm_stencil_02.py (100%) diff --git a/model/atmosphere/advection/tests/stencil_tests/__init__.py b/model/atmosphere/advection/tests/advection_stencil_tests/__init__.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/__init__.py rename to model/atmosphere/advection/tests/advection_stencil_tests/__init__.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_02.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_03.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_03.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py similarity index 99% rename from model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py index 4a2258286c..1b5b9d7c62 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py @@ -118,6 +118,7 @@ def _line_intersect_numpy( return intersect_1, intersect_2 +@pytest.mark.slow_tests class TestDivideFluxAreaListStencil01(StencilTest): PROGRAM = divide_flux_area_list_stencil_01 OUTPUTS = ( @@ -1478,7 +1479,6 @@ def reference( dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, ) - @pytest.mark.slow_tests @pytest.fixture def input_data(self, grid): famask_int = random_mask(grid, EdgeDim, KDim, dtype=int32) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_02.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02a.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02a.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02b.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02b.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02b.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02b.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02c.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02c.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02c.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02c.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_05.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_05.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01a.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01a.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01a.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01b.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01b.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01b.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_02.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_03.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_03.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_04.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_04.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_04.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hor_adv_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_hor_adv_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py similarity index 99% rename from model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py index ba9ba72d47..f883693ab1 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py @@ -27,6 +27,7 @@ ) +@pytest.mark.slow_tests class TestPrepGaussQuadratureCListStencil(StencilTest): PROGRAM = prep_gauss_quadrature_c_list_stencil OUTPUTS = ( @@ -527,7 +528,6 @@ def reference( p_dreg_area=p_dreg_area, ) - @pytest.mark.slow_tests @pytest.fixture def input_data(self, grid): famask_int = constant_field(grid, 1, EdgeDim, KDim, dtype=int32) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_stencil.py similarity index 99% rename from model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_stencil.py index 07c2e16a7e..2252493078 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_stencil.py @@ -21,6 +21,7 @@ from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field +@pytest.mark.slow_tests class TestPrepGaussQuadratureCStencil(StencilTest): PROGRAM = prep_gauss_quadrature_c_stencil OUTPUTS = ( @@ -487,7 +488,6 @@ def reference( p_dreg_area_out=p_dreg_area_out, ) - @pytest.mark.slow_tests @pytest.fixture def input_data(self, grid): p_coords_dreg_v_1_x = random_field(grid, EdgeDim, KDim) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_rbf_intp_edge_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_rbf_intp_edge_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_rbf_intp_edge_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_c_svd_stencil.py similarity index 99% rename from model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_c_svd_stencil.py index 03a627840c..c3c5bfec8c 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_c_svd_stencil.py @@ -27,6 +27,7 @@ ) +@pytest.mark.slow_tests class TestReconLsqCellCSvdStencil(StencilTest): PROGRAM = recon_lsq_cell_c_svd_stencil OUTPUTS = ( @@ -235,7 +236,6 @@ def reference( p_coeff_10_dsl=p_coeff_10_dsl, ) - @pytest.mark.slow_tests @pytest.fixture def input_data(self, grid): p_cc = random_field(grid, CellDim, KDim) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_l_svd_stencil.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_l_svd_stencil.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_l_svd_stencil.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_set_zero_c.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_set_zero_c.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_set_zero_c.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_set_zero_c.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_set_zero_c_k.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_set_zero_c_k.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_set_zero_c_k.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_set_zero_c_k.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_02.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_03.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_03.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_03.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_03.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_04.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_04.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_04.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_04.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura3_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura3_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_vflux_ppm_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_upwind_vflux_ppm_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_vflux_ppm_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_vert_adv_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_vert_adv_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_01.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_01.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_01.py diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_02.py similarity index 100% rename from model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_02.py rename to model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_02.py diff --git a/model/tox.ini b/model/tox.ini index 7aed4fb051..14fdd5cb32 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -30,12 +30,13 @@ allowlist_externals = commands = pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/diffusion/tests/diffusion_stencil_tests --benchmark-skip -n auto {posargs} pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/dycore/tests/dycore_stencil_tests --benchmark-skip -n auto {posargs} + pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/advection/tests/advection_stencil_tests --benchmark-skip -n auto {posargs} [testenv:run_benchmarks] commands = pytest -s -m "not slow_tests" atmosphere/diffusion/tests/diffusion_stencil_tests --benchmark-only {posargs} pytest -s -m "not slow_tests" atmosphere/dycore/tests/dycore_stencil_tests --benchmark-only {posargs} - + pytest -s -m "not slow_tests" atmosphere/advection/tests/advection_stencil_tests --benchmark-only {posargs} [testenv:dev] setenv = From a0724ae1319a9a3f3853c6473997e82443be7018 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 21 Feb 2024 14:48:44 +0100 Subject: [PATCH 39/53] rm stencil --- .../model/atmosphere/advection/set_zero_c.py | 28 ---------------- .../test_set_zero_c.py | 32 ------------------- 2 files changed, 60 deletions(-) delete mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_set_zero_c.py diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py deleted file mode 100644 index 93e21929e3..0000000000 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py +++ /dev/null @@ -1,28 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -from gt4py.next.common import Field, GridType -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import broadcast - -from icon4py.model.common.dimension import CellDim - - -@field_operator -def _set_zero_c() -> Field[[CellDim], float]: - return broadcast(0.0, (CellDim,)) - - -@program(grid_type=GridType.UNSTRUCTURED) -def set_zero_c(field: Field[[CellDim], float]): - _set_zero_c(out=field) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_set_zero_c.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_set_zero_c.py deleted file mode 100644 index f15449b3f4..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_set_zero_c.py +++ /dev/null @@ -1,32 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import pytest - -from icon4py.model.atmosphere.advection.set_zero_c import set_zero_c -from icon4py.model.common.dimension import CellDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestSetZeroCellK(StencilTest): - PROGRAM = set_zero_c - OUTPUTS = ("field",) - - @staticmethod - def reference(grid, **kwargs): - return dict(field=zero_field(grid, CellDim).asnumpy()) - - @pytest.fixture - def input_data(self, grid): - field = random_field(grid, CellDim) - return dict(field=field) From cc72a57b875e9851510719bbd234d08f89213d01 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 26 Feb 2024 13:52:19 +0100 Subject: [PATCH 40/53] fix style --- tools/src/icon4pytools/liskov/parsing/transform.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index ce729bf6de..ed1c8ad93b 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -154,7 +154,9 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: def _transform_stencils(self) -> None: """Identify stencils to transform based on 'optional_modules_to_enable' and applies necessary changes.""" stencils_to_remove = [] - for start_stencil, end_stencil in zip(self.parsed.StartStencil, self.parsed.EndStencil): + for start_stencil, end_stencil in zip( + self.parsed.StartStencil, self.parsed.EndStencil, strict=False + ): if self._should_remove_stencil(start_stencil): stencils_to_remove.extend([start_stencil, end_stencil]) From f2ae727ceac04d0003d43f7934b9d6a283923c95 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 28 Feb 2024 15:51:12 +0100 Subject: [PATCH 41/53] update stencil --- .../advection/hflx_limiter_pd_stencil_02.py | 14 ++----------- .../test_hflx_limiter_pd_stencil_02.py | 21 ++----------------- 2 files changed, 4 insertions(+), 31 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py index 2380296e9e..d6d3681af5 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py @@ -19,34 +19,24 @@ @field_operator def _hflx_limiter_pd_stencil_02( - refin_ctrl: Field[[EdgeDim], int32], r_m: Field[[CellDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], - bound: int32, ) -> Field[[EdgeDim, KDim], float]: - p_mflx_tracer_h_out = where( - refin_ctrl == bound, - p_mflx_tracer_h, - where( + p_mflx_tracer_h_out = where( p_mflx_tracer_h >= 0.0, p_mflx_tracer_h * r_m(E2C[0]), p_mflx_tracer_h * r_m(E2C[1]), - ), - ) + ) return p_mflx_tracer_h_out @program def hflx_limiter_pd_stencil_02( - refin_ctrl: Field[[EdgeDim], int32], r_m: Field[[CellDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], - bound: int32, ): _hflx_limiter_pd_stencil_02( - refin_ctrl, r_m, p_mflx_tracer_h, - bound, out=p_mflx_tracer_h, ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py index 212f95664a..44bcc8df4b 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py @@ -23,40 +23,23 @@ class TestHflxLimiterPdStencil02(StencilTest): OUTPUTS = ("p_mflx_tracer_h",) @staticmethod - def reference(grid, refin_ctrl, r_m, p_mflx_tracer_h, bound, **kwargs): + def reference(grid, r_m, p_mflx_tracer_h, **kwargs): e2c = grid.connectivities[E2CDim] r_m_e2c = r_m[e2c] - refin_ctrl_expanded = np.expand_dims(refin_ctrl, axis=-1) p_mflx_tracer_h_out = np.where( - refin_ctrl_expanded != bound, - np.where( p_mflx_tracer_h >= 0, p_mflx_tracer_h * r_m_e2c[:, 0], p_mflx_tracer_h * r_m_e2c[:, 1], - ), - p_mflx_tracer_h, - ) + ) return dict(p_mflx_tracer_h=p_mflx_tracer_h_out) @pytest.fixture(params=[("no_match", 4), ("everywhere_match", 7), ("partly_match", 4)]) def input_data(self, request, grid): - bound = np.int32(7) - scenario, ctrl_value = request.param - - if scenario == "no_match": - refin_ctrl = constant_field(grid, ctrl_value, EdgeDim, dtype=np.int32) - elif scenario == "everywhere_match": - refin_ctrl = constant_field(grid, bound, EdgeDim, dtype=np.int32) - elif scenario == "partly_match": - refin_ctrl = constant_field(grid, 5, EdgeDim, dtype=np.int32) - refin_ctrl[2:6] = bound r_m = random_field(grid, CellDim, KDim) p_mflx_tracer_h_in = random_field(grid, EdgeDim, KDim) return dict( - refin_ctrl=refin_ctrl, r_m=r_m, p_mflx_tracer_h=p_mflx_tracer_h_in, - bound=bound, ) From fd96ec470537db0e6ca7104233f972190583adaa Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 28 Feb 2024 15:52:14 +0100 Subject: [PATCH 42/53] fix style --- .../advection/hflx_limiter_pd_stencil_02.py | 12 ++++++------ .../test_hflx_limiter_pd_stencil_02.py | 11 +++++------ 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py index d6d3681af5..aeb1ff5db1 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, int32, where +from gt4py.next.ffront.fbuiltins import Field, where from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @@ -22,11 +22,11 @@ def _hflx_limiter_pd_stencil_02( r_m: Field[[CellDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - p_mflx_tracer_h_out = where( - p_mflx_tracer_h >= 0.0, - p_mflx_tracer_h * r_m(E2C[0]), - p_mflx_tracer_h * r_m(E2C[1]), - ) + p_mflx_tracer_h_out = where( + p_mflx_tracer_h >= 0.0, + p_mflx_tracer_h * r_m(E2C[0]), + p_mflx_tracer_h * r_m(E2C[1]), + ) return p_mflx_tracer_h_out diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py index 44bcc8df4b..166f376e41 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py @@ -15,7 +15,7 @@ from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_02 import hflx_limiter_pd_stencil_02 from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, constant_field, random_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field class TestHflxLimiterPdStencil02(StencilTest): @@ -27,15 +27,14 @@ def reference(grid, r_m, p_mflx_tracer_h, **kwargs): e2c = grid.connectivities[E2CDim] r_m_e2c = r_m[e2c] p_mflx_tracer_h_out = np.where( - p_mflx_tracer_h >= 0, - p_mflx_tracer_h * r_m_e2c[:, 0], - p_mflx_tracer_h * r_m_e2c[:, 1], - ) + p_mflx_tracer_h >= 0, + p_mflx_tracer_h * r_m_e2c[:, 0], + p_mflx_tracer_h * r_m_e2c[:, 1], + ) return dict(p_mflx_tracer_h=p_mflx_tracer_h_out) @pytest.fixture(params=[("no_match", 4), ("everywhere_match", 7), ("partly_match", 4)]) def input_data(self, request, grid): - r_m = random_field(grid, CellDim, KDim) p_mflx_tracer_h_in = random_field(grid, EdgeDim, KDim) From d63fdf89c13ef751a541e43e25671a17440d4046 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Fri, 1 Mar 2024 17:42:02 +0100 Subject: [PATCH 43/53] add commentns --- .../advection/btraj_dreg_stencil_03.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py index d2c2fd2c72..8660be2e76 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py @@ -50,20 +50,38 @@ def _btraj_dreg_stencil_03( Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], ]: + # logical switch for MERGE operations: .TRUE. for p_vn >= 0 lvn_pos = where(p_vn >= 0.0, True, False) + # get line and block indices of upwind cell p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1])) p_cell_rel_idx_dsl = where(lvn_pos, int32(0), int32(1)) p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) + # departure points of the departure cell. Point 1 belongs to edge-vertex 1, + # point 2 belongs to edge_vertex 2. + # + # position of vertex 4 (vn > 0) / vertex 2(vn < 0) in normal direction depart_pts_1_x = edge_verts_1_x - p_vn * p_dt + # position of vertex 4 (vn > 0) / vertex 2(vn < 0) in tangential direction depart_pts_1_y = edge_verts_1_y - p_vt * p_dt + # position of vertex 3 in normal direction depart_pts_2_x = edge_verts_2_x - p_vn * p_dt + # position of vertex 3 in tangential direction depart_pts_2_y = edge_verts_2_y - p_vt * p_dt + # determine correct position on tangential plane pos_on_tplane_e_x = where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) pos_on_tplane_e_y = where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) + + # Calculate position of departure region vertices in a translated + # coordinate system. The origin is located at the circumcenter + # of the upwind cell. The distance vectors point from the cell center + # to the vertices. + # + # Take care of correct counterclockwise numbering below + # pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y pos_dreg_vert_c_2_x = where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x @@ -73,15 +91,24 @@ def _btraj_dreg_stencil_03( pos_dreg_vert_c_4_x = where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x pos_dreg_vert_c_4_y = where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y + # In a last step, these distance vectors are transformed into a rotated + # geographical coordinate system, which still has its origin at the circumcenter + # of the upwind cell. Now the coordinate axes point to local East and local + # North. + # + # Determine primal and dual normals of the cell lying in the direction of vn pn_cell_1 = where(lvn_pos, primal_normal_cell_x(E2EC[0]), primal_normal_cell_x(E2EC[1])) pn_cell_2 = where(lvn_pos, primal_normal_cell_y(E2EC[0]), primal_normal_cell_y(E2EC[1])) dn_cell_1 = where(lvn_pos, dual_normal_cell_x(E2EC[0]), dual_normal_cell_x(E2EC[1])) dn_cell_2 = where(lvn_pos, dual_normal_cell_y(E2EC[0]), dual_normal_cell_y(E2EC[1])) + # components in longitudinal direction p_coords_dreg_v_1_lon_dsl = pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 p_coords_dreg_v_2_lon_dsl = pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 p_coords_dreg_v_3_lon_dsl = pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 p_coords_dreg_v_4_lon_dsl = pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 + + # components in latitudinal direction p_coords_dreg_v_1_lat_dsl = pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 p_coords_dreg_v_2_lat_dsl = pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 p_coords_dreg_v_3_lat_dsl = pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 From 6e30b3afd86bea834890ce3f567c3d5e78988c79 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 11 Mar 2024 09:17:16 +0100 Subject: [PATCH 44/53] fix style --- .../model/atmosphere/advection/btraj_dreg_stencil_03.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py index 8660be2e76..2e59e8ac38 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py @@ -58,7 +58,7 @@ def _btraj_dreg_stencil_03( p_cell_rel_idx_dsl = where(lvn_pos, int32(0), int32(1)) p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) - # departure points of the departure cell. Point 1 belongs to edge-vertex 1, + # departure points of the departure cell. Point 1 belongs to edge-vertex 1, # point 2 belongs to edge_vertex 2. # # position of vertex 4 (vn > 0) / vertex 2(vn < 0) in normal direction @@ -74,7 +74,6 @@ def _btraj_dreg_stencil_03( pos_on_tplane_e_x = where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) pos_on_tplane_e_y = where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) - # Calculate position of departure region vertices in a translated # coordinate system. The origin is located at the circumcenter # of the upwind cell. The distance vectors point from the cell center From 888936bf6a3ec8ba868aa45879cca019a9d84421 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 7 May 2024 17:49:41 +0200 Subject: [PATCH 45/53] rm stencil --- .../model/atmosphere/advection/init_zero_c.py | 28 ---------------- .../test_init_zero_c.py | 32 ------------------- 2 files changed, 60 deletions(-) delete mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/init_zero_c.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_init_zero_c.py diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/init_zero_c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/init_zero_c.py deleted file mode 100644 index 10fa5b4998..0000000000 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/init_zero_c.py +++ /dev/null @@ -1,28 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -from gt4py.next.common import Field, GridType -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import broadcast - -from icon4py.model.common.dimension import CellDim - - -@field_operator -def _init_zero_c() -> Field[[CellDim], float]: - return broadcast(0.0, (CellDim,)) - - -@program(grid_type=GridType.UNSTRUCTURED) -def init_zero_c(field: Field[[CellDim], float]): - _init_zero_c(out=field) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_init_zero_c.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_init_zero_c.py deleted file mode 100644 index 08da338d49..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_init_zero_c.py +++ /dev/null @@ -1,32 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import pytest - -from icon4py.model.atmosphere.advection.init_zero_c import init_zero_c -from icon4py.model.common.dimension import CellDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestInitZeroCellK(StencilTest): - PROGRAM = init_zero_c - OUTPUTS = ("field",) - - @staticmethod - def reference(grid, **kwargs): - return dict(field=zero_field(grid, CellDim).asnumpy()) - - @pytest.fixture - def input_data(self, grid): - field = random_field(grid, CellDim) - return dict(field=field) From 8e95d613f0e856ae2d35bffc059af79fb6cde1c6 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 8 May 2024 17:05:00 +0200 Subject: [PATCH 46/53] fix division with zero --- .../divide_flux_area_list_stencil_01.py | 13 ++++++-- .../test_divide_flux_area_list_stencil_01.py | 33 ++++++++++++------- 2 files changed, 32 insertions(+), 14 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py index 2fa1051c61..2801699f67 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py @@ -107,8 +107,17 @@ def line_intersect( line2_p2_lon: Field[[EdgeDim, KDim], float], line2_p2_lat: Field[[EdgeDim, KDim], float], ) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: - m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon) - m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) + + + # avoid division with zero + d1 = line1_p2_lon - line1_p1_lon + d1 = where( d1 !=float(0), d1, line1_p2_lon) + + d2 = line2_p2_lon - line2_p1_lon + d2 = where( d1 !=float(0), d1, line2_p2_lon) + + m1 = (line1_p2_lat - line1_p1_lat) / d1 + m2 = (line2_p2_lat - line2_p1_lat) / d2 intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / (m1 - m2) intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py index 392e01e4a1..817a09c7a1 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py @@ -22,6 +22,7 @@ from icon4py.model.common.test_utils.helpers import ( StencilTest, as_1D_sparse_field, + constant_field, random_field, random_mask, reshape, @@ -109,7 +110,14 @@ def _line_intersect_numpy( line2_p2_lon, line2_p2_lat, ): - m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon) + + d1 = line1_p2_lon - line1_p1_lon + d1 = np.where( d1 !=float(0), d1, line1_p2_lon) + + d2 = line2_p2_lon - line2_p1_lon + d2 = np.where( d1 !=float(0), d1, line2_p2_lon) + + m1 = (line1_p2_lat - line1_p1_lat) / d1 m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / (m1 - m2) @@ -589,6 +597,7 @@ def _apply_case2b_patch2( dreg_patch2_3_lon_vmask, dreg_patch2_3_lat_vmask, ): + dreg_patch2_1_lon_vmask = np.where( mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask ) @@ -1482,19 +1491,19 @@ def reference( def input_data(self, grid): famask_int = random_mask(grid, EdgeDim, KDim, dtype=int32) p_vn = random_field(grid, EdgeDim, KDim) - ptr_v3_lon = random_field(grid, EdgeDim, E2CDim) + ptr_v3_lon = random_field(grid, EdgeDim, E2CDim, low=0.1, high=1.0) ptr_v3_lon_field = as_1D_sparse_field(ptr_v3_lon, ECDim) - ptr_v3_lat = random_field(grid, EdgeDim, E2CDim) + ptr_v3_lat = random_field(grid, EdgeDim, E2CDim, low=0.1, high=1.0) ptr_v3_lat_field = as_1D_sparse_field(ptr_v3_lat, ECDim) - tangent_orientation_dsl = random_field(grid, EdgeDim) - dreg_patch0_1_lon_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_1_lat_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_2_lon_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_2_lat_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_3_lon_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_3_lat_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_4_lon_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_4_lat_dsl = random_field(grid, EdgeDim, KDim) + tangent_orientation_dsl = random_field(grid, EdgeDim, low=0.1, high=1.0) + dreg_patch0_1_lon_dsl = constant_field(grid, 1.0, EdgeDim, KDim) + dreg_patch0_1_lat_dsl = constant_field(grid, 1.0, EdgeDim, KDim) + dreg_patch0_2_lon_dsl = constant_field(grid, 2.0, EdgeDim, KDim) + dreg_patch0_2_lat_dsl = constant_field(grid, 2.0, EdgeDim, KDim) + dreg_patch0_3_lon_dsl = constant_field(grid, 3.0, EdgeDim, KDim) + dreg_patch0_3_lat_dsl = constant_field(grid, 3.0, EdgeDim, KDim) + dreg_patch0_4_lon_dsl = constant_field(grid, 4.0, EdgeDim, KDim) + dreg_patch0_4_lat_dsl = constant_field(grid, 4.0, EdgeDim, KDim) dreg_patch1_1_lon_vmask = zero_field(grid, EdgeDim, KDim) dreg_patch1_1_lat_vmask = zero_field(grid, EdgeDim, KDim) dreg_patch1_2_lon_vmask = zero_field(grid, EdgeDim, KDim) From 55091fe5be70d63cea1e93f769173f525de63185 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 8 May 2024 17:23:54 +0200 Subject: [PATCH 47/53] fix style --- .../advection/divide_flux_area_list_stencil_01.py | 10 ++++------ .../test_divide_flux_area_list_stencil_01.py | 10 ++++------ 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py index 2801699f67..231a364324 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py @@ -107,14 +107,12 @@ def line_intersect( line2_p2_lon: Field[[EdgeDim, KDim], float], line2_p2_lat: Field[[EdgeDim, KDim], float], ) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: - - # avoid division with zero - d1 = line1_p2_lon - line1_p1_lon - d1 = where( d1 !=float(0), d1, line1_p2_lon) + d1 = line1_p2_lon - line1_p1_lon + d1 = where(d1 != float(0), d1, line1_p2_lon) - d2 = line2_p2_lon - line2_p1_lon - d2 = where( d1 !=float(0), d1, line2_p2_lon) + d2 = line2_p2_lon - line2_p1_lon + d2 = where(d1 != float(0), d1, line2_p2_lon) m1 = (line1_p2_lat - line1_p1_lat) / d1 m2 = (line2_p2_lat - line2_p1_lat) / d2 diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py index 817a09c7a1..3d63d513e3 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py @@ -110,12 +110,11 @@ def _line_intersect_numpy( line2_p2_lon, line2_p2_lat, ): + d1 = line1_p2_lon - line1_p1_lon + d1 = np.where(d1 != float(0), d1, line1_p2_lon) - d1 = line1_p2_lon - line1_p1_lon - d1 = np.where( d1 !=float(0), d1, line1_p2_lon) - - d2 = line2_p2_lon - line2_p1_lon - d2 = np.where( d1 !=float(0), d1, line2_p2_lon) + d2 = line2_p2_lon - line2_p1_lon + d2 = np.where(d1 != float(0), d1, line2_p2_lon) m1 = (line1_p2_lat - line1_p1_lat) / d1 m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) @@ -597,7 +596,6 @@ def _apply_case2b_patch2( dreg_patch2_3_lon_vmask, dreg_patch2_3_lat_vmask, ): - dreg_patch2_1_lon_vmask = np.where( mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask ) From bccc0e8af596e1ccfa3d28f12b05a35bca2987cc Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 14 May 2024 09:30:07 +0200 Subject: [PATCH 48/53] fix stencil --- .../atmosphere/advection/divide_flux_area_list_stencil_01.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py index 231a364324..5024dad193 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py @@ -112,7 +112,7 @@ def line_intersect( d1 = where(d1 != float(0), d1, line1_p2_lon) d2 = line2_p2_lon - line2_p1_lon - d2 = where(d1 != float(0), d1, line2_p2_lon) + d2 = where(d2 != float(0), d2, line2_p2_lon) m1 = (line1_p2_lat - line1_p1_lat) / d1 m2 = (line2_p2_lat - line2_p1_lat) / d2 From ded02ac34864b6b5197e2fa27f023e3536d0dad5 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 14 May 2024 09:38:31 +0200 Subject: [PATCH 49/53] fix 2 --- .../test_divide_flux_area_list_stencil_01.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py index 3d63d513e3..034799fbb1 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py @@ -114,10 +114,10 @@ def _line_intersect_numpy( d1 = np.where(d1 != float(0), d1, line1_p2_lon) d2 = line2_p2_lon - line2_p1_lon - d2 = np.where(d1 != float(0), d1, line2_p2_lon) + d2 = np.where(d2 != float(0), d2, line2_p2_lon) m1 = (line1_p2_lat - line1_p1_lat) / d1 - m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) + m2 = (line2_p2_lat - line2_p1_lat) / d2 intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / (m1 - m2) intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) From 6cd3ea16865e88283c90aa8725940ac916c3e23e Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 14 May 2024 11:55:55 +0200 Subject: [PATCH 50/53] dbg test --- .../model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py | 2 +- .../advection_stencil_tests/test_hflx_limiter_mo_stencil_01a.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py index 7ea34b2978..3d7a20bdce 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py @@ -25,7 +25,7 @@ def _hflx_limiter_mo_stencil_01a( ) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: z_mflx_low = 0.5 * ( p_mass_flx_e * (p_cc(E2C[0]) + p_cc(E2C[1])) - - abs(p_mass_flx_e) * (p_cc(E2C[1]) - p_cc(E2C[0])) + # - abs(p_mass_flx_e) * (p_cc(E2C[1]) - p_cc(E2C[0])) ) z_anti = p_mflx_tracer_h - z_mflx_low diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01a.py index 0dbf30ca22..a885048535 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01a.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01a.py @@ -37,7 +37,7 @@ def reference( z_mflx_low = 0.5 * ( p_mass_flx_e * (p_cc_e2c[:, 0] + p_cc_e2c[:, 1]) - - np.absolute(p_mass_flx_e) * (p_cc_e2c[:, 1] - p_cc_e2c[:, 0]) + # - np.absolute(p_mass_flx_e) * (p_cc_e2c[:, 1] - p_cc_e2c[:, 0]) ) z_anti = p_mflx_tracer_h - z_mflx_low From 5e0e069d43bc7106918e718fc5cbdc40c518a1e5 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 15 May 2024 14:27:16 +0200 Subject: [PATCH 51/53] change default --- ci/default.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ci/default.yml b/ci/default.yml index 4a85316030..7c17844fac 100644 --- a/ci/default.yml +++ b/ci/default.yml @@ -17,7 +17,7 @@ test_model_stencils: # exclude slow test configurations - if: $BACKEND == "roundtrip" && $GRID == "icon_grid" when: never - - when: always + - when: on_success test_tools: extends: .test_template From 33b884e8226a01dcce360518d2215099aabdb037 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 15 May 2024 15:19:36 +0200 Subject: [PATCH 52/53] seperate liskov and utests --- .../advection/btraj_dreg_stencil_03.py | 26 - .../divide_flux_area_list_stencil_01.py | 11 +- .../advection/hflx_limiter_mo_stencil_01a.py | 2 +- .../advection/hflx_limiter_pd_stencil_02.py | 18 +- .../advection/hor_adv_stencil_01.py | 8 +- .../model/atmosphere/advection/init_zero_c.py | 28 + .../mo_advection_traj_btraj_compute_o1_dsl.py | 19 +- .../advection/vert_adv_stencil_01.py | 48 +- .../advection/tests/stencil_tests/__init__.py | 12 + .../test_btraj_dreg_stencil_01.py | 52 + .../test_btraj_dreg_stencil_02.py | 66 + .../test_btraj_dreg_stencil_03.py | 233 +++ .../test_divide_flux_area_list_stencil_01.py | 1544 +++++++++++++++++ .../test_divide_flux_area_list_stencil_02.py | 290 ++++ .../test_face_val_ppm_stencil_01.py | 78 + .../test_face_val_ppm_stencil_02.py | 76 + .../test_face_val_ppm_stencil_02a.py | 43 + .../test_face_val_ppm_stencil_02b.py | 35 + .../test_face_val_ppm_stencil_02c.py | 52 + .../test_face_val_ppm_stencil_05.py | 71 + .../test_hflux_ffsl_hybrid_stencil_01a.py | 184 ++ .../test_hflux_ffsl_hybrid_stencil_02.py | 49 + .../test_hflx_limiter_mo_stencil_01a.py | 61 + .../test_hflx_limiter_mo_stencil_01b.py | 120 ++ .../test_hflx_limiter_mo_stencil_02.py | 85 + .../test_hflx_limiter_mo_stencil_03.py | 101 ++ .../test_hflx_limiter_mo_stencil_04.py | 50 + .../test_hflx_limiter_pd_stencil_01.py | 77 + .../test_hflx_limiter_pd_stencil_02.py | 62 + .../stencil_tests/test_hor_adv_stencil_01.py | 74 + .../tests/stencil_tests/test_init_zero_c.py | 32 + .../tests/stencil_tests/test_init_zero_c_k.py | 32 + ..._mo_advection_traj_btraj_compute_o1_dsl.py | 150 ++ ...st_prep_gauss_quadrature_c_list_stencil.py | 636 +++++++ .../test_prep_gauss_quadrature_c_stencil.py | 591 +++++++ .../test_rbf_intp_edge_stencil_01.py | 37 + .../test_recon_lsq_cell_c_svd_stencil.py | 318 ++++ .../test_recon_lsq_cell_l_svd_stencil.py | 77 + .../test_step_advection_stencil_01.py | 64 + .../test_step_advection_stencil_02.py | 61 + .../test_step_advection_stencil_03.py | 50 + .../test_step_advection_stencil_04.py | 49 + .../test_upwind_hflux_miura3_stencil_01.py | 194 +++ ...test_upwind_hflux_miura_cycl_stencil_01.py | 92 + ...test_upwind_hflux_miura_cycl_stencil_02.py | 93 + ...est_upwind_hflux_miura_cycl_stencil_03a.py | 47 + ...est_upwind_hflux_miura_cycl_stencil_03b.py | 50 + .../test_upwind_hflux_miura_stencil_01.py | 92 + .../test_upwind_vflux_ppm_stencil_01.py | 45 + .../stencil_tests/test_vert_adv_stencil_01.py | 64 + .../test_vlimit_prbl_sm_stencil_01.py | 39 + .../test_vlimit_prbl_sm_stencil_02.py | 58 + .../model/common/test_utils/helpers.py | 2 +- model/tox.ini | 3 +- 54 files changed, 6342 insertions(+), 109 deletions(-) create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/init_zero_c.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/__init__.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_01.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_02.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_02.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02b.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02c.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01a.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01b.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_02.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_04.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_01.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_init_zero_c.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_init_zero_c_k.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_rbf_intp_edge_stencil_01.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_l_svd_stencil.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_01.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_02.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_03.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_04.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_upwind_vflux_ppm_stencil_01.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_01.py create mode 100644 model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_02.py diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py index 2e59e8ac38..d2c2fd2c72 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py @@ -50,37 +50,20 @@ def _btraj_dreg_stencil_03( Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], ]: - # logical switch for MERGE operations: .TRUE. for p_vn >= 0 lvn_pos = where(p_vn >= 0.0, True, False) - # get line and block indices of upwind cell p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1])) p_cell_rel_idx_dsl = where(lvn_pos, int32(0), int32(1)) p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) - # departure points of the departure cell. Point 1 belongs to edge-vertex 1, - # point 2 belongs to edge_vertex 2. - # - # position of vertex 4 (vn > 0) / vertex 2(vn < 0) in normal direction depart_pts_1_x = edge_verts_1_x - p_vn * p_dt - # position of vertex 4 (vn > 0) / vertex 2(vn < 0) in tangential direction depart_pts_1_y = edge_verts_1_y - p_vt * p_dt - # position of vertex 3 in normal direction depart_pts_2_x = edge_verts_2_x - p_vn * p_dt - # position of vertex 3 in tangential direction depart_pts_2_y = edge_verts_2_y - p_vt * p_dt - # determine correct position on tangential plane pos_on_tplane_e_x = where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) pos_on_tplane_e_y = where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) - # Calculate position of departure region vertices in a translated - # coordinate system. The origin is located at the circumcenter - # of the upwind cell. The distance vectors point from the cell center - # to the vertices. - # - # Take care of correct counterclockwise numbering below - # pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y pos_dreg_vert_c_2_x = where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x @@ -90,24 +73,15 @@ def _btraj_dreg_stencil_03( pos_dreg_vert_c_4_x = where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x pos_dreg_vert_c_4_y = where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y - # In a last step, these distance vectors are transformed into a rotated - # geographical coordinate system, which still has its origin at the circumcenter - # of the upwind cell. Now the coordinate axes point to local East and local - # North. - # - # Determine primal and dual normals of the cell lying in the direction of vn pn_cell_1 = where(lvn_pos, primal_normal_cell_x(E2EC[0]), primal_normal_cell_x(E2EC[1])) pn_cell_2 = where(lvn_pos, primal_normal_cell_y(E2EC[0]), primal_normal_cell_y(E2EC[1])) dn_cell_1 = where(lvn_pos, dual_normal_cell_x(E2EC[0]), dual_normal_cell_x(E2EC[1])) dn_cell_2 = where(lvn_pos, dual_normal_cell_y(E2EC[0]), dual_normal_cell_y(E2EC[1])) - # components in longitudinal direction p_coords_dreg_v_1_lon_dsl = pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 p_coords_dreg_v_2_lon_dsl = pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 p_coords_dreg_v_3_lon_dsl = pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 p_coords_dreg_v_4_lon_dsl = pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 - - # components in latitudinal direction p_coords_dreg_v_1_lat_dsl = pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 p_coords_dreg_v_2_lat_dsl = pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 p_coords_dreg_v_3_lat_dsl = pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py index 5024dad193..2fa1051c61 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py @@ -107,15 +107,8 @@ def line_intersect( line2_p2_lon: Field[[EdgeDim, KDim], float], line2_p2_lat: Field[[EdgeDim, KDim], float], ) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: - # avoid division with zero - d1 = line1_p2_lon - line1_p1_lon - d1 = where(d1 != float(0), d1, line1_p2_lon) - - d2 = line2_p2_lon - line2_p1_lon - d2 = where(d2 != float(0), d2, line2_p2_lon) - - m1 = (line1_p2_lat - line1_p1_lat) / d1 - m2 = (line2_p2_lat - line2_p1_lat) / d2 + m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon) + m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / (m1 - m2) intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py index 3d7a20bdce..7ea34b2978 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py @@ -25,7 +25,7 @@ def _hflx_limiter_mo_stencil_01a( ) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: z_mflx_low = 0.5 * ( p_mass_flx_e * (p_cc(E2C[0]) + p_cc(E2C[1])) - # - abs(p_mass_flx_e) * (p_cc(E2C[1]) - p_cc(E2C[0])) + - abs(p_mass_flx_e) * (p_cc(E2C[1]) - p_cc(E2C[0])) ) z_anti = p_mflx_tracer_h - z_mflx_low diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py index aeb1ff5db1..2380296e9e 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py @@ -12,31 +12,41 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, where +from gt4py.next.ffront.fbuiltins import Field, int32, where from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator def _hflx_limiter_pd_stencil_02( + refin_ctrl: Field[[EdgeDim], int32], r_m: Field[[CellDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + bound: int32, ) -> Field[[EdgeDim, KDim], float]: p_mflx_tracer_h_out = where( - p_mflx_tracer_h >= 0.0, - p_mflx_tracer_h * r_m(E2C[0]), - p_mflx_tracer_h * r_m(E2C[1]), + refin_ctrl == bound, + p_mflx_tracer_h, + where( + p_mflx_tracer_h >= 0.0, + p_mflx_tracer_h * r_m(E2C[0]), + p_mflx_tracer_h * r_m(E2C[1]), + ), ) return p_mflx_tracer_h_out @program def hflx_limiter_pd_stencil_02( + refin_ctrl: Field[[EdgeDim], int32], r_m: Field[[CellDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + bound: int32, ): _hflx_limiter_pd_stencil_02( + refin_ctrl, r_m, p_mflx_tracer_h, + bound, out=p_mflx_tracer_h, ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py index cd8069e55e..59ad061a9d 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py @@ -27,14 +27,14 @@ def _hor_adv_stencil_01( geofac_div: Field[[CEDim], float], p_dtime: float, ) -> Field[[CellDim, KDim], float]: - tracer_new_hor = ( + tracer_new = ( tracer_now * rhodz_now - p_dtime * deepatmo_divh * neighbor_sum(p_mflx_tracer_h(C2E) * geofac_div(C2CE), axis=C2EDim) ) / rhodz_new - return tracer_new_hor + return tracer_new @program @@ -45,7 +45,7 @@ def hor_adv_stencil_01( rhodz_now: Field[[CellDim, KDim], float], rhodz_new: Field[[CellDim, KDim], float], geofac_div: Field[[CEDim], float], - tracer_new_hor: Field[[CellDim, KDim], float], + tracer_new: Field[[CellDim, KDim], float], p_dtime: float, ): _hor_adv_stencil_01( @@ -56,5 +56,5 @@ def hor_adv_stencil_01( rhodz_new, geofac_div, p_dtime, - out=tracer_new_hor, + out=tracer_new, ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/init_zero_c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/init_zero_c.py new file mode 100644 index 0000000000..10fa5b4998 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/init_zero_c.py @@ -0,0 +1,28 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field, GridType +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import broadcast + +from icon4py.model.common.dimension import CellDim + + +@field_operator +def _init_zero_c() -> Field[[CellDim], float]: + return broadcast(0.0, (CellDim,)) + + +@program(grid_type=GridType.UNSTRUCTURED) +def init_zero_c(field: Field[[CellDim], float]): + _init_zero_c(out=field) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py index f35f0d4564..941cb99661 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py @@ -33,7 +33,6 @@ def _mo_advection_traj_btraj_compute_o1_dsl( dual_normal_cell_2: Field[[ECDim], wpfloat], p_dthalf: wpfloat, ) -> tuple[ - Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], vpfloat], @@ -42,7 +41,6 @@ def _mo_advection_traj_btraj_compute_o1_dsl( lvn_pos = where(p_vn > 0.0, True, False) p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1])) - p_cell_rel_idx_dsl = where(lvn_pos, int32(0), int32(1)) p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) z_ntdistv_bary_1 = -( @@ -69,13 +67,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl( + z_ntdistv_bary_2 * dual_normal_cell_2(E2EC[1]), ) - return ( - p_cell_idx, - p_cell_rel_idx_dsl, - p_cell_blk, - astype(p_distv_bary_1, vpfloat), - astype(p_distv_bary_2, vpfloat), - ) + return p_cell_idx, p_cell_blk, astype(p_distv_bary_1, vpfloat), astype(p_distv_bary_2, vpfloat) @program(grid_type=GridType.UNSTRUCTURED) @@ -91,7 +83,6 @@ def mo_advection_traj_btraj_compute_o1_dsl( primal_normal_cell_2: Field[[ECDim], wpfloat], dual_normal_cell_2: Field[[ECDim], wpfloat], p_cell_idx: Field[[EdgeDim, KDim], int32], - p_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], p_cell_blk: Field[[EdgeDim, KDim], int32], p_distv_bary_1: Field[[EdgeDim, KDim], vpfloat], p_distv_bary_2: Field[[EdgeDim, KDim], vpfloat], @@ -113,13 +104,7 @@ def mo_advection_traj_btraj_compute_o1_dsl( primal_normal_cell_2, dual_normal_cell_2, p_dthalf, - out=( - p_cell_idx, - p_cell_rel_idx_dsl, - p_cell_blk, - p_distv_bary_1, - p_distv_bary_2, - ), + out=(p_cell_idx, p_cell_blk, p_distv_bary_1, p_distv_bary_2), domain={ EdgeDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end), diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py index dfe82b7064..44f933b18f 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py @@ -12,13 +12,13 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where +from gt4py.next.ffront.fbuiltins import Field from icon4py.model.common.dimension import CellDim, KDim, Koff @field_operator -def _vert_adv_stencil_01a( +def _vert_adv_stencil_01( tracer_now: Field[[CellDim, KDim], float], rhodz_now: Field[[CellDim, KDim], float], p_mflx_tracer_v: Field[[CellDim, KDim], float], @@ -35,42 +35,6 @@ def _vert_adv_stencil_01a( return tracer_new -@field_operator -def _vert_adv_stencil_01( - tracer_now: Field[[CellDim, KDim], float], - rhodz_now: Field[[CellDim, KDim], float], - p_mflx_tracer_v: Field[[CellDim, KDim], float], - deepatmo_divzl: Field[[KDim], float], - deepatmo_divzu: Field[[KDim], float], - rhodz_new: Field[[CellDim, KDim], float], - k: Field[[KDim], int32], - p_dtime: float, - ivadv_tracer: int32, - iadv_slev_jt: int32, -) -> Field[[CellDim, KDim], float]: - k = broadcast(k, (CellDim, KDim)) - - tracer_new = ( - where( - (iadv_slev_jt <= k), - _vert_adv_stencil_01a( - tracer_now, - rhodz_now, - p_mflx_tracer_v, - deepatmo_divzl, - deepatmo_divzu, - rhodz_new, - p_dtime, - ), - tracer_now, - ) - if (ivadv_tracer != int32(0)) - else tracer_now - ) - - return tracer_new - - @program(grid_type=GridType.UNSTRUCTURED) def vert_adv_stencil_01( tracer_now: Field[[CellDim, KDim], float], @@ -79,11 +43,8 @@ def vert_adv_stencil_01( deepatmo_divzl: Field[[KDim], float], deepatmo_divzu: Field[[KDim], float], rhodz_new: Field[[CellDim, KDim], float], - k: Field[[KDim], int32], - p_dtime: float, - ivadv_tracer: int32, - iadv_slev_jt: int32, tracer_new: Field[[CellDim, KDim], float], + p_dtime: float, ): _vert_adv_stencil_01( tracer_now, @@ -92,9 +53,6 @@ def vert_adv_stencil_01( deepatmo_divzl, deepatmo_divzu, rhodz_new, - k, p_dtime, - ivadv_tracer, - iadv_slev_jt, out=tracer_new, ) diff --git a/model/atmosphere/advection/tests/stencil_tests/__init__.py b/model/atmosphere/advection/tests/stencil_tests/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_01.py new file mode 100644 index 0000000000..229ac10410 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_01.py @@ -0,0 +1,52 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_01 import btraj_dreg_stencil_01 +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestBtrajDregStencil01(StencilTest): + PROGRAM = btraj_dreg_stencil_01 + OUTPUTS = ("lvn_sys_pos",) + + @staticmethod + def reference( + grid, lcounterclock: bool, p_vn: np.array, tangent_orientation: np.array, **kwargs + ): + tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) + + tangent_orientation = np.broadcast_to(tangent_orientation, p_vn.shape) + + lvn_sys_pos_true = np.where(tangent_orientation * p_vn >= 0.0, True, False) + + mask_lcounterclock = np.broadcast_to(lcounterclock, p_vn.shape) + + lvn_sys_pos = np.where(mask_lcounterclock, lvn_sys_pos_true, False) + + return dict(lvn_sys_pos=lvn_sys_pos) + + @pytest.fixture + def input_data(self, grid): + lcounterclock = True + p_vn = random_field(grid, EdgeDim, KDim) + tangent_orientation = random_field(grid, EdgeDim) + lvn_sys_pos = zero_field(grid, EdgeDim, KDim, dtype=bool) + return dict( + lcounterclock=lcounterclock, + p_vn=p_vn, + tangent_orientation=tangent_orientation, + lvn_sys_pos=lvn_sys_pos, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_02.py new file mode 100644 index 0000000000..d602eae212 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_02.py @@ -0,0 +1,66 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_02 import btraj_dreg_stencil_02 +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + numpy_to_1D_sparse_field, + random_field, + zero_field, +) + + +class TestBtrajDregStencil02(StencilTest): + PROGRAM = btraj_dreg_stencil_02 + OUTPUTS = ("opt_famask_dsl",) + + @staticmethod + def reference(grid, p_vn: np.array, p_vt: np.array, p_dt: float, **kwargs): + lvn_pos = np.where(p_vn >= 0.0, True, False) + + traj_length = np.sqrt(p_vn**2 + p_vt**2) * p_dt + + edge_cell_length = np.expand_dims( + np.asarray(grid.connectivities[E2CDim], dtype=float), axis=-1 + ) + e2c_length = np.where(lvn_pos, edge_cell_length[:, 0], edge_cell_length[:, 1]) + + opt_famask_dsl = np.where( + traj_length > (1.25 * np.broadcast_to(e2c_length, p_vn.shape)), + int32(1), + int32(0), + ) + + return dict(opt_famask_dsl=opt_famask_dsl) + + @pytest.fixture + def input_data(self, grid): + p_vn = random_field(grid, EdgeDim, KDim) + p_vt = random_field(grid, EdgeDim, KDim) + edge_cell_length = np.asarray(grid.connectivities[E2CDim], dtype=float) + edge_cell_length_new = numpy_to_1D_sparse_field(edge_cell_length, ECDim) + p_dt = 1.0 + opt_famask_dsl = zero_field(grid, EdgeDim, KDim, dtype=int32) + + return dict( + p_vn=p_vn, + p_vt=p_vt, + edge_cell_length=edge_cell_length_new, + p_dt=p_dt, + opt_famask_dsl=opt_famask_dsl, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py new file mode 100644 index 0000000000..927bccd654 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py @@ -0,0 +1,233 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_03 import btraj_dreg_stencil_03 +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + constant_field, + numpy_to_1D_sparse_field, + random_field, + reshape, +) + + +class TestBtrajDregStencil03(StencilTest): + PROGRAM = btraj_dreg_stencil_03 + OUTPUTS = ( + "p_cell_idx", + "p_cell_rel_idx_dsl", + "p_cell_blk", + "p_coords_dreg_v_1_lon_dsl", + "p_coords_dreg_v_2_lon_dsl", + "p_coords_dreg_v_3_lon_dsl", + "p_coords_dreg_v_4_lon_dsl", + "p_coords_dreg_v_1_lat_dsl", + "p_coords_dreg_v_2_lat_dsl", + "p_coords_dreg_v_3_lat_dsl", + "p_coords_dreg_v_4_lat_dsl", + ) + + @staticmethod + def reference( + grid, + p_vn: np.array, + p_vt: np.array, + cell_idx: np.array, + cell_blk: np.array, + edge_verts_1_x: np.array, + edge_verts_2_x: np.array, + edge_verts_1_y: np.array, + edge_verts_2_y: np.array, + pos_on_tplane_e_1_x: np.array, + pos_on_tplane_e_2_x: np.array, + pos_on_tplane_e_1_y: np.array, + pos_on_tplane_e_2_y: np.array, + primal_normal_cell_x: np.array, + primal_normal_cell_y: np.array, + dual_normal_cell_x: np.array, + dual_normal_cell_y: np.array, + lvn_sys_pos: np.array, + p_dt: float, + **kwargs, + ): + e2c_shape = grid.connectivities[E2CDim].shape + cell_idx = reshape(cell_idx, e2c_shape) + cell_blk = reshape(cell_blk, e2c_shape) + primal_normal_cell_x = reshape(primal_normal_cell_x, e2c_shape) + dual_normal_cell_x = reshape(dual_normal_cell_x, e2c_shape) + primal_normal_cell_y = reshape(primal_normal_cell_y, e2c_shape) + dual_normal_cell_y = reshape(dual_normal_cell_y, e2c_shape) + + lvn_pos = np.where(p_vn >= 0.0, True, False) + cell_idx = np.expand_dims(cell_idx, axis=-1) + cell_blk = np.expand_dims(cell_blk, axis=-1) + primal_normal_cell_x = np.expand_dims(primal_normal_cell_x, axis=-1) + dual_normal_cell_x = np.expand_dims(dual_normal_cell_x, axis=-1) + primal_normal_cell_y = np.expand_dims(primal_normal_cell_y, axis=-1) + dual_normal_cell_y = np.expand_dims(dual_normal_cell_y, axis=-1) + edge_verts_1_x = np.expand_dims(edge_verts_1_x, axis=-1) + edge_verts_1_y = np.expand_dims(edge_verts_1_y, axis=-1) + edge_verts_2_x = np.expand_dims(edge_verts_2_x, axis=-1) + edge_verts_2_y = np.expand_dims(edge_verts_2_y, axis=-1) + pos_on_tplane_e_1_x = np.expand_dims(pos_on_tplane_e_1_x, axis=-1) + pos_on_tplane_e_1_y = np.expand_dims(pos_on_tplane_e_1_y, axis=-1) + pos_on_tplane_e_2_x = np.expand_dims(pos_on_tplane_e_2_x, axis=-1) + pos_on_tplane_e_2_y = np.expand_dims(pos_on_tplane_e_2_y, axis=-1) + + p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) + p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) + p_cell_rel_idx_dsl = np.where(lvn_pos, int32(0), int32(1)) + + depart_pts_1_x = np.broadcast_to(edge_verts_1_x, p_vn.shape) - p_vn * p_dt + depart_pts_1_y = np.broadcast_to(edge_verts_1_y, p_vn.shape) - p_vt * p_dt + depart_pts_2_x = np.broadcast_to(edge_verts_2_x, p_vn.shape) - p_vn * p_dt + depart_pts_2_y = np.broadcast_to(edge_verts_2_y, p_vn.shape) - p_vt * p_dt + + pos_on_tplane_e_x = np.where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) + pos_on_tplane_e_y = np.where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) + + pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x + pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y + pos_dreg_vert_c_2_x = ( + np.where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x + ) + pos_dreg_vert_c_2_y = ( + np.where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y + ) + pos_dreg_vert_c_3_x = depart_pts_2_x - pos_on_tplane_e_x + pos_dreg_vert_c_3_y = depart_pts_2_y - pos_on_tplane_e_y + pos_dreg_vert_c_4_x = ( + np.where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x + ) + pos_dreg_vert_c_4_y = ( + np.where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y + ) + + pn_cell_1 = np.where(lvn_pos, primal_normal_cell_x[:, 0], primal_normal_cell_x[:, 1]) + pn_cell_2 = np.where(lvn_pos, primal_normal_cell_y[:, 0], primal_normal_cell_y[:, 1]) + dn_cell_1 = np.where(lvn_pos, dual_normal_cell_x[:, 0], dual_normal_cell_x[:, 1]) + dn_cell_2 = np.where(lvn_pos, dual_normal_cell_y[:, 0], dual_normal_cell_y[:, 1]) + + p_coords_dreg_v_1_lon_dsl = ( + pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 + ) + p_coords_dreg_v_2_lon_dsl = ( + pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 + ) + p_coords_dreg_v_3_lon_dsl = ( + pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 + ) + p_coords_dreg_v_4_lon_dsl = ( + pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 + ) + p_coords_dreg_v_1_lat_dsl = ( + pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 + ) + p_coords_dreg_v_2_lat_dsl = ( + pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 + ) + p_coords_dreg_v_3_lat_dsl = ( + pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 + ) + p_coords_dreg_v_4_lat_dsl = ( + pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 + ) + return { + "p_cell_idx": p_cell_idx, + "p_cell_rel_idx_dsl": p_cell_rel_idx_dsl, + "p_cell_blk": p_cell_blk, + "p_coords_dreg_v_1_lon_dsl": p_coords_dreg_v_1_lon_dsl, + "p_coords_dreg_v_2_lon_dsl": p_coords_dreg_v_2_lon_dsl, + "p_coords_dreg_v_3_lon_dsl": p_coords_dreg_v_3_lon_dsl, + "p_coords_dreg_v_4_lon_dsl": p_coords_dreg_v_4_lon_dsl, + "p_coords_dreg_v_1_lat_dsl": p_coords_dreg_v_1_lat_dsl, + "p_coords_dreg_v_2_lat_dsl": p_coords_dreg_v_2_lat_dsl, + "p_coords_dreg_v_3_lat_dsl": p_coords_dreg_v_3_lat_dsl, + "p_coords_dreg_v_4_lat_dsl": p_coords_dreg_v_4_lat_dsl, + } + + @pytest.fixture + def input_data(self, grid): + p_vn = random_field(grid, EdgeDim, KDim) + p_vt = random_field(grid, EdgeDim, KDim) + cell_idx = np.asarray(grid.connectivities[E2CDim], dtype=int32) + cell_idx_new = numpy_to_1D_sparse_field(cell_idx, ECDim) + cell_blk = constant_field(grid, 1, EdgeDim, E2CDim, dtype=int32) + cell_blk_new = as_1D_sparse_field(cell_blk, ECDim) + + edge_verts_1_x = random_field(grid, EdgeDim) + edge_verts_2_x = random_field(grid, EdgeDim) + edge_verts_1_y = random_field(grid, EdgeDim) + edge_verts_2_y = random_field(grid, EdgeDim) + pos_on_tplane_e_1_x = random_field(grid, EdgeDim) + pos_on_tplane_e_2_x = random_field(grid, EdgeDim) + pos_on_tplane_e_1_y = random_field(grid, EdgeDim) + pos_on_tplane_e_2_y = random_field(grid, EdgeDim) + primal_normal_cell_x = random_field(grid, EdgeDim, E2CDim) + primal_normal_cell_x_new = as_1D_sparse_field(primal_normal_cell_x, ECDim) + dual_normal_cell_x = random_field(grid, EdgeDim, E2CDim) + dual_normal_cell_x_new = as_1D_sparse_field(dual_normal_cell_x, ECDim) + primal_normal_cell_y = random_field(grid, EdgeDim, E2CDim) + primal_normal_cell_y_new = as_1D_sparse_field(primal_normal_cell_y, ECDim) + dual_normal_cell_y = random_field(grid, EdgeDim, E2CDim) + dual_normal_cell_y_new = as_1D_sparse_field(dual_normal_cell_y, ECDim) + lvn_sys_pos = constant_field(grid, True, EdgeDim, KDim, dtype=bool) + p_dt = 2.0 + p_cell_idx = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) + p_cell_rel_idx_dsl = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) + p_cell_blk = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) + p_coords_dreg_v_1_lon_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_2_lon_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_3_lon_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_4_lon_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_1_lat_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_2_lat_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_3_lat_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_4_lat_dsl = random_field(grid, EdgeDim, KDim) + + return dict( + p_vn=p_vn, + p_vt=p_vt, + cell_idx=cell_idx_new, + cell_blk=cell_blk_new, + edge_verts_1_x=edge_verts_1_x, + edge_verts_2_x=edge_verts_2_x, + edge_verts_1_y=edge_verts_1_y, + edge_verts_2_y=edge_verts_2_y, + pos_on_tplane_e_1_x=pos_on_tplane_e_1_x, + pos_on_tplane_e_2_x=pos_on_tplane_e_2_x, + pos_on_tplane_e_1_y=pos_on_tplane_e_1_y, + pos_on_tplane_e_2_y=pos_on_tplane_e_2_y, + primal_normal_cell_x=primal_normal_cell_x_new, + primal_normal_cell_y=primal_normal_cell_y_new, + dual_normal_cell_x=dual_normal_cell_x_new, + dual_normal_cell_y=dual_normal_cell_y_new, + lvn_sys_pos=lvn_sys_pos, + p_dt=p_dt, + p_cell_idx=p_cell_idx, + p_cell_rel_idx_dsl=p_cell_rel_idx_dsl, + p_cell_blk=p_cell_blk, + p_coords_dreg_v_1_lon_dsl=p_coords_dreg_v_1_lon_dsl, + p_coords_dreg_v_2_lon_dsl=p_coords_dreg_v_2_lon_dsl, + p_coords_dreg_v_3_lon_dsl=p_coords_dreg_v_3_lon_dsl, + p_coords_dreg_v_4_lon_dsl=p_coords_dreg_v_4_lon_dsl, + p_coords_dreg_v_1_lat_dsl=p_coords_dreg_v_1_lat_dsl, + p_coords_dreg_v_2_lat_dsl=p_coords_dreg_v_2_lat_dsl, + p_coords_dreg_v_3_lat_dsl=p_coords_dreg_v_3_lat_dsl, + p_coords_dreg_v_4_lat_dsl=p_coords_dreg_v_4_lat_dsl, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py new file mode 100644 index 0000000000..392e01e4a1 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py @@ -0,0 +1,1544 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.divide_flux_area_list_stencil_01 import ( + divide_flux_area_list_stencil_01, +) +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + random_mask, + reshape, + zero_field, +) + + +# Check whether lines inters. +def _ccw_numpy( + p0_lon, + p0_lat, + p1_lon, + p1_lat, + p2_lon, + p2_lat, +): + dx1 = p1_lon - p0_lon + dy1 = p1_lat - p0_lat + + dx2 = p2_lon - p0_lon + dy2 = p2_lat - p0_lat + + dx1dy2 = dx1 * dy2 + dy1dx2 = dy1 * dx2 + + lccw = np.where(dx1dy2 > dy1dx2, True, False) + ccw_out = np.where(lccw, int32(1), int32(-1)) # 1: clockwise, -1: counterclockwise + return ccw_out + + +# Check whether two lines intersect +def _lintersect_numpy( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, +): + intersect1 = _ccw_numpy( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p1_lon, + line2_p1_lat, + ) * _ccw_numpy( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p2_lon, + line2_p2_lat, + ) + intersect2 = _ccw_numpy( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p1_lon, + line1_p1_lat, + ) * _ccw_numpy( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p2_lon, + line1_p2_lat, + ) + lintersect_out = np.where((intersect1 + intersect2) == -2, True, False) + + return lintersect_out + + +# Compute intersection point of two lines in 2D +def _line_intersect_numpy( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, +): + m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon) + m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) + + intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / (m1 - m2) + intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) + + return intersect_1, intersect_2 + + +@pytest.mark.slow_tests +class TestDivideFluxAreaListStencil01(StencilTest): + PROGRAM = divide_flux_area_list_stencil_01 + OUTPUTS = ( + "dreg_patch0_1_lon_dsl", + "dreg_patch0_1_lat_dsl", + "dreg_patch0_2_lon_dsl", + "dreg_patch0_2_lat_dsl", + "dreg_patch0_3_lon_dsl", + "dreg_patch0_3_lat_dsl", + "dreg_patch0_4_lon_dsl", + "dreg_patch0_4_lat_dsl", + "dreg_patch1_1_lon_vmask", + "dreg_patch1_1_lat_vmask", + "dreg_patch1_2_lon_vmask", + "dreg_patch1_2_lat_vmask", + "dreg_patch1_3_lon_vmask", + "dreg_patch1_3_lat_vmask", + "dreg_patch1_4_lon_vmask", + "dreg_patch1_4_lat_vmask", + "dreg_patch2_1_lon_vmask", + "dreg_patch2_1_lat_vmask", + "dreg_patch2_2_lon_vmask", + "dreg_patch2_2_lat_vmask", + "dreg_patch2_3_lon_vmask", + "dreg_patch2_3_lat_vmask", + "dreg_patch2_4_lon_vmask", + "dreg_patch2_4_lat_vmask", + ) + + @staticmethod + def _generate_flux_area_geometry( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + p_vn, + ptr_v3_lon_e, + ptr_v3_lat_e, + ): + arrival_pts_1_lon_dsl = dreg_patch0_1_lon_dsl + arrival_pts_1_lat_dsl = dreg_patch0_1_lat_dsl + arrival_pts_2_lon_dsl = dreg_patch0_2_lon_dsl + arrival_pts_2_lat_dsl = dreg_patch0_2_lat_dsl + depart_pts_1_lon_dsl = dreg_patch0_4_lon_dsl + depart_pts_1_lat_dsl = dreg_patch0_4_lat_dsl + depart_pts_2_lon_dsl = dreg_patch0_3_lon_dsl + depart_pts_2_lat_dsl = dreg_patch0_3_lat_dsl + + lvn_pos = np.where(p_vn >= 0.0, True, False) + + fl_line_p1_lon = depart_pts_1_lon_dsl + fl_line_p1_lat = depart_pts_1_lat_dsl + fl_line_p2_lon = depart_pts_2_lon_dsl + fl_line_p2_lat = depart_pts_2_lat_dsl + + tri_line1_p1_lon = arrival_pts_1_lon_dsl + tri_line1_p1_lat = arrival_pts_1_lat_dsl + tri_line1_p2_lon = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), + ) + tri_line1_p2_lat = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), + ) + + tri_line2_p1_lon = arrival_pts_2_lon_dsl + tri_line2_p1_lat = arrival_pts_2_lat_dsl + tri_line2_p2_lon = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), + ) + tri_line2_p2_lat = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), + ) + + return ( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ) + + @staticmethod + def _apply_case1_patch0( + mask_case1, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + ps1_x, + ps1_y, + ps2_x, + ps2_y, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ): + dreg_patch0_1_lon_dsl = arrival_pts_1_lon_dsl + dreg_patch0_1_lat_dsl = arrival_pts_1_lat_dsl + dreg_patch0_2_lon_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + arrival_pts_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + arrival_pts_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case1, ps2_x, depart_pts_2_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case1, ps2_y, depart_pts_2_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + depart_pts_1_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + depart_pts_1_lat_dsl, + ) + + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + @staticmethod + def _apply_case1_patch1( + mask_case1, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + ps1_x, + ps1_y, + ): + dreg_patch1_1_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_1_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_4_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_4_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_2_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 + ) + dreg_patch1_3_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 + ) + dreg_patch1_3_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 + ) + + return ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) + + @staticmethod + def _apply_case1_patch2( + mask_case1, + lvn_sys_pos, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ps2_x, + ps2_y, + ): + # Case 1 - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_1_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_4_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_4_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_2_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0 + ) + dreg_patch2_2_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0 + ) + dreg_patch2_3_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0 + ) + dreg_patch2_3_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0 + ) + + return ( + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + ) + + @staticmethod + def _apply_case2a_patch0( + mask_case2a, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + ps1_x, + ps1_y, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ): + dreg_patch0_1_lon_dsl = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + @staticmethod + def _apply_case2a_patch1( + mask_case2a, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + ps1_x, + ps1_y, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ): + dreg_patch1_1_lon_vmask = np.where( + mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask + ) + dreg_patch1_1_lat_vmask = np.where( + mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask + ) + dreg_patch1_4_lon_vmask = np.where( + mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask + ) + dreg_patch1_4_lat_vmask = np.where( + mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask + ) + dreg_patch1_2_lon_vmask = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), + dreg_patch1_2_lon_vmask, + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), + dreg_patch1_2_lat_vmask, + ) + dreg_patch1_3_lon_vmask = np.where( + mask_case2a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), + dreg_patch1_3_lon_vmask, + ) + dreg_patch1_3_lat_vmask = np.where( + mask_case2a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), + dreg_patch1_3_lat_vmask, + ) + + return ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) + + @staticmethod + def _apply_case2b_patch0( + mask_case2b, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + ps2_x, + ps2_y, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ): + dreg_patch0_1_lon_dsl = np.where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + @staticmethod + def _apply_case2b_patch1( + mask_case2b, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + ): + zeros_array = np.zeros_like(mask_case2b) + + dreg_patch1_1_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_4_lat_vmask) + + return ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + ) + + @staticmethod + def _apply_case2b_patch2( + mask_case2b, + lvn_sys_pos, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ps2_x, + ps2_y, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + ): + dreg_patch2_1_lon_vmask = np.where( + mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask + ) + dreg_patch2_1_lat_vmask = np.where( + mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask + ) + dreg_patch2_4_lon_vmask = np.where( + mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask + ) + dreg_patch2_4_lat_vmask = np.where( + mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask + ) + dreg_patch2_2_lon_vmask = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), + dreg_patch2_2_lon_vmask, + ) + dreg_patch2_2_lat_vmask = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), + dreg_patch2_2_lat_vmask, + ) + dreg_patch2_3_lon_vmask = np.where( + mask_case2b, + np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), + dreg_patch2_3_lon_vmask, + ) + dreg_patch2_3_lat_vmask = np.where( + mask_case2b, + np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), + dreg_patch2_3_lat_vmask, + ) + + return ( + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + ) + + @staticmethod + def _apply_case3a_patch0( + mask_case3a, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + lvn_sys_pos, + ps2_x, + ps2_y, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ): + dreg_patch0_1_lon_dsl = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + @staticmethod + def _apply_case3a_patch1( + mask_case3a, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + pi1_x, + pi1_y, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ): + dreg_patch1_1_lon_vmask = np.where( + mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask + ) + dreg_patch1_1_lat_vmask = np.where( + mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask + ) + dreg_patch1_2_lon_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), + dreg_patch1_2_lon_vmask, + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), + dreg_patch1_2_lat_vmask, + ) + dreg_patch1_3_lon_vmask = np.where( + mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask + ) + dreg_patch1_3_lat_vmask = np.where( + mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask + ) + dreg_patch1_4_lon_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), + dreg_patch1_4_lon_vmask, + ) + dreg_patch1_4_lat_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), + dreg_patch1_4_lat_vmask, + ) + + return ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) + + @staticmethod + def _apply_case3b_patch0( + mask_case3b, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + pi2_x, + pi2_y, + lvn_sys_pos, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + ): + dreg_patch0_1_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_4_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) + dreg_patch0_4_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), + dreg_patch0_3_lon_dsl, + ) + dreg_patch0_3_lat_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), + dreg_patch0_3_lat_dsl, + ) + + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + ) + + @staticmethod + def _apply_case3b_patch2( + mask_case3b, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + pi2_x, + pi2_y, + lvn_sys_pos, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + ): + dreg_patch2_1_lon_vmask = np.where( + mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask + ) + dreg_patch2_1_lat_vmask = np.where( + mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask + ) + dreg_patch2_2_lon_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), + dreg_patch2_2_lon_vmask, + ) + dreg_patch2_2_lat_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), + dreg_patch2_2_lat_vmask, + ) + dreg_patch2_3_lon_vmask = np.where( + mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask + ) + dreg_patch2_3_lat_vmask = np.where( + mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask + ) + dreg_patch2_4_lon_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), + dreg_patch2_4_lon_vmask, + ) + dreg_patch2_4_lat_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), + dreg_patch2_4_lat_vmask, + ) + + return ( + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + ) + + @classmethod + def reference( + cls, + grid, + famask_int, + p_vn, + ptr_v3_lon, + ptr_v3_lat, + tangent_orientation_dsl, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + **kwargs, + ): + e2c = grid.connectivities[E2CDim] + ptr_v3_lon = reshape(ptr_v3_lon, e2c.shape) + ptr_v3_lon_e = np.expand_dims(ptr_v3_lon, axis=-1) + ptr_v3_lat = reshape(ptr_v3_lat, e2c.shape) + ptr_v3_lat_e = np.expand_dims(ptr_v3_lat, axis=-1) + tangent_orientation_dsl = np.expand_dims(tangent_orientation_dsl, axis=-1) + + result_tuple = cls._generate_flux_area_geometry( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + p_vn, + ptr_v3_lon_e, + ptr_v3_lat_e, + ) + + ( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ) = result_tuple + + # Create first mask does departure-line segment intersects with A1V3 + lintersect_line1 = _lintersect_numpy( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + # Create first mask does departure-line segment intersects with A2V3 + lintersect_line2 = _lintersect_numpy( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + + lvn_sys_pos = np.where( + (p_vn * np.broadcast_to(tangent_orientation_dsl, p_vn.shape)) >= 0.0, + True, + False, + ) + famask_bool = np.where(famask_int == int32(1), True, False) + mask_case1 = np.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) + ps1_x, ps1_y = _line_intersect_numpy( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + ps2_x, ps2_y = _line_intersect_numpy( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + + # ------------------------------------------------- Case 1 + # Case 1 - patch 0 + ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) = cls._apply_case1_patch0( + mask_case1, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + ps1_x, + ps1_y, + ps2_x, + ps2_y, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ) + # Case 1 - patch 1 + ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) = cls._apply_case1_patch1( + mask_case1, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + ps1_x, + ps1_y, + ) + # Case 1 - patch 2 + ( + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + ) = cls._apply_case1_patch2( + mask_case1, + lvn_sys_pos, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ps2_x, + ps2_y, + ) + + # ------------------------------------------------- Case 2a + mask_case2a = np.logical_and.reduce( + [lintersect_line1, np.logical_not(lintersect_line2), famask_bool] + ) + # Case 2a - patch 0 + result_tuple_patch0 = cls._apply_case2a_patch0( + mask_case2a, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + ps1_x, + ps1_y, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) = result_tuple_patch0 + # Case 2a - patch 1 + result_tuple_patch1 = cls._apply_case2a_patch1( + mask_case2a, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + ps1_x, + ps1_y, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) + + ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) = result_tuple_patch1 + # Case 2a - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lat_vmask) + + # -------------------------------------------------- Case 2b + mask_case2b = np.logical_and.reduce( + [lintersect_line2, np.logical_not(lintersect_line1), famask_bool] + ) + # Case 2b - patch 0 + result_tuple_patch0_case2b = cls._apply_case2b_patch0( + mask_case2b, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + ps2_x, + ps2_y, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) = result_tuple_patch0_case2b + + # Case 2b - patch 1 + result_tuple_patch1_case2b = cls._apply_case2b_patch1( + mask_case2b, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + ) + + ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + ) = result_tuple_patch1_case2b + + # Case 2b - patch 2 + result_tuple_patch2_case2b = cls._apply_case2b_patch2( + mask_case2b, + lvn_sys_pos, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ps2_x, + ps2_y, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + ) + + ( + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + ) = result_tuple_patch2_case2b + + # flux area edge 1 and 2 + fl_e1_p1_lon = arrival_pts_1_lon_dsl + fl_e1_p1_lat = arrival_pts_1_lat_dsl + fl_e1_p2_lon = depart_pts_1_lon_dsl + fl_e1_p2_lat = depart_pts_1_lat_dsl + fl_e2_p1_lon = arrival_pts_2_lon_dsl + fl_e2_p1_lat = arrival_pts_2_lat_dsl + fl_e2_p2_lon = depart_pts_2_lon_dsl + fl_e2_p2_lat = depart_pts_2_lat_dsl + + # ----------------------------------------------- Case 3a + # Check whether flux area edge 2 intersects with triangle edge 1 + lintersect_e2_line1 = _lintersect_numpy( + fl_e2_p1_lon, + fl_e2_p1_lat, + fl_e2_p2_lon, + fl_e2_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + mask_case3a = np.logical_and(lintersect_e2_line1, famask_bool) + pi1_x, pi1_y = _line_intersect_numpy( + fl_e2_p1_lon, + fl_e2_p1_lat, + fl_e2_p2_lon, + fl_e2_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + # Case 3a - patch 0 + result = cls._apply_case3a_patch0( + mask_case3a, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + lvn_sys_pos, + ps2_x, + ps2_y, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) = result + + # Case 3a - patch 1 + ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) = cls._apply_case3a_patch1( + mask_case3a, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + pi1_x, + pi1_y, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) + # Case 3a - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lat_vmask) + + # ------------------------------------------------ Case 3b + # Check whether flux area edge 1 intersects with triangle edge 2 + lintersect_e1_line2 = _lintersect_numpy( + fl_e1_p1_lon, + fl_e1_p1_lat, + fl_e1_p2_lon, + fl_e1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + mask_case3b = lintersect_e1_line2 & famask_bool + pi2_x, pi2_y = _line_intersect_numpy( + fl_e1_p1_lon, + fl_e1_p1_lat, + fl_e1_p2_lon, + fl_e1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + # Case 3b - patch 0 + ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + ) = cls._apply_case3b_patch0( + mask_case3b, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + pi2_x, + pi2_y, + lvn_sys_pos, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + ) + # Case 3b - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) + + # Case 3b - patch 2 + ( + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + ) = cls._apply_case3b_patch2( + mask_case3b, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + pi2_x, + pi2_y, + lvn_sys_pos, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + ) + # --------------------------------------------- Case 4 + # NB: Next line acts as the "ELSE IF", indices that already previously matched one of the above conditions + # can't be overwritten by this new condition. + indices_previously_matched = np.logical_or.reduce( + [mask_case3b, mask_case3a, mask_case2b, mask_case2a, mask_case1] + ) + # mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices + mask_case4 = np.logical_and.reduce( + [famask_bool, np.logical_not(indices_previously_matched)] + ) + # Case 4 - patch 0 - no change + # Case 4 - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lat_vmask) + # Case 4 - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_4_lon_vmask) + + return dict( + dreg_patch0_1_lon_dsl=dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl=dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl=dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl=dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl=dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl=dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl=dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl=dreg_patch0_4_lat_dsl, + dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, + ) + + @pytest.fixture + def input_data(self, grid): + famask_int = random_mask(grid, EdgeDim, KDim, dtype=int32) + p_vn = random_field(grid, EdgeDim, KDim) + ptr_v3_lon = random_field(grid, EdgeDim, E2CDim) + ptr_v3_lon_field = as_1D_sparse_field(ptr_v3_lon, ECDim) + ptr_v3_lat = random_field(grid, EdgeDim, E2CDim) + ptr_v3_lat_field = as_1D_sparse_field(ptr_v3_lat, ECDim) + tangent_orientation_dsl = random_field(grid, EdgeDim) + dreg_patch0_1_lon_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_1_lat_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_2_lon_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_2_lat_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_3_lon_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_3_lat_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_4_lon_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_4_lat_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch1_1_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_1_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_2_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_2_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_3_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_3_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_4_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_4_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_1_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_1_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_2_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_2_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_3_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_3_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_4_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_4_lat_vmask = zero_field(grid, EdgeDim, KDim) + return dict( + famask_int=famask_int, + p_vn=p_vn, + ptr_v3_lon=ptr_v3_lon_field, + ptr_v3_lat=ptr_v3_lat_field, + tangent_orientation_dsl=tangent_orientation_dsl, + dreg_patch0_1_lon_dsl=dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl=dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl=dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl=dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl=dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl=dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl=dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl=dreg_patch0_4_lat_dsl, + dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_02.py new file mode 100644 index 0000000000..6b1daabc13 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_02.py @@ -0,0 +1,290 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.divide_flux_area_list_stencil_02 import ( + divide_flux_area_list_stencil_02, +) +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + random_mask, + reshape, +) + + +class TestDivideFluxAreaListStencil02(StencilTest): + PROGRAM = divide_flux_area_list_stencil_02 + OUTPUTS = ( + "dreg_patch1_1_lon_vmask", + "dreg_patch1_1_lat_vmask", + "dreg_patch1_2_lon_vmask", + "dreg_patch1_2_lat_vmask", + "dreg_patch1_3_lon_vmask", + "dreg_patch1_3_lat_vmask", + "dreg_patch1_4_lon_vmask", + "dreg_patch1_4_lat_vmask", + "dreg_patch2_1_lon_vmask", + "dreg_patch2_1_lat_vmask", + "dreg_patch2_2_lon_vmask", + "dreg_patch2_2_lat_vmask", + "dreg_patch2_3_lon_vmask", + "dreg_patch2_3_lat_vmask", + "dreg_patch2_4_lon_vmask", + "dreg_patch2_4_lat_vmask", + "patch1_cell_idx_vmask", + "patch1_cell_blk_vmask", + "patch2_cell_idx_vmask", + "patch2_cell_blk_vmask", + ) + + @staticmethod + def reference( + grid, + famask_int: np.array, + p_vn: np.array, + bf_cc_patch1_lon: np.array, + bf_cc_patch1_lat: np.array, + bf_cc_patch2_lon: np.array, + bf_cc_patch2_lat: np.array, + butterfly_idx_patch1_vnpos: np.array, + butterfly_idx_patch1_vnneg: np.array, + butterfly_blk_patch1_vnpos: np.array, + butterfly_blk_patch1_vnneg: np.array, + butterfly_idx_patch2_vnpos: np.array, + butterfly_idx_patch2_vnneg: np.array, + butterfly_blk_patch2_vnpos: np.array, + butterfly_blk_patch2_vnneg: np.array, + dreg_patch1_1_lon_vmask: np.array, + dreg_patch1_1_lat_vmask: np.array, + dreg_patch1_2_lon_vmask: np.array, + dreg_patch1_2_lat_vmask: np.array, + dreg_patch1_3_lon_vmask: np.array, + dreg_patch1_3_lat_vmask: np.array, + dreg_patch1_4_lon_vmask: np.array, + dreg_patch1_4_lat_vmask: np.array, + dreg_patch2_1_lon_vmask: np.array, + dreg_patch2_1_lat_vmask: np.array, + dreg_patch2_2_lon_vmask: np.array, + dreg_patch2_2_lat_vmask: np.array, + dreg_patch2_3_lon_vmask: np.array, + dreg_patch2_3_lat_vmask: np.array, + dreg_patch2_4_lon_vmask: np.array, + dreg_patch2_4_lat_vmask: np.array, + **kwargs, + ): + e2c = grid.connectivities[E2CDim] + famask_bool = np.where(famask_int == int32(1), True, False) + lvn_pos = np.where(p_vn >= np.broadcast_to(0.0, p_vn.shape), True, False) + # Translation of patch 1 and patch 2 in system relative to respective cell + bf_cc_patch1_lon = reshape(bf_cc_patch1_lon, e2c.shape) + bf_cc_patch1_lon_e = np.expand_dims(bf_cc_patch1_lon, axis=-1) + bf_cc_patch1_lat = reshape(bf_cc_patch1_lat, e2c.shape) + bf_cc_patch1_lat_e = np.expand_dims(bf_cc_patch1_lat, axis=-1) + bf_cc_patch2_lon = reshape(bf_cc_patch2_lon, e2c.shape) + bf_cc_patch2_lon_e = np.expand_dims(bf_cc_patch2_lon, axis=-1) + bf_cc_patch2_lat = reshape(bf_cc_patch2_lat, e2c.shape) + bf_cc_patch2_lat_e = np.expand_dims(bf_cc_patch2_lat, axis=-1) + + bf_cc_patch1_lon = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch1_lon_e[:, 0], bf_cc_patch1_lon_e[:, 1]), + 0.0, + ) + bf_cc_patch1_lat = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch1_lat_e[:, 0], bf_cc_patch1_lat_e[:, 1]), + 0.0, + ) + bf_cc_patch2_lon = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch2_lon_e[:, 0], bf_cc_patch2_lon_e[:, 1]), + 0.0, + ) + bf_cc_patch2_lat = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch2_lat_e[:, 0], bf_cc_patch2_lat_e[:, 1]), + 0.0, + ) + + # patch1 in translated system + dreg_patch1_1_lon_vmask = dreg_patch1_1_lon_vmask - bf_cc_patch1_lon + dreg_patch1_1_lat_vmask = dreg_patch1_1_lat_vmask - bf_cc_patch1_lat + dreg_patch1_2_lon_vmask = dreg_patch1_2_lon_vmask - bf_cc_patch1_lon + dreg_patch1_2_lat_vmask = dreg_patch1_2_lat_vmask - bf_cc_patch1_lat + dreg_patch1_3_lon_vmask = dreg_patch1_3_lon_vmask - bf_cc_patch1_lon + dreg_patch1_3_lat_vmask = dreg_patch1_3_lat_vmask - bf_cc_patch1_lat + dreg_patch1_4_lon_vmask = dreg_patch1_4_lon_vmask - bf_cc_patch1_lon + dreg_patch1_4_lat_vmask = dreg_patch1_4_lat_vmask - bf_cc_patch1_lat + # patch2 in translated system + dreg_patch2_1_lon_vmask = dreg_patch2_1_lon_vmask - bf_cc_patch2_lon + dreg_patch2_1_lat_vmask = dreg_patch2_1_lat_vmask - bf_cc_patch2_lat + dreg_patch2_2_lon_vmask = dreg_patch2_2_lon_vmask - bf_cc_patch2_lon + dreg_patch2_2_lat_vmask = dreg_patch2_2_lat_vmask - bf_cc_patch2_lat + dreg_patch2_3_lon_vmask = dreg_patch2_3_lon_vmask - bf_cc_patch2_lon + dreg_patch2_3_lat_vmask = dreg_patch2_3_lat_vmask - bf_cc_patch2_lat + dreg_patch2_4_lon_vmask = dreg_patch2_4_lon_vmask - bf_cc_patch2_lon + dreg_patch2_4_lat_vmask = dreg_patch2_4_lat_vmask - bf_cc_patch2_lat + + # Store global index of the underlying grid cell + # Adapt dimensions to fit ofr multiple levels + butterfly_idx_patch1_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch1_vnpos, axis=-1), p_vn.shape + ) + butterfly_idx_patch1_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch1_vnneg, axis=-1), p_vn.shape + ) + butterfly_idx_patch2_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch2_vnpos, axis=-1), p_vn.shape + ) + butterfly_idx_patch2_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch2_vnneg, axis=-1), p_vn.shape + ) + butterfly_blk_patch1_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch1_vnpos, axis=-1), p_vn.shape + ) + butterfly_blk_patch1_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch1_vnneg, axis=-1), p_vn.shape + ) + butterfly_blk_patch2_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch2_vnpos, axis=-1), p_vn.shape + ) + butterfly_blk_patch2_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch2_vnneg, axis=-1), p_vn.shape + ) + patch1_cell_idx_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_idx_patch1_vnpos_3d, butterfly_idx_patch1_vnneg_3d), + int32(0), + ) + patch2_cell_idx_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_idx_patch2_vnpos_3d, butterfly_idx_patch2_vnneg_3d), + int32(0), + ) + patch1_cell_blk_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_blk_patch1_vnpos_3d, butterfly_blk_patch1_vnneg_3d), + int32(0), + ) + patch2_cell_blk_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_blk_patch2_vnpos_3d, butterfly_blk_patch2_vnneg_3d), + int32(0), + ) + + return dict( + dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, + patch1_cell_idx_vmask=patch1_cell_idx_vmask, + patch1_cell_blk_vmask=patch1_cell_blk_vmask, + patch2_cell_idx_vmask=patch2_cell_idx_vmask, + patch2_cell_blk_vmask=patch2_cell_blk_vmask, + ) + + @pytest.fixture + def input_data(self, grid): + famask_int = random_mask(grid, EdgeDim, KDim, dtype=int32) + p_vn = random_field(grid, EdgeDim, KDim) + bf_cc_patch1_lon = random_field(grid, EdgeDim, E2CDim) + bf_cc_patch1_lon_field = as_1D_sparse_field(bf_cc_patch1_lon, ECDim) + bf_cc_patch1_lat = random_field(grid, EdgeDim, E2CDim) + bf_cc_patch1_lat_field = as_1D_sparse_field(bf_cc_patch1_lat, ECDim) + bf_cc_patch2_lon = random_field(grid, EdgeDim, E2CDim) + bf_cc_patch2_lon_field = as_1D_sparse_field(bf_cc_patch2_lon, ECDim) + bf_cc_patch2_lat = random_field(grid, EdgeDim, E2CDim) + bf_cc_patch2_lat_field = as_1D_sparse_field(bf_cc_patch2_lat, ECDim) + butterfly_idx_patch1_vnpos = random_mask(grid, EdgeDim, dtype=int32) + butterfly_idx_patch1_vnneg = random_mask(grid, EdgeDim, dtype=int32) + butterfly_blk_patch1_vnpos = random_mask(grid, EdgeDim, dtype=int32) + butterfly_blk_patch1_vnneg = random_mask(grid, EdgeDim, dtype=int32) + butterfly_idx_patch2_vnpos = random_mask(grid, EdgeDim, dtype=int32) + butterfly_idx_patch2_vnneg = random_mask(grid, EdgeDim, dtype=int32) + butterfly_blk_patch2_vnpos = random_mask(grid, EdgeDim, dtype=int32) + butterfly_blk_patch2_vnneg = random_mask(grid, EdgeDim, dtype=int32) + dreg_patch1_1_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_1_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_2_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_2_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_3_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_3_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_4_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_4_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_1_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_1_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_2_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_2_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_3_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_3_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_4_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_4_lat_vmask = random_field(grid, EdgeDim, KDim) + patch1_cell_idx_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) + patch1_cell_blk_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) + patch2_cell_idx_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) + patch2_cell_blk_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) + + return dict( + famask_int=famask_int, + p_vn=p_vn, + bf_cc_patch1_lon=bf_cc_patch1_lon_field, + bf_cc_patch1_lat=bf_cc_patch1_lat_field, + bf_cc_patch2_lon=bf_cc_patch2_lon_field, + bf_cc_patch2_lat=bf_cc_patch2_lat_field, + butterfly_idx_patch1_vnpos=butterfly_idx_patch1_vnpos, + butterfly_idx_patch1_vnneg=butterfly_idx_patch1_vnneg, + butterfly_blk_patch1_vnpos=butterfly_blk_patch1_vnpos, + butterfly_blk_patch1_vnneg=butterfly_blk_patch1_vnneg, + butterfly_idx_patch2_vnpos=butterfly_idx_patch2_vnpos, + butterfly_idx_patch2_vnneg=butterfly_idx_patch2_vnneg, + butterfly_blk_patch2_vnpos=butterfly_blk_patch2_vnpos, + butterfly_blk_patch2_vnneg=butterfly_blk_patch2_vnneg, + dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, + patch1_cell_idx_vmask=patch1_cell_idx_vmask, + patch1_cell_blk_vmask=patch1_cell_blk_vmask, + patch2_cell_idx_vmask=patch2_cell_idx_vmask, + patch2_cell_blk_vmask=patch2_cell_blk_vmask, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py new file mode 100644 index 0000000000..222a53e035 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py @@ -0,0 +1,78 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import numpy as np +import pytest +from gt4py.next import as_field +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + Output, + StencilTest, + _shape, + random_field, + zero_field, +) + + +class TestFaceValPpmStencil01(StencilTest): + PROGRAM = face_val_ppm_stencil_01 + OUTPUTS = ( + Output( + "z_slope", refslice=(slice(None), slice(None, -1)), gtslice=(slice(None), slice(1, -1)) + ), + ) + + @staticmethod + def reference( + grid, p_cc: np.array, p_cellhgt_mc_now: np.array, k: np.array, elev: int32, **kwargs + ): + zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] + ) + zfac = (p_cc[:, 2:] - p_cc[:, 1:-1]) / (p_cellhgt_mc_now[:, 2:] + p_cellhgt_mc_now[:, 1:-1]) + z_slope_a = ( + p_cellhgt_mc_now[:, 1:-1] + / (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 2:]) + ) * ( + (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac + + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 2:]) * zfac_m1 + ) + + zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] + ) + zfac = (p_cc[:, 1:-1] - p_cc[:, 1:-1]) / ( + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1] + ) + z_slope_b = ( + p_cellhgt_mc_now[:, 1:-1] + / (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1]) + ) * ( + (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac + + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 1:-1]) * zfac_m1 + ) + + z_slope = np.where(k[1:-1] < elev, z_slope_a, z_slope_b) + return dict(z_slope=z_slope) + + @pytest.fixture + def input_data(self, grid): + z_slope = zero_field(grid, CellDim, KDim) + p_cc = random_field(grid, CellDim, KDim, extend={KDim: 1}) + p_cellhgt_mc_now = random_field(grid, CellDim, KDim, extend={KDim: 1}) + k = as_field((KDim,), np.arange(0, _shape(grid, KDim, extend={KDim: 1})[0], dtype=int32)) + elev = k[-2].as_scalar() + + return dict(p_cc=p_cc, p_cellhgt_mc_now=p_cellhgt_mc_now, k=k, elev=elev, z_slope=z_slope) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py new file mode 100644 index 0000000000..b9f27fb7fe --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py @@ -0,0 +1,76 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next import as_field +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02 import face_val_ppm_stencil_02 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, _shape, random_field, zero_field + + +class TestFaceValPpmStencil02(StencilTest): + PROGRAM = face_val_ppm_stencil_02 + OUTPUTS = ("p_face",) + + @staticmethod + def reference( + grid, + p_cc: np.array, + p_cellhgt_mc_now: np.array, + p_face_in: np.array, + k: np.array, + slev: int32, + elev: int32, + slevp1: int32, + elevp1: int32, + **kwargs, + ): + p_face_a = p_face_in + p_face_a[:, 1:] = p_cc[:, 1:] * ( + 1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) + ) + (p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:])) * ( + (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1] + ) + + p_face = np.where((k == slevp1) | (k == elev), p_face_a, p_face_in) + p_face = np.where((k == slev), p_cc, p_face) + p_face[:, 1:] = np.where((k[1:] == elevp1), p_cc[:, :-1], p_face[:, 1:]) + return dict(p_face=p_face) + + @pytest.fixture + def input_data(self, grid): + p_cc = random_field(grid, CellDim, KDim) + p_cellhgt_mc_now = random_field(grid, CellDim, KDim) + p_face_in = random_field(grid, CellDim, KDim) + p_face = zero_field(grid, CellDim, KDim) + + k = as_field((KDim,), np.arange(0, _shape(grid, KDim)[0], dtype=int32)) + slev = int32(1) + slevp1 = slev + int32(1) + elev = int32(k[-3].as_scalar()) + elevp1 = elev + int32(1) + + return dict( + p_cc=p_cc, + p_cellhgt_mc_now=p_cellhgt_mc_now, + p_face_in=p_face_in, + k=k, + slev=slev, + elev=elev, + slevp1=slevp1, + elevp1=elevp1, + p_face=p_face, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py new file mode 100644 index 0000000000..d1d8d5d900 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py @@ -0,0 +1,43 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02a import face_val_ppm_stencil_02a +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field + + +outslice = (slice(None), slice(1, None)) + + +class TestFaceValPpmStencil02a(StencilTest): + PROGRAM = face_val_ppm_stencil_02a + OUTPUTS = (Output("p_face", refslice=outslice, gtslice=outslice),) + + @staticmethod + def reference(grid, p_cc: np.array, p_cellhgt_mc_now: np.array, **kwargs): + p_face = p_cc.copy() + p_face[:, 1:] = p_cc[:, 1:] * ( + 1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) + ) + (p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:])) * ( + (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1] + ) + return dict(p_face=p_face) + + @pytest.fixture + def input_data(self, grid): + p_face = random_field(grid, CellDim, KDim) + p_cc = random_field(grid, CellDim, KDim) + p_cellhgt_mc_now = random_field(grid, CellDim, KDim) + return dict(p_cc=p_cc, p_cellhgt_mc_now=p_cellhgt_mc_now, p_face=p_face) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02b.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02b.py new file mode 100644 index 0000000000..1a0a540106 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02b.py @@ -0,0 +1,35 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02b import face_val_ppm_stencil_02b +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestFaceValPpmStencil02b(StencilTest): + PROGRAM = face_val_ppm_stencil_02b + OUTPUTS = ("p_face",) + + @staticmethod + def reference(grid, p_cc: np.array, **kwargs): + p_face = p_cc.copy() + return dict(p_face=p_face) + + @pytest.fixture + def input_data(self, grid): + p_cc = random_field(grid, CellDim, KDim) + p_face = random_field(grid, CellDim, KDim) + return dict(p_cc=p_cc, p_face=p_face) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02c.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02c.py new file mode 100644 index 0000000000..bdc5529eaa --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02c.py @@ -0,0 +1,52 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02c import face_val_ppm_stencil_02c +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestFaceValPpmStencil02c(StencilTest): + PROGRAM = face_val_ppm_stencil_02c + OUTPUTS = ("p_face",) + + @staticmethod + def reference( + grid, + p_cc: np.array, + p_face: np.array, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, + ): + subset = (slice(horizontal_start, horizontal_end), slice(vertical_start, vertical_end)) + p_face = p_face.copy() + p_face[subset] = np.roll(p_cc, shift=1, axis=1)[subset] + return dict(p_face=p_face) + + @pytest.fixture + def input_data(self, grid): + p_cc = random_field(grid, CellDim, KDim) + p_face = random_field(grid, CellDim, KDim) + return dict( + p_cc=p_cc, + p_face=p_face, + horizontal_start=int32(0), + horizontal_end=int32(grid.num_cells), + vertical_start=int32(1), + vertical_end=int32(grid.num_levels), + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py new file mode 100644 index 0000000000..6d7115b942 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py @@ -0,0 +1,71 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_05 import face_val_ppm_stencil_05 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field, zero_field + + +class TestFaceValPpmStencil05(StencilTest): + PROGRAM = face_val_ppm_stencil_05 + OUTPUTS = (Output("p_face", gtslice=(slice(None), slice(2, None))),) + + @staticmethod + def reference(grid, p_cc: np.array, p_cellhgt_mc_now: np.array, z_slope: np.array, **kwargs): + p_cellhgt_mc_now_k_minus_1 = p_cellhgt_mc_now[:, 1:-2] + p_cellhgt_mc_now_k_minus_2 = p_cellhgt_mc_now[:, 0:-3] + p_cellhgt_mc_now_k_plus_1 = p_cellhgt_mc_now[:, 3:] + p_cellhgt_mc_now = p_cellhgt_mc_now[:, 2:-1] + + p_cc_k_minus_1 = p_cc[:, 1:-1] + p_cc = p_cc[:, 2:] + z_slope_k_minus_1 = z_slope[:, 1:-1] + z_slope = z_slope[:, 2:] + + zgeo1 = p_cellhgt_mc_now_k_minus_1 / (p_cellhgt_mc_now_k_minus_1 + p_cellhgt_mc_now) + zgeo2 = 1.0 / ( + p_cellhgt_mc_now_k_minus_2 + + p_cellhgt_mc_now_k_minus_1 + + p_cellhgt_mc_now + + p_cellhgt_mc_now_k_plus_1 + ) + zgeo3 = (p_cellhgt_mc_now_k_minus_2 + p_cellhgt_mc_now_k_minus_1) / ( + 2.0 * p_cellhgt_mc_now_k_minus_1 + p_cellhgt_mc_now + ) + zgeo4 = (p_cellhgt_mc_now_k_plus_1 + p_cellhgt_mc_now) / ( + 2 * p_cellhgt_mc_now + p_cellhgt_mc_now_k_minus_1 + ) + + p_face = ( + p_cc_k_minus_1 + + zgeo1 * (p_cc - p_cc_k_minus_1) + + zgeo2 + * ( + (2 * p_cellhgt_mc_now * zgeo1) * (zgeo3 - zgeo4) * (p_cc - p_cc_k_minus_1) + - zgeo3 * p_cellhgt_mc_now_k_minus_1 * z_slope + + zgeo4 * p_cellhgt_mc_now * z_slope_k_minus_1 + ) + ) + return dict(p_face=p_face) + + @pytest.fixture + def input_data(self, grid): + p_cc = random_field(grid, CellDim, KDim) + p_cellhgt_mc_now = random_field(grid, CellDim, KDim, extend={KDim: 1}) + z_slope = random_field(grid, CellDim, KDim) + p_face = zero_field(grid, CellDim, KDim) + + return dict(p_cc=p_cc, p_cellhgt_mc_now=p_cellhgt_mc_now, z_slope=z_slope, p_face=p_face) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py new file mode 100644 index 0000000000..acfa579ff3 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py @@ -0,0 +1,184 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_01a import ( + hflux_ffsl_hybrid_stencil_01a, +) +from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + constant_field, + random_field, + zero_field, +) + + +class TestHfluxFfslHybridStencil01a(StencilTest): + PROGRAM = hflux_ffsl_hybrid_stencil_01a + OUTPUTS = ("p_out_e_hybrid_1a",) + + @staticmethod + def reference( + grid, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, + z_lsq_coeff_4: np.array, + z_lsq_coeff_5: np.array, + z_lsq_coeff_6: np.array, + z_lsq_coeff_7: np.array, + z_lsq_coeff_8: np.array, + z_lsq_coeff_9: np.array, + z_lsq_coeff_10: np.array, + z_quad_vector_sum0_1: np.array, + z_quad_vector_sum0_2: np.array, + z_quad_vector_sum0_3: np.array, + z_quad_vector_sum0_4: np.array, + z_quad_vector_sum0_5: np.array, + z_quad_vector_sum0_6: np.array, + z_quad_vector_sum0_7: np.array, + z_quad_vector_sum0_8: np.array, + z_quad_vector_sum0_9: np.array, + z_quad_vector_sum0_10: np.array, + patch0_cell_rel_idx_dsl: np.array, + **kwargs, + ): + e2c = grid.connectivities[E2CDim] + z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] + z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] + z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] + z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] + z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] + z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] + z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] + z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] + z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] + z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] + + p_out_e_hybrid_1a = ( + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_e2c[:, 1], + z_lsq_coeff_1_e2c[:, 0], + ) + * z_quad_vector_sum0_1 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_e2c[:, 1], + z_lsq_coeff_2_e2c[:, 0], + ) + * z_quad_vector_sum0_2 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_e2c[:, 1], + z_lsq_coeff_3_e2c[:, 0], + ) + * z_quad_vector_sum0_3 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_4_e2c[:, 1], + z_lsq_coeff_4_e2c[:, 0], + ) + * z_quad_vector_sum0_4 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_5_e2c[:, 1], + z_lsq_coeff_5_e2c[:, 0], + ) + * z_quad_vector_sum0_5 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_6_e2c[:, 1], + z_lsq_coeff_6_e2c[:, 0], + ) + * z_quad_vector_sum0_6 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_7_e2c[:, 1], + z_lsq_coeff_7_e2c[:, 0], + ) + * z_quad_vector_sum0_7 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_8_e2c[:, 1], + z_lsq_coeff_8_e2c[:, 0], + ) + * z_quad_vector_sum0_8 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_9_e2c[:, 1], + z_lsq_coeff_9_e2c[:, 0], + ) + * z_quad_vector_sum0_9 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_10_e2c[:, 1], + z_lsq_coeff_10_e2c[:, 0], + ) + * z_quad_vector_sum0_10 + ) + + return dict(p_out_e_hybrid_1a=p_out_e_hybrid_1a) + + @pytest.fixture + def input_data(self, grid): + z_lsq_coeff_1 = random_field(grid, CellDim, KDim) + z_lsq_coeff_2 = random_field(grid, CellDim, KDim) + z_lsq_coeff_3 = random_field(grid, CellDim, KDim) + z_lsq_coeff_4 = random_field(grid, CellDim, KDim) + z_lsq_coeff_5 = random_field(grid, CellDim, KDim) + z_lsq_coeff_6 = random_field(grid, CellDim, KDim) + z_lsq_coeff_7 = random_field(grid, CellDim, KDim) + z_lsq_coeff_8 = random_field(grid, CellDim, KDim) + z_lsq_coeff_9 = random_field(grid, CellDim, KDim) + z_lsq_coeff_10 = random_field(grid, CellDim, KDim) + z_quad_vector_sum0_1 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_2 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_3 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_4 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_5 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_6 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_7 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_8 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_9 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_10 = random_field(grid, EdgeDim, KDim) + patch0_cell_rel_idx_dsl = constant_field(grid, 1, EdgeDim, KDim, dtype=int32) + p_out_e_hybrid_1a = zero_field(grid, EdgeDim, KDim) + return dict( + z_lsq_coeff_1=z_lsq_coeff_1, + z_lsq_coeff_2=z_lsq_coeff_2, + z_lsq_coeff_3=z_lsq_coeff_3, + z_lsq_coeff_4=z_lsq_coeff_4, + z_lsq_coeff_5=z_lsq_coeff_5, + z_lsq_coeff_6=z_lsq_coeff_6, + z_lsq_coeff_7=z_lsq_coeff_7, + z_lsq_coeff_8=z_lsq_coeff_8, + z_lsq_coeff_9=z_lsq_coeff_9, + z_lsq_coeff_10=z_lsq_coeff_10, + z_quad_vector_sum0_1=z_quad_vector_sum0_1, + z_quad_vector_sum0_2=z_quad_vector_sum0_2, + z_quad_vector_sum0_3=z_quad_vector_sum0_3, + z_quad_vector_sum0_4=z_quad_vector_sum0_4, + z_quad_vector_sum0_5=z_quad_vector_sum0_5, + z_quad_vector_sum0_6=z_quad_vector_sum0_6, + z_quad_vector_sum0_7=z_quad_vector_sum0_7, + z_quad_vector_sum0_8=z_quad_vector_sum0_8, + z_quad_vector_sum0_9=z_quad_vector_sum0_9, + z_quad_vector_sum0_10=z_quad_vector_sum0_10, + patch0_cell_rel_idx_dsl=patch0_cell_rel_idx_dsl, + p_out_e_hybrid_1a=p_out_e_hybrid_1a, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py new file mode 100644 index 0000000000..5c57bf79b9 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py @@ -0,0 +1,49 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_02 import ( + hflux_ffsl_hybrid_stencil_02, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestHfluxFfslHybridStencil02(StencilTest): + PROGRAM = hflux_ffsl_hybrid_stencil_02 + OUTPUTS = ("p_out_e_hybrid_2",) + + @staticmethod + def reference( + grid, + p_out_e_hybrid_2: np.array, + p_mass_flx_e: np.array, + z_dreg_area: np.array, + **kwargs, + ): + p_out_e_hybrid_2 = p_mass_flx_e * p_out_e_hybrid_2 / z_dreg_area + + return dict(p_out_e_hybrid_2=p_out_e_hybrid_2) + + @pytest.fixture + def input_data(self, grid): + p_out_e_hybrid_2 = random_field(grid, EdgeDim, KDim) + p_mass_flx_e = random_field(grid, EdgeDim, KDim) + z_dreg_area = random_field(grid, EdgeDim, KDim) + return dict( + p_mass_flx_e=p_mass_flx_e, + z_dreg_area=z_dreg_area, + p_out_e_hybrid_2=p_out_e_hybrid_2, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01a.py new file mode 100644 index 0000000000..0dbf30ca22 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01a.py @@ -0,0 +1,61 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01a import ( + hflx_limiter_mo_stencil_01a, +) +from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestHflxLimiterMoStencil01a(StencilTest): + PROGRAM = hflx_limiter_mo_stencil_01a + OUTPUTS = ( + "z_mflx_low", + "z_anti", + ) + + @staticmethod + def reference( + grid, p_mflx_tracer_h: np.ndarray, p_mass_flx_e: np.ndarray, p_cc: np.ndarray, **kwargs + ): + e2c = grid.connectivities[E2CDim] + p_cc_e2c = p_cc[e2c] + + z_mflx_low = 0.5 * ( + p_mass_flx_e * (p_cc_e2c[:, 0] + p_cc_e2c[:, 1]) + - np.absolute(p_mass_flx_e) * (p_cc_e2c[:, 1] - p_cc_e2c[:, 0]) + ) + + z_anti = p_mflx_tracer_h - z_mflx_low + + return dict(z_mflx_low=z_mflx_low, z_anti=z_anti) + + @pytest.fixture + def input_data(self, grid): + p_mflx_tracer_h = random_field(grid, EdgeDim, KDim) + p_mass_flx_e = random_field(grid, EdgeDim, KDim) + p_cc = random_field(grid, CellDim, KDim) + z_mflx_low = zero_field(grid, EdgeDim, KDim) + z_anti = zero_field(grid, EdgeDim, KDim) + + return dict( + p_mflx_tracer_h=p_mflx_tracer_h, + p_mass_flx_e=p_mass_flx_e, + p_cc=p_cc, + z_mflx_low=z_mflx_low, + z_anti=z_anti, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01b.py new file mode 100644 index 0000000000..70181f6105 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01b.py @@ -0,0 +1,120 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01b import ( + hflx_limiter_mo_stencil_01b, +) +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + reshape, + zero_field, +) + + +class TestHflxLimiterMoStencil01b(StencilTest): + PROGRAM = hflx_limiter_mo_stencil_01b + OUTPUTS = ( + "z_mflx_anti_in", + "z_mflx_anti_out", + "z_tracer_new_low", + "z_tracer_max", + "z_tracer_min", + ) + + @staticmethod + def reference( + grid, + geofac_div: np.ndarray, + p_rhodz_now: np.ndarray, + p_rhodz_new: np.ndarray, + z_mflx_low: np.ndarray, + z_anti: np.ndarray, + p_cc: np.ndarray, + p_dtime: float, + **kwargs, + ): + c2e = grid.connectivities[C2EDim] + z_anti_c2e = z_anti[c2e] + + geofac_div = reshape(geofac_div, c2e.shape) + geofac_div = np.expand_dims(geofac_div, axis=-1) + + zero_array = np.zeros(p_rhodz_now.shape) + + z_mflx_anti_1 = p_dtime * geofac_div[:, 0] / p_rhodz_new * z_anti_c2e[:, 0] + z_mflx_anti_2 = p_dtime * geofac_div[:, 1] / p_rhodz_new * z_anti_c2e[:, 1] + z_mflx_anti_3 = p_dtime * geofac_div[:, 2] / p_rhodz_new * z_anti_c2e[:, 2] + + z_mflx_anti_in = -1.0 * ( + np.minimum(zero_array, z_mflx_anti_1) + + np.minimum(zero_array, z_mflx_anti_2) + + np.minimum(zero_array, z_mflx_anti_3) + ) + + z_mflx_anti_out = ( + np.maximum(zero_array, z_mflx_anti_1) + + np.maximum(zero_array, z_mflx_anti_2) + + np.maximum(zero_array, z_mflx_anti_3) + ) + + z_fluxdiv_c = np.sum(z_mflx_low[c2e] * geofac_div, axis=1) + + z_tracer_new_low = (p_cc * p_rhodz_now - p_dtime * z_fluxdiv_c) / p_rhodz_new + z_tracer_max = np.maximum(p_cc, z_tracer_new_low) + z_tracer_min = np.minimum(p_cc, z_tracer_new_low) + + return dict( + z_mflx_anti_in=z_mflx_anti_in, + z_mflx_anti_out=z_mflx_anti_out, + z_tracer_new_low=z_tracer_new_low, + z_tracer_max=z_tracer_max, + z_tracer_min=z_tracer_min, + ) + + @pytest.fixture + def input_data(self, grid): + geofac_div = random_field(grid, CellDim, C2EDim) + geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) + p_rhodz_now = random_field(grid, CellDim, KDim) + p_rhodz_new = random_field(grid, CellDim, KDim) + z_mflx_low = random_field(grid, EdgeDim, KDim) + z_anti = random_field(grid, EdgeDim, KDim) + p_cc = random_field(grid, CellDim, KDim) + p_dtime = 5.0 + + z_mflx_anti_in = zero_field(grid, CellDim, KDim) + z_mflx_anti_out = zero_field(grid, CellDim, KDim) + z_tracer_new_low = zero_field(grid, CellDim, KDim) + z_tracer_max = zero_field(grid, CellDim, KDim) + z_tracer_min = zero_field(grid, CellDim, KDim) + + return dict( + geofac_div=geofac_div_new, + p_rhodz_now=p_rhodz_now, + p_rhodz_new=p_rhodz_new, + z_mflx_low=z_mflx_low, + z_anti=z_anti, + p_cc=p_cc, + p_dtime=p_dtime, + z_mflx_anti_in=z_mflx_anti_in, + z_mflx_anti_out=z_mflx_anti_out, + z_tracer_new_low=z_tracer_new_low, + z_tracer_max=z_tracer_max, + z_tracer_min=z_tracer_min, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_02.py new file mode 100644 index 0000000000..4925d83987 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_02.py @@ -0,0 +1,85 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import numpy as np +import pytest +from numpy import int32 + +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_02 import hflx_limiter_mo_stencil_02 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + constant_field, + random_field, + zero_field, +) + + +class TestHflxLimiterMoStencil02(StencilTest): + PROGRAM = hflx_limiter_mo_stencil_02 + OUTPUTS = ("z_tracer_new_low", "z_tracer_max", "z_tracer_min") + + @staticmethod + def reference( + grid, + refin_ctrl: np.ndarray, + p_cc: np.ndarray, + z_tracer_new_low: np.ndarray, + z_tracer_max: np.ndarray, + z_tracer_min: np.ndarray, + lo_bound: float, + hi_bound: float, + **kwargs, + ): + refin_ctrl = np.expand_dims(refin_ctrl, axis=1) + condition = np.logical_or( + np.equal(refin_ctrl, lo_bound * np.ones(refin_ctrl.shape, dtype=int32)), + np.equal(refin_ctrl, hi_bound * np.ones(refin_ctrl.shape, dtype=int32)), + ) + z_tracer_new_out = np.where( + condition, + np.minimum(1.1 * p_cc, np.maximum(0.9 * p_cc, z_tracer_new_low)), + z_tracer_new_low, + ) + z_tracer_max_out = np.where(condition, np.maximum(p_cc, z_tracer_new_out), z_tracer_max) + z_tracer_min_out = np.where(condition, np.minimum(p_cc, z_tracer_new_out), z_tracer_min) + return dict( + z_tracer_new_low=z_tracer_new_out, + z_tracer_max=z_tracer_max_out, + z_tracer_min=z_tracer_min_out, + ) + + @pytest.fixture() + def input_data(self, grid): + hi_bound, lo_bound = 3, 1 + refin_ctrl = constant_field(grid, 2, CellDim, dtype=int32) + p_cc = random_field(grid, CellDim, KDim) + z_tracer_new_low_in = random_field(grid, CellDim, KDim) + z_tracer_max_in = random_field(grid, CellDim, KDim) + z_tracer_min_in = random_field(grid, CellDim, KDim) + + z_tracer_new_low_out = zero_field(grid, CellDim, KDim) + z_tracer_max_out = zero_field(grid, CellDim, KDim) + z_tracer_min_out = zero_field(grid, CellDim, KDim) + + return dict( + refin_ctrl=refin_ctrl, + p_cc=p_cc, + z_tracer_new_low=z_tracer_new_low_in, + z_tracer_max=z_tracer_max_in, + z_tracer_min=z_tracer_min_in, + lo_bound=lo_bound, + hi_bound=hi_bound, + z_tracer_new_low_out=z_tracer_new_low_out, + z_tracer_max_out=z_tracer_max_out, + z_tracer_min_out=z_tracer_min_out, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py new file mode 100644 index 0000000000..f8819ee83e --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py @@ -0,0 +1,101 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_03 import ( + hflx_limiter_mo_stencil_03, + hflx_limiter_mo_stencil_03_min_max, +) +from icon4py.model.common.dimension import C2E2CDim, CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestHflxLimiterMoStencil03MinMax(StencilTest): + PROGRAM = hflx_limiter_mo_stencil_03_min_max + OUTPUTS = ("z_max", "z_min") + + @staticmethod + def reference(grid, z_tracer_max, z_tracer_min, beta_fct, r_beta_fct, **kwargs): + c2e2c = grid.connectivities[C2E2CDim] + z_max = beta_fct * np.maximum(np.max(z_tracer_max[c2e2c], axis=1), z_tracer_max) + z_min = r_beta_fct * np.minimum(np.min(z_tracer_min[c2e2c], axis=1), z_tracer_min) + return dict(z_max=z_max, z_min=z_min) + + @pytest.fixture + def input_data(self, grid): + z_tracer_max = random_field(grid, CellDim, KDim) + z_tracer_min = random_field(grid, CellDim, KDim) + beta_fct = 0.9 + r_beta_fct = 0.3 + z_max = zero_field(grid, CellDim, KDim) + z_min = zero_field(grid, CellDim, KDim) + return dict( + z_tracer_max=z_tracer_max, + z_tracer_min=z_tracer_min, + beta_fct=beta_fct, + r_beta_fct=r_beta_fct, + z_max=z_max, + z_min=z_min, + ) + + +class TestHflxLimiterMoStencil03(StencilTest): + PROGRAM = hflx_limiter_mo_stencil_03 + OUTPUTS = ("r_p", "r_m") + + @staticmethod + def reference( + grid, + z_tracer_max, + z_tracer_min, + beta_fct, + r_beta_fct, + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + dbl_eps, + **kwargs, + ): + res = TestHflxLimiterMoStencil03MinMax.reference( + grid, z_tracer_max, z_tracer_min, beta_fct, r_beta_fct + ) + z_max, z_min = res["z_max"], res["z_min"] + r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) + r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out + dbl_eps) + return dict(r_p=r_p, r_m=r_m) + + @pytest.fixture + def input_data(self, grid): + z_tracer_max = random_field(grid, CellDim, KDim) + z_tracer_min = random_field(grid, CellDim, KDim) + beta_fct = 0.4 + r_beta_fct = 0.6 + z_mflx_anti_in = random_field(grid, CellDim, KDim) + z_mflx_anti_out = random_field(grid, CellDim, KDim) + z_tracer_new_low = random_field(grid, CellDim, KDim) + dbl_eps = 1e-5 + r_p = zero_field(grid, CellDim, KDim) + r_m = zero_field(grid, CellDim, KDim) + return dict( + z_tracer_max=z_tracer_max, + z_tracer_min=z_tracer_min, + beta_fct=beta_fct, + r_beta_fct=r_beta_fct, + z_mflx_anti_in=z_mflx_anti_in, + z_mflx_anti_out=z_mflx_anti_out, + z_tracer_new_low=z_tracer_new_low, + dbl_eps=dbl_eps, + r_p=r_p, + r_m=r_m, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_04.py new file mode 100644 index 0000000000..d1b5e3244d --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_04.py @@ -0,0 +1,50 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_04 import hflx_limiter_mo_stencil_04 +from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestHflxLimiterMoStencil04(StencilTest): + PROGRAM = hflx_limiter_mo_stencil_04 + OUTPUTS = ("p_mflx_tracer_h",) + + @staticmethod + def reference( + grid, z_anti: np.ndarray, r_m: np.ndarray, r_p: np.ndarray, z_mflx_low: np.ndarray, **kwargs + ): + r_frac = np.where( + z_anti >= 0, + np.minimum( + r_m[grid.connectivities[E2CDim][:, 0]], r_p[grid.connectivities[E2CDim][:, 1]] + ), + np.minimum( + r_m[grid.connectivities[E2CDim][:, 1]], r_p[grid.connectivities[E2CDim][:, 0]] + ), + ) + return dict(p_mflx_tracer_h=z_mflx_low + np.minimum(1.0, r_frac) * z_anti) + + @pytest.fixture + def input_data(self, grid): + z_anti = random_field(grid, EdgeDim, KDim, low=-2.0, high=2.0) + r_m = random_field(grid, CellDim, KDim) + r_p = random_field(grid, CellDim, KDim) + z_mflx_low = random_field(grid, EdgeDim, KDim) + p_mflx_tracer_h = zero_field(grid, EdgeDim, KDim) + return dict( + z_anti=z_anti, r_m=r_m, r_p=r_p, z_mflx_low=z_mflx_low, p_mflx_tracer_h=p_mflx_tracer_h + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_01.py new file mode 100644 index 0000000000..9ebfe4692a --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_01.py @@ -0,0 +1,77 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_01 import hflx_limiter_pd_stencil_01 +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + reshape, + zero_field, +) + + +class TestHflxLimiterPdStencil01(StencilTest): + PROGRAM = hflx_limiter_pd_stencil_01 + OUTPUTS = ("r_m",) + + @staticmethod + def reference( + grid, + geofac_div: np.ndarray, + p_cc: np.ndarray, + p_rhodz_now: np.ndarray, + p_mflx_tracer_h: np.ndarray, + p_dtime, + dbl_eps, + **kwargs, + ): + geofac_div = reshape(geofac_div, grid.connectivities[C2EDim].shape) + geofac_div = np.expand_dims(geofac_div, axis=-1) + p_m_0 = np.maximum( + 0.0, p_mflx_tracer_h[grid.connectivities[C2EDim][:, 0]] * geofac_div[:, 0] * p_dtime + ) + p_m_1 = np.maximum( + 0.0, p_mflx_tracer_h[grid.connectivities[C2EDim][:, 1]] * geofac_div[:, 1] * p_dtime + ) + p_m_2 = np.maximum( + 0.0, p_mflx_tracer_h[grid.connectivities[C2EDim][:, 2]] * geofac_div[:, 2] * p_dtime + ) + + p_m = p_m_0 + p_m_1 + p_m_2 + r_m = np.minimum(1.0, p_cc * p_rhodz_now / (p_m + dbl_eps)) + + return dict(r_m=r_m) + + @pytest.fixture + def input_data(self, grid): + geofac_div = random_field(grid, CellDim, C2EDim) + geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) + p_cc = random_field(grid, CellDim, KDim) + p_rhodz_now = random_field(grid, CellDim, KDim) + p_mflx_tracer_h = random_field(grid, EdgeDim, KDim) + r_m = zero_field(grid, CellDim, KDim) + p_dtime = np.float64(5) + dbl_eps = np.float64(1e-9) + return dict( + geofac_div=geofac_div_new, + p_cc=p_cc, + p_rhodz_now=p_rhodz_now, + p_mflx_tracer_h=p_mflx_tracer_h, + p_dtime=p_dtime, + dbl_eps=dbl_eps, + r_m=r_m, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py new file mode 100644 index 0000000000..212f95664a --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py @@ -0,0 +1,62 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_02 import hflx_limiter_pd_stencil_02 +from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, constant_field, random_field + + +class TestHflxLimiterPdStencil02(StencilTest): + PROGRAM = hflx_limiter_pd_stencil_02 + OUTPUTS = ("p_mflx_tracer_h",) + + @staticmethod + def reference(grid, refin_ctrl, r_m, p_mflx_tracer_h, bound, **kwargs): + e2c = grid.connectivities[E2CDim] + r_m_e2c = r_m[e2c] + refin_ctrl_expanded = np.expand_dims(refin_ctrl, axis=-1) + p_mflx_tracer_h_out = np.where( + refin_ctrl_expanded != bound, + np.where( + p_mflx_tracer_h >= 0, + p_mflx_tracer_h * r_m_e2c[:, 0], + p_mflx_tracer_h * r_m_e2c[:, 1], + ), + p_mflx_tracer_h, + ) + return dict(p_mflx_tracer_h=p_mflx_tracer_h_out) + + @pytest.fixture(params=[("no_match", 4), ("everywhere_match", 7), ("partly_match", 4)]) + def input_data(self, request, grid): + bound = np.int32(7) + scenario, ctrl_value = request.param + + if scenario == "no_match": + refin_ctrl = constant_field(grid, ctrl_value, EdgeDim, dtype=np.int32) + elif scenario == "everywhere_match": + refin_ctrl = constant_field(grid, bound, EdgeDim, dtype=np.int32) + elif scenario == "partly_match": + refin_ctrl = constant_field(grid, 5, EdgeDim, dtype=np.int32) + refin_ctrl[2:6] = bound + + r_m = random_field(grid, CellDim, KDim) + p_mflx_tracer_h_in = random_field(grid, EdgeDim, KDim) + + return dict( + refin_ctrl=refin_ctrl, + r_m=r_m, + p_mflx_tracer_h=p_mflx_tracer_h_in, + bound=bound, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py new file mode 100644 index 0000000000..51fb24ffb9 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py @@ -0,0 +1,74 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.hor_adv_stencil_01 import hor_adv_stencil_01 +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + reshape, + zero_field, +) + + +class TestHorAdvStencil01(StencilTest): + PROGRAM = hor_adv_stencil_01 + OUTPUTS = ("tracer_new",) + + @staticmethod + def reference( + grid, + p_mflx_tracer_h: np.array, + deepatmo_divh: np.array, + tracer_now: np.array, + rhodz_now: np.array, + rhodz_new: np.array, + geofac_div: np.array, + p_dtime, + **kwargs, + ) -> np.array: + geofac_div = reshape(geofac_div, grid.connectivities[C2EDim].shape) + geofac_div = np.expand_dims(geofac_div, axis=-1) + tracer_new = ( + tracer_now * rhodz_now + - p_dtime + * deepatmo_divh + * np.sum(p_mflx_tracer_h[grid.connectivities[C2EDim]] * geofac_div, axis=1) + ) / rhodz_new + return dict(tracer_new=tracer_new) + + @pytest.fixture + def input_data(self, grid): + p_mflx_tracer_h = random_field(grid, EdgeDim, KDim) + deepatmo_divh = random_field(grid, KDim) + tracer_now = random_field(grid, CellDim, KDim) + rhodz_now = random_field(grid, CellDim, KDim) + rhodz_new = random_field(grid, CellDim, KDim) + geofac_div = random_field(grid, CellDim, C2EDim) + geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) + p_dtime = np.float64(5.0) + tracer_new = zero_field(grid, CellDim, KDim) + return dict( + p_mflx_tracer_h=p_mflx_tracer_h, + deepatmo_divh=deepatmo_divh, + tracer_now=tracer_now, + rhodz_now=rhodz_now, + rhodz_new=rhodz_new, + geofac_div=geofac_div_new, + p_dtime=p_dtime, + tracer_new=tracer_new, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_init_zero_c.py b/model/atmosphere/advection/tests/stencil_tests/test_init_zero_c.py new file mode 100644 index 0000000000..08da338d49 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_init_zero_c.py @@ -0,0 +1,32 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import pytest + +from icon4py.model.atmosphere.advection.init_zero_c import init_zero_c +from icon4py.model.common.dimension import CellDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestInitZeroCellK(StencilTest): + PROGRAM = init_zero_c + OUTPUTS = ("field",) + + @staticmethod + def reference(grid, **kwargs): + return dict(field=zero_field(grid, CellDim).asnumpy()) + + @pytest.fixture + def input_data(self, grid): + field = random_field(grid, CellDim) + return dict(field=field) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_init_zero_c_k.py b/model/atmosphere/advection/tests/stencil_tests/test_init_zero_c_k.py new file mode 100644 index 0000000000..0f0e3f5d74 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_init_zero_c_k.py @@ -0,0 +1,32 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import pytest + +from icon4py.model.atmosphere.advection.init_zero_c_k import init_zero_c_k +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestInitZeroCK(StencilTest): + PROGRAM = init_zero_c_k + OUTPUTS = ("field",) + + @staticmethod + def reference(grid, **kwargs): + return dict(field=zero_field(grid, CellDim, KDim).asnumpy()) + + @pytest.fixture + def input_data(self, grid): + field = random_field(grid, CellDim, KDim) + return dict(field=field) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py new file mode 100644 index 0000000000..05ca5d197a --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py @@ -0,0 +1,150 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.mo_advection_traj_btraj_compute_o1_dsl import ( + mo_advection_traj_btraj_compute_o1_dsl, +) +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + constant_field, + numpy_to_1D_sparse_field, + random_field, +) + + +class TestMoAdvectionTrajBtrajComputeO1Dsl(StencilTest): + PROGRAM = mo_advection_traj_btraj_compute_o1_dsl + OUTPUTS = ("p_cell_idx", "p_cell_blk", "p_distv_bary_1", "p_distv_bary_2") + + @staticmethod + def reference( + grid, + p_vn: np.array, + p_vt: np.array, + cell_idx: np.array, + cell_blk: np.array, + pos_on_tplane_e_1: np.array, + pos_on_tplane_e_2: np.array, + primal_normal_cell_1: np.array, + dual_normal_cell_1: np.array, + primal_normal_cell_2: np.array, + dual_normal_cell_2: np.array, + p_dthalf: float, + **kwargs, + ) -> dict: + e2c = grid.connectivities[E2CDim] + cell_idx = cell_idx.reshape(e2c.shape) + cell_blk = cell_blk.reshape(e2c.shape) + pos_on_tplane_e_1 = pos_on_tplane_e_1.reshape(e2c.shape) + pos_on_tplane_e_2 = pos_on_tplane_e_2.reshape(e2c.shape) + primal_normal_cell_1 = primal_normal_cell_1.reshape(e2c.shape) + primal_normal_cell_2 = primal_normal_cell_2.reshape(e2c.shape) + dual_normal_cell_1 = dual_normal_cell_1.reshape(e2c.shape) + dual_normal_cell_2 = dual_normal_cell_2.reshape(e2c.shape) + + lvn_pos = np.where(p_vn > 0.0, True, False) + cell_idx = np.expand_dims(cell_idx, axis=-1) + cell_blk = np.expand_dims(cell_blk, axis=-1) + pos_on_tplane_e_1 = np.expand_dims(pos_on_tplane_e_1, axis=-1) + pos_on_tplane_e_2 = np.expand_dims(pos_on_tplane_e_2, axis=-1) + primal_normal_cell_1 = np.expand_dims(primal_normal_cell_1, axis=-1) + dual_normal_cell_1 = np.expand_dims(dual_normal_cell_1, axis=-1) + primal_normal_cell_2 = np.expand_dims(primal_normal_cell_2, axis=-1) + dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) + + p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) + p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) + + z_ntdistv_bary_1 = -( + p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + ) + z_ntdistv_bary_2 = -( + p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + ) + + p_distv_bary_1 = np.where( + lvn_pos, + z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] + + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], + ) + + p_distv_bary_2 = np.where( + lvn_pos, + z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] + + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], + ) + + return dict( + p_cell_idx=p_cell_idx, + p_cell_blk=p_cell_blk, + p_distv_bary_1=p_distv_bary_1, + p_distv_bary_2=p_distv_bary_2, + ) + + @pytest.fixture + def input_data(self, grid): + p_vn = random_field(grid, EdgeDim, KDim) + p_vt = random_field(grid, EdgeDim, KDim) + cell_idx = np.asarray(grid.connectivities[E2CDim], dtype=int32) + cell_idx_new = numpy_to_1D_sparse_field(cell_idx, ECDim) + cell_blk = constant_field(grid, 1, EdgeDim, E2CDim, dtype=int32) + cell_blk_new = as_1D_sparse_field(cell_blk, ECDim) + pos_on_tplane_e_1 = random_field(grid, EdgeDim, E2CDim) + pos_on_tplane_e_1_new = as_1D_sparse_field(pos_on_tplane_e_1, ECDim) + pos_on_tplane_e_2 = random_field(grid, EdgeDim, E2CDim) + pos_on_tplane_e_2_new = as_1D_sparse_field(pos_on_tplane_e_2, ECDim) + primal_normal_cell_1 = random_field(grid, EdgeDim, E2CDim) + primal_normal_cell_1_new = as_1D_sparse_field(primal_normal_cell_1, ECDim) + dual_normal_cell_1 = random_field(grid, EdgeDim, E2CDim) + dual_normal_cell_1_new = as_1D_sparse_field(dual_normal_cell_1, ECDim) + primal_normal_cell_2 = random_field(grid, EdgeDim, E2CDim) + primal_normal_cell_2_new = as_1D_sparse_field(primal_normal_cell_2, ECDim) + dual_normal_cell_2 = random_field(grid, EdgeDim, E2CDim) + dual_normal_cell_2_new = as_1D_sparse_field(dual_normal_cell_2, ECDim) + p_cell_idx = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) + p_cell_blk = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) + p_distv_bary_1 = random_field(grid, EdgeDim, KDim) + p_distv_bary_2 = random_field(grid, EdgeDim, KDim) + p_dthalf = 2.0 + + return dict( + p_vn=p_vn, + p_vt=p_vt, + cell_idx=cell_idx_new, + cell_blk=cell_blk_new, + pos_on_tplane_e_1=pos_on_tplane_e_1_new, + pos_on_tplane_e_2=pos_on_tplane_e_2_new, + primal_normal_cell_1=primal_normal_cell_1_new, + dual_normal_cell_1=dual_normal_cell_1_new, + primal_normal_cell_2=primal_normal_cell_2_new, + dual_normal_cell_2=dual_normal_cell_2_new, + p_cell_idx=p_cell_idx, + p_cell_blk=p_cell_blk, + p_distv_bary_1=p_distv_bary_1, + p_distv_bary_2=p_distv_bary_2, + p_dthalf=p_dthalf, + horizontal_start=int32(0), + horizontal_end=int32(grid.num_edges), + vertical_start=int32(0), + vertical_end=int32(grid.num_levels), + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py new file mode 100644 index 0000000000..f883693ab1 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py @@ -0,0 +1,636 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.prep_gauss_quadrature_c_list_stencil import ( + prep_gauss_quadrature_c_list_stencil, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + constant_field, + random_field, + zero_field, +) + + +@pytest.mark.slow_tests +class TestPrepGaussQuadratureCListStencil(StencilTest): + PROGRAM = prep_gauss_quadrature_c_list_stencil + OUTPUTS = ( + "p_quad_vector_sum_1", + "p_quad_vector_sum_2", + "p_quad_vector_sum_3", + "p_quad_vector_sum_4", + "p_quad_vector_sum_5", + "p_quad_vector_sum_6", + "p_quad_vector_sum_7", + "p_quad_vector_sum_8", + "p_quad_vector_sum_9", + "p_quad_vector_sum_10", + "p_dreg_area", + ) + + @staticmethod + def _compute_wgt_t_detjac( + wgt_zeta_1, + wgt_eta_1, + wgt_zeta_2, + wgt_eta_2, + eta_1, + eta_2, + eta_3, + eta_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + famask_int, + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + dbl_eps, + ): + z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 + z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 + z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 + z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 + + z_eta_1_1, z_eta_2_1, z_eta_3_1, z_eta_4_1 = ( + 1.0 - eta_1, + 1.0 - eta_2, + 1.0 - eta_3, + 1.0 - eta_4, + ) + z_eta_1_2, z_eta_2_2, z_eta_3_2, z_eta_4_2 = ( + 1.0 + eta_1, + 1.0 + eta_2, + 1.0 + eta_3, + 1.0 + eta_4, + ) + z_eta_1_3, z_eta_2_3, z_eta_3_3, z_eta_4_3 = ( + 1.0 - zeta_1, + 1.0 - zeta_2, + 1.0 - zeta_3, + 1.0 - zeta_4, + ) + z_eta_1_4, z_eta_2_4, z_eta_3_4, z_eta_4_4 = ( + 1.0 + zeta_1, + 1.0 + zeta_2, + 1.0 + zeta_3, + 1.0 + zeta_4, + ) + + famask_bool = np.where(famask_int == np.int32(1), True, False) + + p_coords_dreg_v_1_x = np.where(famask_bool, p_coords_dreg_v_1_x, 0.0) + p_coords_dreg_v_2_x = np.where(famask_bool, p_coords_dreg_v_2_x, 0.0) + p_coords_dreg_v_3_x = np.where(famask_bool, p_coords_dreg_v_3_x, 0.0) + p_coords_dreg_v_4_x = np.where(famask_bool, p_coords_dreg_v_4_x, 0.0) + p_coords_dreg_v_1_y = np.where(famask_bool, p_coords_dreg_v_1_y, 0.0) + p_coords_dreg_v_2_y = np.where(famask_bool, p_coords_dreg_v_2_y, 0.0) + p_coords_dreg_v_3_y = np.where(famask_bool, p_coords_dreg_v_3_y, 0.0) + p_coords_dreg_v_4_y = np.where(famask_bool, p_coords_dreg_v_4_y, 0.0) + + wgt_t_detjac_1 = np.where( + famask_bool, + dbl_eps + + z_wgt_1 + * ( + ( + z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + + wgt_t_detjac_2 = np.where( + famask_bool, + dbl_eps + + z_wgt_2 + * ( + ( + z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + + wgt_t_detjac_3 = np.where( + famask_bool, + dbl_eps + + z_wgt_3 + * ( + ( + z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_4 = np.where( + famask_bool, + dbl_eps + + z_wgt_4 + * ( + ( + z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + + return wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 + + @staticmethod + def _compute_z_gauss_points( + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + ): + z_gauss_pts_1_x = ( + shape_func_1_1 * p_coords_dreg_v_1_x + + shape_func_2_1 * p_coords_dreg_v_2_x + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) + return ( + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ) + + @staticmethod + def _compute_vector_sums( + wgt_t_detjac_1, + wgt_t_detjac_2, + wgt_t_detjac_3, + wgt_t_detjac_4, + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ): + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + return ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + ) + + @classmethod + def reference( + cls, + grid, + famask_int: np.array, + p_coords_dreg_v_1_x: np.array, + p_coords_dreg_v_2_x: np.array, + p_coords_dreg_v_3_x: np.array, + p_coords_dreg_v_4_x: np.array, + p_coords_dreg_v_1_y: np.array, + p_coords_dreg_v_2_y: np.array, + p_coords_dreg_v_3_y: np.array, + p_coords_dreg_v_4_y: np.array, + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, + p_dreg_area_in: np.array, + **kwargs, + ): + wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 = cls._compute_wgt_t_detjac( + wgt_zeta_1, + wgt_eta_1, + wgt_zeta_2, + wgt_eta_2, + eta_1, + eta_2, + eta_3, + eta_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + famask_int, + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + dbl_eps, + ) + + ( + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ) = cls._compute_z_gauss_points( + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + ) + + ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + ) = cls._compute_vector_sums( + wgt_t_detjac_1, + wgt_t_detjac_2, + wgt_t_detjac_3, + wgt_t_detjac_4, + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ) + + p_dreg_area = p_dreg_area_in + p_quad_vector_sum_1 + return dict( + p_quad_vector_sum_1=p_quad_vector_sum_1, + p_quad_vector_sum_2=p_quad_vector_sum_2, + p_quad_vector_sum_3=p_quad_vector_sum_3, + p_quad_vector_sum_4=p_quad_vector_sum_4, + p_quad_vector_sum_5=p_quad_vector_sum_5, + p_quad_vector_sum_6=p_quad_vector_sum_6, + p_quad_vector_sum_7=p_quad_vector_sum_7, + p_quad_vector_sum_8=p_quad_vector_sum_8, + p_quad_vector_sum_9=p_quad_vector_sum_9, + p_quad_vector_sum_10=p_quad_vector_sum_10, + p_dreg_area=p_dreg_area, + ) + + @pytest.fixture + def input_data(self, grid): + famask_int = constant_field(grid, 1, EdgeDim, KDim, dtype=int32) + p_coords_dreg_v_1_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_2_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_3_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_4_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_1_y = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_2_y = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_3_y = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_4_y = random_field(grid, EdgeDim, KDim) + shape_func_1_1 = 0.001 + shape_func_2_1 = 0.001 + shape_func_3_1 = 0.001 + shape_func_4_1 = 0.001 + shape_func_1_2 = 0.001 + shape_func_2_2 = 0.001 + shape_func_3_2 = 0.001 + shape_func_4_2 = 0.001 + shape_func_1_3 = 0.001 + shape_func_2_3 = 0.001 + shape_func_3_3 = 0.001 + shape_func_4_3 = 0.001 + shape_func_1_4 = 0.001 + shape_func_2_4 = 0.001 + shape_func_3_4 = 0.001 + shape_func_4_4 = 0.001 + zeta_1 = 0.002 + zeta_2 = 0.002 + zeta_3 = 0.002 + zeta_4 = 0.002 + eta_1 = 0.5 + eta_2 = 0.5 + eta_3 = 0.5 + eta_4 = 0.5 + wgt_zeta_1 = 0.003 + wgt_zeta_2 = 0.003 + wgt_eta_1 = 0.002 + wgt_eta_2 = 0.007 + dbl_eps = np.float64(0.1) + eps = 0.1 + p_dreg_area_in = random_field(grid, EdgeDim, KDim) + p_quad_vector_sum_1 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_2 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_3 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_4 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_5 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_6 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_7 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_8 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_9 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_10 = zero_field(grid, EdgeDim, KDim) + p_dreg_area = zero_field(grid, EdgeDim, KDim) + return dict( + famask_int=famask_int, + p_coords_dreg_v_1_x=p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x=p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x=p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x=p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y=p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y=p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y=p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y=p_coords_dreg_v_4_y, + shape_func_1_1=shape_func_1_1, + shape_func_2_1=shape_func_2_1, + shape_func_3_1=shape_func_3_1, + shape_func_4_1=shape_func_4_1, + shape_func_1_2=shape_func_1_2, + shape_func_2_2=shape_func_2_2, + shape_func_3_2=shape_func_3_2, + shape_func_4_2=shape_func_4_2, + shape_func_1_3=shape_func_1_3, + shape_func_2_3=shape_func_2_3, + shape_func_3_3=shape_func_3_3, + shape_func_4_3=shape_func_4_3, + shape_func_1_4=shape_func_1_4, + shape_func_2_4=shape_func_2_4, + shape_func_3_4=shape_func_3_4, + shape_func_4_4=shape_func_4_4, + zeta_1=zeta_1, + zeta_2=zeta_2, + zeta_3=zeta_3, + zeta_4=zeta_4, + eta_1=eta_1, + eta_2=eta_2, + eta_3=eta_3, + eta_4=eta_4, + wgt_zeta_1=wgt_zeta_1, + wgt_zeta_2=wgt_zeta_2, + wgt_eta_1=wgt_eta_1, + wgt_eta_2=wgt_eta_2, + dbl_eps=dbl_eps, + eps=eps, + p_dreg_area_in=p_dreg_area_in, + p_quad_vector_sum_1=p_quad_vector_sum_1, + p_quad_vector_sum_2=p_quad_vector_sum_2, + p_quad_vector_sum_3=p_quad_vector_sum_3, + p_quad_vector_sum_4=p_quad_vector_sum_4, + p_quad_vector_sum_5=p_quad_vector_sum_5, + p_quad_vector_sum_6=p_quad_vector_sum_6, + p_quad_vector_sum_7=p_quad_vector_sum_7, + p_quad_vector_sum_8=p_quad_vector_sum_8, + p_quad_vector_sum_9=p_quad_vector_sum_9, + p_quad_vector_sum_10=p_quad_vector_sum_10, + p_dreg_area=p_dreg_area, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py new file mode 100644 index 0000000000..a6b7b6b732 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py @@ -0,0 +1,591 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.prep_gauss_quadrature_c_stencil import ( + prep_gauss_quadrature_c_stencil, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +@pytest.mark.slow_tests +class TestPrepGaussQuadratureCStencil(StencilTest): + PROGRAM = prep_gauss_quadrature_c_stencil + OUTPUTS = ( + "p_quad_vector_sum_1", + "p_quad_vector_sum_2", + "p_quad_vector_sum_3", + "p_quad_vector_sum_4", + "p_quad_vector_sum_5", + "p_quad_vector_sum_6", + "p_quad_vector_sum_7", + "p_quad_vector_sum_8", + "p_quad_vector_sum_9", + "p_quad_vector_sum_10", + "p_dreg_area_out", + ) + + @staticmethod + def _compute_wgt_t_detjac( + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + ): + # Compute z_wgt + z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 + z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 + z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 + z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 + + # Compute z_eta + z_eta_1_1, z_eta_2_1, z_eta_3_1, z_eta_4_1 = ( + 1.0 - eta_1, + 1.0 - eta_2, + 1.0 - eta_3, + 1.0 - eta_4, + ) + z_eta_1_2, z_eta_2_2, z_eta_3_2, z_eta_4_2 = ( + 1.0 + eta_1, + 1.0 + eta_2, + 1.0 + eta_3, + 1.0 + eta_4, + ) + z_eta_1_3, z_eta_2_3, z_eta_3_3, z_eta_4_3 = ( + 1.0 - zeta_1, + 1.0 - zeta_2, + 1.0 - zeta_3, + 1.0 - zeta_4, + ) + z_eta_1_4, z_eta_2_4, z_eta_3_4, z_eta_4_4 = ( + 1.0 + zeta_1, + 1.0 + zeta_2, + 1.0 + zeta_3, + 1.0 + zeta_4, + ) + + wgt_t_detjac_1 = dbl_eps + z_wgt_1 * ( + ( + z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_2 = dbl_eps + z_wgt_2 * ( + ( + z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_3 = dbl_eps + z_wgt_3 * ( + ( + z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_4 = dbl_eps + z_wgt_4 * ( + ( + z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + return wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 + + @staticmethod + def _compute_z_gauss_points( + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + ): + z_gauss_pts_1_x = ( + shape_func_1_1 * p_coords_dreg_v_1_x + + shape_func_2_1 * p_coords_dreg_v_2_x + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) + return ( + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ) + + @staticmethod + def _compute_vector_sums( + wgt_t_detjac_1, + wgt_t_detjac_2, + wgt_t_detjac_3, + wgt_t_detjac_4, + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ): + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + return ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + ) + + @classmethod + def reference( + cls, + grid, + p_coords_dreg_v_1_x: np.array, + p_coords_dreg_v_2_x: np.array, + p_coords_dreg_v_3_x: np.array, + p_coords_dreg_v_4_x: np.array, + p_coords_dreg_v_1_y: np.array, + p_coords_dreg_v_2_y: np.array, + p_coords_dreg_v_3_y: np.array, + p_coords_dreg_v_4_y: np.array, + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, + **kwargs, + ): + wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 = cls._compute_wgt_t_detjac( + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + ) + + ( + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ) = cls._compute_z_gauss_points( + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + ) + + ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + ) = cls._compute_vector_sums( + wgt_t_detjac_1, + wgt_t_detjac_2, + wgt_t_detjac_3, + wgt_t_detjac_4, + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ) + + z_area = p_quad_vector_sum_1 + p_dreg_area_out = np.where( + z_area >= 0.0, + np.maximum(eps, np.absolute(z_area)), + -np.maximum(eps, np.absolute(z_area)), + ) + return dict( + p_quad_vector_sum_1=p_quad_vector_sum_1, + p_quad_vector_sum_2=p_quad_vector_sum_2, + p_quad_vector_sum_3=p_quad_vector_sum_3, + p_quad_vector_sum_4=p_quad_vector_sum_4, + p_quad_vector_sum_5=p_quad_vector_sum_5, + p_quad_vector_sum_6=p_quad_vector_sum_6, + p_quad_vector_sum_7=p_quad_vector_sum_7, + p_quad_vector_sum_8=p_quad_vector_sum_8, + p_quad_vector_sum_9=p_quad_vector_sum_9, + p_quad_vector_sum_10=p_quad_vector_sum_10, + p_dreg_area_out=p_dreg_area_out, + ) + + @pytest.fixture + def input_data(self, grid): + p_coords_dreg_v_1_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_2_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_3_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_4_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_1_y = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_2_y = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_3_y = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_4_y = random_field(grid, EdgeDim, KDim) + shape_func_1_1 = 0.001 + shape_func_2_1 = 0.001 + shape_func_3_1 = 0.001 + shape_func_4_1 = 0.001 + shape_func_1_2 = 0.001 + shape_func_2_2 = 0.001 + shape_func_3_2 = 0.001 + shape_func_4_2 = 0.001 + shape_func_1_3 = 0.001 + shape_func_2_3 = 0.001 + shape_func_3_3 = 0.001 + shape_func_4_3 = 0.001 + shape_func_1_4 = 0.001 + shape_func_2_4 = 0.001 + shape_func_3_4 = 0.001 + shape_func_4_4 = 0.001 + zeta_1 = 0.002 + zeta_2 = 0.002 + zeta_3 = 0.002 + zeta_4 = 0.002 + eta_1 = 0.5 + eta_2 = 0.5 + eta_3 = 0.5 + eta_4 = 0.5 + wgt_zeta_1 = 0.003 + wgt_zeta_2 = 0.003 + wgt_eta_1 = 0.002 + wgt_eta_2 = 0.007 + dbl_eps = np.float64(0.1) + eps = 0.1 + p_quad_vector_sum_1 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_2 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_3 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_4 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_5 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_6 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_7 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_8 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_9 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_10 = zero_field(grid, EdgeDim, KDim) + p_dreg_area_out = zero_field(grid, EdgeDim, KDim) + return dict( + p_coords_dreg_v_1_x=p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x=p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x=p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x=p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y=p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y=p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y=p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y=p_coords_dreg_v_4_y, + shape_func_1_1=shape_func_1_1, + shape_func_2_1=shape_func_2_1, + shape_func_3_1=shape_func_3_1, + shape_func_4_1=shape_func_4_1, + shape_func_1_2=shape_func_1_2, + shape_func_2_2=shape_func_2_2, + shape_func_3_2=shape_func_3_2, + shape_func_4_2=shape_func_4_2, + shape_func_1_3=shape_func_1_3, + shape_func_2_3=shape_func_2_3, + shape_func_3_3=shape_func_3_3, + shape_func_4_3=shape_func_4_3, + shape_func_1_4=shape_func_1_4, + shape_func_2_4=shape_func_2_4, + shape_func_3_4=shape_func_3_4, + shape_func_4_4=shape_func_4_4, + zeta_1=zeta_1, + zeta_2=zeta_2, + zeta_3=zeta_3, + zeta_4=zeta_4, + eta_1=eta_1, + eta_2=eta_2, + eta_3=eta_3, + eta_4=eta_4, + wgt_zeta_1=wgt_zeta_1, + wgt_zeta_2=wgt_zeta_2, + wgt_eta_1=wgt_eta_1, + wgt_eta_2=wgt_eta_2, + dbl_eps=dbl_eps, + eps=eps, + p_quad_vector_sum_1=p_quad_vector_sum_1, + p_quad_vector_sum_2=p_quad_vector_sum_2, + p_quad_vector_sum_3=p_quad_vector_sum_3, + p_quad_vector_sum_4=p_quad_vector_sum_4, + p_quad_vector_sum_5=p_quad_vector_sum_5, + p_quad_vector_sum_6=p_quad_vector_sum_6, + p_quad_vector_sum_7=p_quad_vector_sum_7, + p_quad_vector_sum_8=p_quad_vector_sum_8, + p_quad_vector_sum_9=p_quad_vector_sum_9, + p_quad_vector_sum_10=p_quad_vector_sum_10, + p_dreg_area_out=p_dreg_area_out, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_rbf_intp_edge_stencil_01.py new file mode 100644 index 0000000000..47882226d9 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_rbf_intp_edge_stencil_01.py @@ -0,0 +1,37 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.rbf_intp_edge_stencil_01 import rbf_intp_edge_stencil_01 +from icon4py.model.common.dimension import E2C2EDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestRbfIntpEdgeStencil01(StencilTest): + PROGRAM = rbf_intp_edge_stencil_01 + OUTPUTS = ("p_vt_out",) + + @staticmethod + def reference(grid, p_vn_in: np.array, ptr_coeff: np.array, **kwargs) -> np.array: + ptr_coeff = np.expand_dims(ptr_coeff, axis=-1) + p_vt_out = np.sum(p_vn_in[grid.connectivities[E2C2EDim]] * ptr_coeff, axis=1) + return dict(p_vt_out=p_vt_out) + + @pytest.fixture + def input_data(self, grid): + p_vn_in = random_field(grid, EdgeDim, KDim) + ptr_coeff = random_field(grid, EdgeDim, E2C2EDim) + p_vt_out = zero_field(grid, EdgeDim, KDim) + return dict(p_vn_in=p_vn_in, ptr_coeff=ptr_coeff, p_vt_out=p_vt_out) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py new file mode 100644 index 0000000000..c3c5bfec8c --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py @@ -0,0 +1,318 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.recon_lsq_cell_c_svd_stencil import ( + recon_lsq_cell_c_svd_stencil, +) +from icon4py.model.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + reshape, + zero_field, +) + + +@pytest.mark.slow_tests +class TestReconLsqCellCSvdStencil(StencilTest): + PROGRAM = recon_lsq_cell_c_svd_stencil + OUTPUTS = ( + "p_coeff_1_dsl", + "p_coeff_2_dsl", + "p_coeff_3_dsl", + "p_coeff_4_dsl", + "p_coeff_5_dsl", + "p_coeff_6_dsl", + "p_coeff_7_dsl", + "p_coeff_8_dsl", + "p_coeff_9_dsl", + "p_coeff_10_dsl", + ) + + @staticmethod + def reference( + grid, + p_cc: np.array, + lsq_pseudoinv_1: np.array, + lsq_pseudoinv_2: np.array, + lsq_pseudoinv_3: np.array, + lsq_pseudoinv_4: np.array, + lsq_pseudoinv_5: np.array, + lsq_pseudoinv_6: np.array, + lsq_pseudoinv_7: np.array, + lsq_pseudoinv_8: np.array, + lsq_pseudoinv_9: np.array, + lsq_moments_1: np.array, + lsq_moments_2: np.array, + lsq_moments_3: np.array, + lsq_moments_4: np.array, + lsq_moments_5: np.array, + lsq_moments_6: np.array, + lsq_moments_7: np.array, + lsq_moments_8: np.array, + lsq_moments_9: np.array, + **kwargs, + ): + c2e2c2e2c = grid.connectivities[C2E2C2E2CDim] + lsq_moments_1 = np.expand_dims(lsq_moments_1, axis=-1) + lsq_moments_2 = np.expand_dims(lsq_moments_2, axis=-1) + lsq_moments_3 = np.expand_dims(lsq_moments_3, axis=-1) + lsq_moments_4 = np.expand_dims(lsq_moments_4, axis=-1) + lsq_moments_5 = np.expand_dims(lsq_moments_5, axis=-1) + lsq_moments_6 = np.expand_dims(lsq_moments_6, axis=-1) + lsq_moments_7 = np.expand_dims(lsq_moments_7, axis=-1) + lsq_moments_8 = np.expand_dims(lsq_moments_8, axis=-1) + lsq_moments_9 = np.expand_dims(lsq_moments_9, axis=-1) + lsq_moments_1 = np.broadcast_to(lsq_moments_1, p_cc.shape) + lsq_moments_2 = np.broadcast_to(lsq_moments_2, p_cc.shape) + lsq_moments_3 = np.broadcast_to(lsq_moments_3, p_cc.shape) + lsq_moments_4 = np.broadcast_to(lsq_moments_4, p_cc.shape) + lsq_moments_5 = np.broadcast_to(lsq_moments_5, p_cc.shape) + lsq_moments_6 = np.broadcast_to(lsq_moments_6, p_cc.shape) + lsq_moments_7 = np.broadcast_to(lsq_moments_7, p_cc.shape) + lsq_moments_8 = np.broadcast_to(lsq_moments_8, p_cc.shape) + lsq_pseudoinv_9 = reshape(lsq_pseudoinv_9, c2e2c2e2c.shape) + lsq_pseudoinv_9 = np.expand_dims(lsq_pseudoinv_9, axis=-1) + lsq_pseudoinv_8 = reshape(lsq_pseudoinv_8, c2e2c2e2c.shape) + lsq_pseudoinv_8 = np.expand_dims(lsq_pseudoinv_8, axis=-1) + lsq_pseudoinv_7 = reshape(lsq_pseudoinv_7, c2e2c2e2c.shape) + lsq_pseudoinv_7 = np.expand_dims(lsq_pseudoinv_7, axis=-1) + lsq_pseudoinv_6 = reshape(lsq_pseudoinv_6, c2e2c2e2c.shape) + lsq_pseudoinv_6 = np.expand_dims(lsq_pseudoinv_6, axis=-1) + lsq_pseudoinv_5 = reshape(lsq_pseudoinv_5, c2e2c2e2c.shape) + lsq_pseudoinv_5 = np.expand_dims(lsq_pseudoinv_5, axis=-1) + lsq_pseudoinv_4 = reshape(lsq_pseudoinv_4, c2e2c2e2c.shape) + lsq_pseudoinv_4 = np.expand_dims(lsq_pseudoinv_4, axis=-1) + lsq_pseudoinv_3 = reshape(lsq_pseudoinv_3, c2e2c2e2c.shape) + lsq_pseudoinv_3 = np.expand_dims(lsq_pseudoinv_3, axis=-1) + lsq_pseudoinv_2 = reshape(lsq_pseudoinv_2, c2e2c2e2c.shape) + lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) + lsq_pseudoinv_1 = reshape(lsq_pseudoinv_1, c2e2c2e2c.shape) + lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) + + p_coeff_10_dsl = ( + lsq_pseudoinv_9[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_9[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_9[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_9[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_9[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_9[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_9[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_9[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_9[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_9_dsl = ( + lsq_pseudoinv_8[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_8[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_8[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_8[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_8[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_8[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_8[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_8[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_8[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_8_dsl = ( + lsq_pseudoinv_7[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_7[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_7[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_7[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_7[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_7[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_7[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_7[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_7[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_7_dsl = ( + lsq_pseudoinv_6[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_6[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_6[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_6[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_6[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_6[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_6[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_6[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_6[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_6_dsl = ( + lsq_pseudoinv_5[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_5[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_5[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_5[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_5[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_5[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_5[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_5[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_5[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_5_dsl = ( + lsq_pseudoinv_4[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_4[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_4[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_4[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_4[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_4[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_4[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_4[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_4[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_4_dsl = ( + lsq_pseudoinv_3[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_3[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_3[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_3[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_3[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_3[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_3[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_3[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_3[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_3_dsl = ( + lsq_pseudoinv_2[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_2[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_2[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_2[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_2[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_2[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_2[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_2[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_2[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_2_dsl = ( + lsq_pseudoinv_1[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_1[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_1[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_1[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_1[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_1[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_1[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_1[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_1[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_1_dsl = p_cc - ( + p_coeff_2_dsl * lsq_moments_1 + + p_coeff_3_dsl * lsq_moments_2 + + p_coeff_4_dsl * lsq_moments_3 + + p_coeff_5_dsl * lsq_moments_4 + + p_coeff_6_dsl * lsq_moments_5 + + p_coeff_7_dsl * lsq_moments_6 + + p_coeff_8_dsl * lsq_moments_7 + + p_coeff_9_dsl * lsq_moments_8 + + p_coeff_10_dsl * lsq_moments_9 + ) + return dict( + p_coeff_1_dsl=p_coeff_1_dsl, + p_coeff_2_dsl=p_coeff_2_dsl, + p_coeff_3_dsl=p_coeff_3_dsl, + p_coeff_4_dsl=p_coeff_4_dsl, + p_coeff_5_dsl=p_coeff_5_dsl, + p_coeff_6_dsl=p_coeff_6_dsl, + p_coeff_7_dsl=p_coeff_7_dsl, + p_coeff_8_dsl=p_coeff_8_dsl, + p_coeff_9_dsl=p_coeff_9_dsl, + p_coeff_10_dsl=p_coeff_10_dsl, + ) + + @pytest.fixture + def input_data(self, grid): + p_cc = random_field(grid, CellDim, KDim) + lsq_pseudoinv_1_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_2_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_3_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_4_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_5_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_6_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_7_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_8_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_9_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_moments_1 = random_field(grid, CellDim) + lsq_moments_2 = random_field(grid, CellDim) + lsq_moments_3 = random_field(grid, CellDim) + lsq_moments_4 = random_field(grid, CellDim) + lsq_moments_5 = random_field(grid, CellDim) + lsq_moments_6 = random_field(grid, CellDim) + lsq_moments_7 = random_field(grid, CellDim) + lsq_moments_8 = random_field(grid, CellDim) + lsq_moments_9 = random_field(grid, CellDim) + p_coeff_1_dsl = zero_field(grid, CellDim, KDim) + p_coeff_2_dsl = zero_field(grid, CellDim, KDim) + p_coeff_3_dsl = zero_field(grid, CellDim, KDim) + p_coeff_4_dsl = zero_field(grid, CellDim, KDim) + p_coeff_5_dsl = zero_field(grid, CellDim, KDim) + p_coeff_6_dsl = zero_field(grid, CellDim, KDim) + p_coeff_7_dsl = zero_field(grid, CellDim, KDim) + p_coeff_8_dsl = zero_field(grid, CellDim, KDim) + p_coeff_9_dsl = zero_field(grid, CellDim, KDim) + p_coeff_10_dsl = zero_field(grid, CellDim, KDim) + return dict( + p_cc=p_cc, + lsq_pseudoinv_1=lsq_pseudoinv_1_field, + lsq_pseudoinv_2=lsq_pseudoinv_2_field, + lsq_pseudoinv_3=lsq_pseudoinv_3_field, + lsq_pseudoinv_4=lsq_pseudoinv_4_field, + lsq_pseudoinv_5=lsq_pseudoinv_5_field, + lsq_pseudoinv_6=lsq_pseudoinv_6_field, + lsq_pseudoinv_7=lsq_pseudoinv_7_field, + lsq_pseudoinv_8=lsq_pseudoinv_8_field, + lsq_pseudoinv_9=lsq_pseudoinv_9_field, + lsq_moments_1=lsq_moments_1, + lsq_moments_2=lsq_moments_2, + lsq_moments_3=lsq_moments_3, + lsq_moments_4=lsq_moments_4, + lsq_moments_5=lsq_moments_5, + lsq_moments_6=lsq_moments_6, + lsq_moments_7=lsq_moments_7, + lsq_moments_8=lsq_moments_8, + lsq_moments_9=lsq_moments_9, + p_coeff_1_dsl=p_coeff_1_dsl, + p_coeff_2_dsl=p_coeff_2_dsl, + p_coeff_3_dsl=p_coeff_3_dsl, + p_coeff_4_dsl=p_coeff_4_dsl, + p_coeff_5_dsl=p_coeff_5_dsl, + p_coeff_6_dsl=p_coeff_6_dsl, + p_coeff_7_dsl=p_coeff_7_dsl, + p_coeff_8_dsl=p_coeff_8_dsl, + p_coeff_9_dsl=p_coeff_9_dsl, + p_coeff_10_dsl=p_coeff_10_dsl, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_l_svd_stencil.py new file mode 100644 index 0000000000..298b6e71f3 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_l_svd_stencil.py @@ -0,0 +1,77 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.recon_lsq_cell_l_svd_stencil import ( + recon_lsq_cell_l_svd_stencil, +) +from icon4py.model.common.dimension import C2E2CDim, CECDim, CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + reshape, + zero_field, +) + + +class TestReconLsqCellLSvdStencil(StencilTest): + PROGRAM = recon_lsq_cell_l_svd_stencil + OUTPUTS = ( + "p_coeff_1_dsl", + "p_coeff_2_dsl", + "p_coeff_3_dsl", + ) + + @staticmethod + def reference( + grid, + p_cc: np.array, + lsq_pseudoinv_1: np.array, + lsq_pseudoinv_2: np.array, + **kwargs, + ): + c2e2c = grid.connectivities[C2E2CDim] + p_cc_e = np.expand_dims(p_cc, axis=1) + n_diff = p_cc[c2e2c] - p_cc_e + lsq_pseudoinv_1 = reshape(lsq_pseudoinv_1, c2e2c.shape) + lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) + lsq_pseudoinv_2 = reshape(lsq_pseudoinv_2, c2e2c.shape) + lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) + p_coeff_1_dsl = p_cc + p_coeff_2_dsl = np.sum(lsq_pseudoinv_1 * n_diff, axis=1) + p_coeff_3_dsl = np.sum(lsq_pseudoinv_2 * n_diff, axis=1) + return dict( + p_coeff_1_dsl=p_coeff_1_dsl, + p_coeff_2_dsl=p_coeff_2_dsl, + p_coeff_3_dsl=p_coeff_3_dsl, + ) + + @pytest.fixture + def input_data(self, grid): + p_cc = random_field(grid, CellDim, KDim) + lsq_pseudoinv_1_field = as_1D_sparse_field(random_field(grid, CellDim, C2E2CDim), CECDim) + lsq_pseudoinv_2_field = as_1D_sparse_field(random_field(grid, CellDim, C2E2CDim), CECDim) + p_coeff_1_dsl = zero_field(grid, CellDim, KDim) + p_coeff_2_dsl = zero_field(grid, CellDim, KDim) + p_coeff_3_dsl = zero_field(grid, CellDim, KDim) + return dict( + p_cc=p_cc, + lsq_pseudoinv_1=lsq_pseudoinv_1_field, + lsq_pseudoinv_2=lsq_pseudoinv_2_field, + p_coeff_1_dsl=p_coeff_1_dsl, + p_coeff_2_dsl=p_coeff_2_dsl, + p_coeff_3_dsl=p_coeff_3_dsl, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_01.py new file mode 100644 index 0000000000..f45130e824 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_01.py @@ -0,0 +1,64 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.step_advection_stencil_01 import step_advection_stencil_01 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field, zero_field + + +class TestStepAdvectionStencil01(StencilTest): + PROGRAM = step_advection_stencil_01 + OUTPUTS = ( + Output( + "rhodz_ast2", + refslice=(slice(None), slice(None, -1)), + gtslice=(slice(None), slice(None, -1)), + ), + ) + + @staticmethod + def reference( + grid, + rhodz_ast: np.array, + p_mflx_contra_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, + p_dtime, + **kwargs, + ): + tmp = p_dtime * ( + p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu + ) + rhodz_ast2 = rhodz_ast + tmp + + return dict(rhodz_ast2=rhodz_ast2) + + @pytest.fixture + def input_data(self, grid): + rhodz_ast = random_field(grid, CellDim, KDim) + p_mflx_contra_v = random_field(grid, CellDim, KDim, extend={KDim: 1}) + deepatmo_divzl = random_field(grid, KDim) + deepatmo_divzu = random_field(grid, KDim) + rhodz_ast2 = zero_field(grid, CellDim, KDim) + p_dtime = 0.1 + return dict( + rhodz_ast=rhodz_ast, + p_mflx_contra_v=p_mflx_contra_v, + deepatmo_divzl=deepatmo_divzl, + deepatmo_divzu=deepatmo_divzu, + p_dtime=p_dtime, + rhodz_ast2=rhodz_ast2, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_02.py new file mode 100644 index 0000000000..c34cd6f037 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_02.py @@ -0,0 +1,61 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.step_advection_stencil_02 import step_advection_stencil_02 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field, zero_field + + +class TestStepAdvectionStencil02(StencilTest): + PROGRAM = step_advection_stencil_02 + OUTPUTS = ( + Output( + "rhodz_ast2", + refslice=(slice(None), slice(None, -1)), + gtslice=(slice(None), slice(None, -1)), + ), + ) + + @staticmethod + def reference( + grid, + p_rhodz_new: np.array, + p_mflx_contra_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, + p_dtime: float, + **kwargs, + ): + tmp = p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu + rhodz_ast2 = np.maximum(0.1 * p_rhodz_new, p_rhodz_new) - p_dtime * tmp + return dict(rhodz_ast2=rhodz_ast2) + + @pytest.fixture + def input_data(self, grid): + p_rhodz_new = random_field(grid, CellDim, KDim) + p_mflx_contra_v = random_field(grid, CellDim, KDim, extend={KDim: 1}) + deepatmo_divzl = random_field(grid, KDim) + deepatmo_divzu = random_field(grid, KDim) + p_dtime = 0.1 + rhodz_ast2 = zero_field(grid, CellDim, KDim) + return dict( + p_rhodz_new=p_rhodz_new, + p_mflx_contra_v=p_mflx_contra_v, + deepatmo_divzl=deepatmo_divzl, + deepatmo_divzu=deepatmo_divzu, + p_dtime=p_dtime, + rhodz_ast2=rhodz_ast2, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_03.py b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_03.py new file mode 100644 index 0000000000..2e1f83ab08 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_03.py @@ -0,0 +1,50 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.step_advection_stencil_03 import step_advection_stencil_03 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestStepAdvectionStencil03(StencilTest): + PROGRAM = step_advection_stencil_03 + OUTPUTS = ("p_tracer_new",) + + @staticmethod + def reference( + grid, + p_tracer_now: np.array, + p_grf_tend_tracer: np.array, + p_dtime, + **kwargs, + ): + p_tracer_new = p_tracer_now + p_dtime * p_grf_tend_tracer + p_tracer_new = np.where(p_tracer_new < 0.0, 0.0, p_tracer_new) + + return dict(p_tracer_new=p_tracer_new) + + @pytest.fixture + def input_data(self, grid): + p_tracer_now = random_field(grid, CellDim, KDim) + p_grf_tend_tracer = random_field(grid, CellDim, KDim) + p_tracer_new = random_field(grid, CellDim, KDim) + p_dtime = np.float64(5.0) + return dict( + p_tracer_now=p_tracer_now, + p_grf_tend_tracer=p_grf_tend_tracer, + p_dtime=p_dtime, + p_tracer_new=p_tracer_new, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_04.py b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_04.py new file mode 100644 index 0000000000..6541824d3a --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_04.py @@ -0,0 +1,49 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.step_advection_stencil_04 import step_advection_stencil_04 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestStepAdvectionStencil04(StencilTest): + PROGRAM = step_advection_stencil_04 + OUTPUTS = ("opt_ddt_tracer_adv",) + + @staticmethod + def reference( + grid, + p_tracer_now: np.array, + p_tracer_new: np.array, + p_dtime, + **kwargs, + ): + opt_ddt_tracer_adv = (p_tracer_new - p_tracer_now) / p_dtime + + return dict(opt_ddt_tracer_adv=opt_ddt_tracer_adv) + + @pytest.fixture + def input_data(self, grid): + p_tracer_now = random_field(grid, CellDim, KDim) + p_tracer_new = random_field(grid, CellDim, KDim) + opt_ddt_tracer_adv = zero_field(grid, CellDim, KDim) + p_dtime = np.float64(5.0) + return dict( + p_tracer_now=p_tracer_now, + p_tracer_new=p_tracer_new, + p_dtime=p_dtime, + opt_ddt_tracer_adv=opt_ddt_tracer_adv, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py new file mode 100644 index 0000000000..a538479e1e --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py @@ -0,0 +1,194 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.upwind_hflux_miura3_stencil_01 import ( + upwind_hflux_miura3_stencil_01, +) +from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + random_mask, + zero_field, +) + + +class TestUpwindHfluxMiura3Stencil01(StencilTest): + PROGRAM = upwind_hflux_miura3_stencil_01 + OUTPUTS = ("p_out_e_miura3",) + + @staticmethod + def reference( + grid, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, + z_lsq_coeff_4: np.array, + z_lsq_coeff_5: np.array, + z_lsq_coeff_6: np.array, + z_lsq_coeff_7: np.array, + z_lsq_coeff_8: np.array, + z_lsq_coeff_9: np.array, + z_lsq_coeff_10: np.array, + z_quad_vector_sum_1: np.array, + z_quad_vector_sum_2: np.array, + z_quad_vector_sum_3: np.array, + z_quad_vector_sum_4: np.array, + z_quad_vector_sum_5: np.array, + z_quad_vector_sum_6: np.array, + z_quad_vector_sum_7: np.array, + z_quad_vector_sum_8: np.array, + z_quad_vector_sum_9: np.array, + z_quad_vector_sum_10: np.array, + z_dreg_area: np.array, + p_mass_flx_e: np.array, + cell_rel_idx_dsl: np.array, + **kwargs, + ): + e2c = grid.connectivities[E2CDim] + z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] + z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] + z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] + z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] + z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] + z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] + z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] + z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] + z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] + z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] + + p_out_e_miura3 = ( + ( + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_e2c[:, 1], + z_lsq_coeff_1_e2c[:, 0], + ) + * z_quad_vector_sum_1 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_e2c[:, 1], + z_lsq_coeff_2_e2c[:, 0], + ) + * z_quad_vector_sum_2 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_e2c[:, 1], + z_lsq_coeff_3_e2c[:, 0], + ) + * z_quad_vector_sum_3 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_4_e2c[:, 1], + z_lsq_coeff_4_e2c[:, 0], + ) + * z_quad_vector_sum_4 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_5_e2c[:, 1], + z_lsq_coeff_5_e2c[:, 0], + ) + * z_quad_vector_sum_5 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_6_e2c[:, 1], + z_lsq_coeff_6_e2c[:, 0], + ) + * z_quad_vector_sum_6 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_7_e2c[:, 1], + z_lsq_coeff_7_e2c[:, 0], + ) + * z_quad_vector_sum_7 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_8_e2c[:, 1], + z_lsq_coeff_8_e2c[:, 0], + ) + * z_quad_vector_sum_8 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_9_e2c[:, 1], + z_lsq_coeff_9_e2c[:, 0], + ) + * z_quad_vector_sum_9 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_10_e2c[:, 1], + z_lsq_coeff_10_e2c[:, 0], + ) + * z_quad_vector_sum_10 + ) + / z_dreg_area + * p_mass_flx_e + ) + + return dict(p_out_e_miura3=p_out_e_miura3) + + @pytest.fixture + def input_data(self, grid): + z_lsq_coeff_1 = random_field(grid, CellDim, KDim) + z_lsq_coeff_2 = random_field(grid, CellDim, KDim) + z_lsq_coeff_3 = random_field(grid, CellDim, KDim) + z_lsq_coeff_4 = random_field(grid, CellDim, KDim) + z_lsq_coeff_5 = random_field(grid, CellDim, KDim) + z_lsq_coeff_6 = random_field(grid, CellDim, KDim) + z_lsq_coeff_7 = random_field(grid, CellDim, KDim) + z_lsq_coeff_8 = random_field(grid, CellDim, KDim) + z_lsq_coeff_9 = random_field(grid, CellDim, KDim) + z_lsq_coeff_10 = random_field(grid, CellDim, KDim) + z_quad_vector_sum_1 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_2 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_3 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_4 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_5 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_6 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_7 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_8 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_9 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_10 = random_field(grid, EdgeDim, KDim) + p_mass_flx_e = random_field(grid, EdgeDim, KDim) + z_dreg_area = random_field(grid, EdgeDim, KDim) + cell_rel_idx_dsl = random_mask(grid, EdgeDim, KDim, dtype=int32) + p_out_e_miura3 = zero_field(grid, EdgeDim, KDim) + return dict( + z_lsq_coeff_1=z_lsq_coeff_1, + z_lsq_coeff_2=z_lsq_coeff_2, + z_lsq_coeff_3=z_lsq_coeff_3, + z_lsq_coeff_4=z_lsq_coeff_4, + z_lsq_coeff_5=z_lsq_coeff_5, + z_lsq_coeff_6=z_lsq_coeff_6, + z_lsq_coeff_7=z_lsq_coeff_7, + z_lsq_coeff_8=z_lsq_coeff_8, + z_lsq_coeff_9=z_lsq_coeff_9, + z_lsq_coeff_10=z_lsq_coeff_10, + z_quad_vector_sum_1=z_quad_vector_sum_1, + z_quad_vector_sum_2=z_quad_vector_sum_2, + z_quad_vector_sum_3=z_quad_vector_sum_3, + z_quad_vector_sum_4=z_quad_vector_sum_4, + z_quad_vector_sum_5=z_quad_vector_sum_5, + z_quad_vector_sum_6=z_quad_vector_sum_6, + z_quad_vector_sum_7=z_quad_vector_sum_7, + z_quad_vector_sum_8=z_quad_vector_sum_8, + z_quad_vector_sum_9=z_quad_vector_sum_9, + z_quad_vector_sum_10=z_quad_vector_sum_10, + z_dreg_area=z_dreg_area, + p_mass_flx_e=p_mass_flx_e, + cell_rel_idx_dsl=cell_rel_idx_dsl, + p_out_e_miura3=p_out_e_miura3, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py new file mode 100644 index 0000000000..1e33e6a7a7 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py @@ -0,0 +1,92 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_01 import ( + upwind_hflux_miura_cycl_stencil_01, +) +from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + random_mask, + zero_field, +) + + +class TestUpwindHfluxMiuraCyclStencil01(StencilTest): + PROGRAM = upwind_hflux_miura_cycl_stencil_01 + OUTPUTS = ("z_tracer_mflx_dsl",) + + @staticmethod + def reference( + grid, + z_lsq_coeff_1_dsl: np.array, + z_lsq_coeff_2_dsl: np.array, + z_lsq_coeff_3_dsl: np.array, + distv_bary_1: np.array, + distv_bary_2: np.array, + p_mass_flx_e: np.array, + cell_rel_idx_dsl: np.array, + **kwargs, + ): + e2c = grid.connectivities[E2CDim] + z_lsq_coeff_1_dsl_e2c = z_lsq_coeff_1_dsl[e2c] + z_lsq_coeff_2_dsl_e2c = z_lsq_coeff_2_dsl[e2c] + z_lsq_coeff_3_dsl_e2c = z_lsq_coeff_3_dsl[e2c] + + z_tracer_mflx_dsl = ( + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_dsl_e2c[:, 1], + z_lsq_coeff_1_dsl_e2c[:, 0], + ) + + distv_bary_1 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_dsl_e2c[:, 1], + z_lsq_coeff_2_dsl_e2c[:, 0], + ) + + distv_bary_2 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_dsl_e2c[:, 1], + z_lsq_coeff_3_dsl_e2c[:, 0], + ) + ) * p_mass_flx_e + + return dict(z_tracer_mflx_dsl=z_tracer_mflx_dsl) + + @pytest.fixture + def input_data(self, grid): + z_lsq_coeff_1_dsl = random_field(grid, CellDim, KDim) + z_lsq_coeff_2_dsl = random_field(grid, CellDim, KDim) + z_lsq_coeff_3_dsl = random_field(grid, CellDim, KDim) + distv_bary_1 = random_field(grid, EdgeDim, KDim) + distv_bary_2 = random_field(grid, EdgeDim, KDim) + p_mass_flx_e = random_field(grid, EdgeDim, KDim) + cell_rel_idx_dsl = random_mask(grid, EdgeDim, KDim, dtype=int32) + z_tracer_mflx_dsl = zero_field(grid, EdgeDim, KDim) + return dict( + z_lsq_coeff_1_dsl=z_lsq_coeff_1_dsl, + z_lsq_coeff_2_dsl=z_lsq_coeff_2_dsl, + z_lsq_coeff_3_dsl=z_lsq_coeff_3_dsl, + distv_bary_1=distv_bary_1, + distv_bary_2=distv_bary_2, + p_mass_flx_e=p_mass_flx_e, + cell_rel_idx_dsl=cell_rel_idx_dsl, + z_tracer_mflx_dsl=z_tracer_mflx_dsl, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py new file mode 100644 index 0000000000..dbb64d6807 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py @@ -0,0 +1,93 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_02 import ( + upwind_hflux_miura_cycl_stencil_02, +) +from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestUpwindHfluxMiuraCyclStencil02(StencilTest): + PROGRAM = upwind_hflux_miura_cycl_stencil_02 + OUTPUTS = ( + "z_rhofluxdiv_c_out", + "z_fluxdiv_c_dsl", + "z_rho_new_dsl", + "z_tracer_new_dsl", + ) + + @staticmethod + def reference( + grid, + nsub: int32, + p_mass_flx_e: np.array, + geofac_div: np.array, + z_rhofluxdiv_c: np.array, + z_tracer_mflx: np.array, + z_rho_now: np.array, + z_tracer_now: np.array, + z_dtsub: float, + **kwargs, + ): + c2e = grid.connectivities[C2EDim] + p_mass_flx_e_c2e = p_mass_flx_e[c2e] + geofac_div = np.expand_dims(geofac_div, axis=-1) + z_tracer_mflx_c2e = z_tracer_mflx[c2e] + + z_rhofluxdiv_c_out = ( + np.sum(p_mass_flx_e_c2e * geofac_div, axis=1) if nsub == int32(1) else z_rhofluxdiv_c + ) + z_fluxdiv_c_dsl = np.sum(z_tracer_mflx_c2e * geofac_div, axis=1) + z_rho_new_dsl = z_rho_now - z_dtsub * z_rhofluxdiv_c_out + z_tracer_new_dsl = (z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl) / z_rho_new_dsl + + return dict( + z_rhofluxdiv_c_out=z_rhofluxdiv_c_out, + z_fluxdiv_c_dsl=z_fluxdiv_c_dsl, + z_rho_new_dsl=z_rho_new_dsl, + z_tracer_new_dsl=z_tracer_new_dsl, + ) + + @pytest.fixture + def input_data(self, grid): + nsub = int32(1) + p_mass_flx_e = random_field(grid, EdgeDim, KDim) + geofac_div = random_field(grid, CellDim, C2EDim) + z_rhofluxdiv_c = random_field(grid, CellDim, KDim) + z_tracer_mflx = random_field(grid, EdgeDim, KDim) + z_rho_now = random_field(grid, CellDim, KDim) + z_tracer_now = random_field(grid, CellDim, KDim) + z_dtsub = 0.5 + z_rhofluxdiv_c_out = zero_field(grid, CellDim, KDim) + z_fluxdiv_c_dsl = zero_field(grid, CellDim, KDim) + z_rho_new_dsl = zero_field(grid, CellDim, KDim) + z_tracer_new_dsl = zero_field(grid, CellDim, KDim) + return dict( + nsub=nsub, + p_mass_flx_e=p_mass_flx_e, + geofac_div=geofac_div, + z_rhofluxdiv_c=z_rhofluxdiv_c, + z_tracer_mflx=z_tracer_mflx, + z_rho_now=z_rho_now, + z_tracer_now=z_tracer_now, + z_dtsub=z_dtsub, + z_rhofluxdiv_c_out=z_rhofluxdiv_c_out, + z_fluxdiv_c_dsl=z_fluxdiv_c_dsl, + z_rho_new_dsl=z_rho_new_dsl, + z_tracer_new_dsl=z_tracer_new_dsl, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py new file mode 100644 index 0000000000..8c1a99e509 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py @@ -0,0 +1,47 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03a import ( + upwind_hflux_miura_cycl_stencil_03a, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestUpwindHfluxMiuraCyclStencil03a(StencilTest): + PROGRAM = upwind_hflux_miura_cycl_stencil_03a + OUTPUTS = ("p_out_e",) + + @staticmethod + def reference( + grid, + z_tracer_mflx_1_dsl: np.array, + z_tracer_mflx_2_dsl: np.array, + **kwargs, + ): + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl) / float(2) + return dict(p_out_e=p_out_e) + + @pytest.fixture + def input_data(self, grid): + z_tracer_mflx_1_dsl = random_field(grid, EdgeDim, KDim) + z_tracer_mflx_2_dsl = random_field(grid, EdgeDim, KDim) + p_out_e = zero_field(grid, EdgeDim, KDim) + return dict( + z_tracer_mflx_1_dsl=z_tracer_mflx_1_dsl, + z_tracer_mflx_2_dsl=z_tracer_mflx_2_dsl, + p_out_e=p_out_e, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py new file mode 100644 index 0000000000..481fd93c84 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py @@ -0,0 +1,50 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03b import ( + upwind_hflux_miura_cycl_stencil_03b, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestUpwindHfluxMiuraCyclStencil03b(StencilTest): + PROGRAM = upwind_hflux_miura_cycl_stencil_03b + OUTPUTS = ("p_out_e",) + + @staticmethod + def reference( + grid, + z_tracer_mflx_1_dsl: np.array, + z_tracer_mflx_2_dsl: np.array, + z_tracer_mflx_3_dsl: np.array, + **kwargs, + ): + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float(3) + return dict(p_out_e=p_out_e) + + @pytest.fixture + def input_data(self, grid): + z_tracer_mflx_1_dsl = random_field(grid, EdgeDim, KDim) + z_tracer_mflx_2_dsl = random_field(grid, EdgeDim, KDim) + z_tracer_mflx_3_dsl = random_field(grid, EdgeDim, KDim) + p_out_e = zero_field(grid, EdgeDim, KDim) + return dict( + z_tracer_mflx_1_dsl=z_tracer_mflx_1_dsl, + z_tracer_mflx_2_dsl=z_tracer_mflx_2_dsl, + z_tracer_mflx_3_dsl=z_tracer_mflx_3_dsl, + p_out_e=p_out_e, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py new file mode 100644 index 0000000000..59e3c026fe --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py @@ -0,0 +1,92 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.upwind_hflux_miura_stencil_01 import ( + upwind_hflux_miura_stencil_01, +) +from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + constant_field, + random_field, + zero_field, +) + + +class TestUpwindHfluxMiuraStencil01(StencilTest): + PROGRAM = upwind_hflux_miura_stencil_01 + OUTPUTS = ("p_out_e",) + + @staticmethod + def reference( + grid, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, + distv_bary_1: np.array, + distv_bary_2: np.array, + p_mass_flx_e: np.array, + cell_rel_idx_dsl: np.array, + **kwargs, + ): + e2c = grid.connectivities[E2CDim] + z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] + z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] + z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] + + p_out_e = ( + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_e2c[:, 1], + z_lsq_coeff_1_e2c[:, 0], + ) + + distv_bary_1 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_e2c[:, 1], + z_lsq_coeff_2_e2c[:, 0], + ) + + distv_bary_2 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_e2c[:, 1], + z_lsq_coeff_3_e2c[:, 0], + ) + ) * p_mass_flx_e + + return dict(p_out_e=p_out_e) + + @pytest.fixture + def input_data(self, grid): + z_lsq_coeff_1 = random_field(grid, CellDim, KDim) + z_lsq_coeff_2 = random_field(grid, CellDim, KDim) + z_lsq_coeff_3 = random_field(grid, CellDim, KDim) + distv_bary_1 = random_field(grid, EdgeDim, KDim) + distv_bary_2 = random_field(grid, EdgeDim, KDim) + p_mass_flx_e = random_field(grid, EdgeDim, KDim) + cell_rel_idx_dsl = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) + p_out_e = zero_field(grid, EdgeDim, KDim) + return dict( + z_lsq_coeff_1=z_lsq_coeff_1, + z_lsq_coeff_2=z_lsq_coeff_2, + z_lsq_coeff_3=z_lsq_coeff_3, + distv_bary_1=distv_bary_1, + distv_bary_2=distv_bary_2, + p_mass_flx_e=p_mass_flx_e, + cell_rel_idx_dsl=cell_rel_idx_dsl, + p_out_e=p_out_e, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_vflux_ppm_stencil_01.py new file mode 100644 index 0000000000..a8f9c3303b --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_vflux_ppm_stencil_01.py @@ -0,0 +1,45 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.upwind_vflux_ppm_stencil_01 import ( + upwind_vflux_ppm_stencil_01, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestUpwindVfluxPpmStencil01(StencilTest): + PROGRAM = upwind_vflux_ppm_stencil_01 + OUTPUTS = ("z_delta_q", "z_a1") + + @staticmethod + def reference( + grid, z_face_up: np.ndarray, z_face_low: np.ndarray, p_cc: np.ndarray, **kwargs + ) -> tuple[np.ndarray]: + z_delta_q = 0.5 * (z_face_up - z_face_low) + z_a1 = p_cc - 0.5 * (z_face_up + z_face_low) + return dict(z_delta_q=z_delta_q, z_a1=z_a1) + + @pytest.fixture + def input_data(self, grid): + z_face_up = random_field(grid, CellDim, KDim) + z_face_low = random_field(grid, CellDim, KDim) + p_cc = random_field(grid, CellDim, KDim) + z_delta_q = zero_field(grid, CellDim, KDim) + z_a1 = zero_field(grid, CellDim, KDim) + return dict( + z_face_up=z_face_up, z_face_low=z_face_low, p_cc=p_cc, z_delta_q=z_delta_q, z_a1=z_a1 + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py new file mode 100644 index 0000000000..598d496e3c --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py @@ -0,0 +1,64 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.vert_adv_stencil_01 import vert_adv_stencil_01 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestVertAdvStencil01(StencilTest): + PROGRAM = vert_adv_stencil_01 + OUTPUTS = ("tracer_new",) + + @staticmethod + def reference( + grid, + tracer_now: np.array, + rhodz_now: np.array, + p_mflx_tracer_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, + rhodz_new: np.array, + p_dtime, + **kwargs, + ) -> np.array: + tracer_new = ( + tracer_now * rhodz_now + + p_dtime + * (p_mflx_tracer_v[:, 1:] * deepatmo_divzl - p_mflx_tracer_v[:, :-1] * deepatmo_divzu) + ) / rhodz_new + return dict(tracer_new=tracer_new) + + @pytest.fixture + def input_data(self, grid): + tracer_now = random_field(grid, CellDim, KDim) + rhodz_now = random_field(grid, CellDim, KDim) + p_mflx_tracer_v = random_field(grid, CellDim, KDim, extend={KDim: 1}) + deepatmo_divzl = random_field(grid, KDim) + deepatmo_divzu = random_field(grid, KDim) + rhodz_new = random_field(grid, CellDim, KDim) + p_dtime = np.float64(5.0) + tracer_new = zero_field(grid, CellDim, KDim) + return dict( + tracer_now=tracer_now, + rhodz_now=rhodz_now, + p_mflx_tracer_v=p_mflx_tracer_v, + deepatmo_divzl=deepatmo_divzl, + deepatmo_divzu=deepatmo_divzu, + rhodz_new=rhodz_new, + p_dtime=p_dtime, + tracer_new=tracer_new, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_01.py new file mode 100644 index 0000000000..f6fc3c2ff8 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_01.py @@ -0,0 +1,39 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.v_limit_prbl_sm_stencil_01 import v_limit_prbl_sm_stencil_01 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestVLimitPrblSmStencil01(StencilTest): + PROGRAM = v_limit_prbl_sm_stencil_01 + OUTPUTS = ("l_limit",) + + @staticmethod + def reference(grid, p_face: np.array, p_cc: np.array, **kwargs): + z_delta = p_face[:, :-1] - p_face[:, 1:] + z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:, :-1] + p_face[:, 1:])) + l_limit = np.where(np.abs(z_delta) < -1 * z_a6i, int32(1), int32(0)) + return dict(l_limit=l_limit) + + @pytest.fixture + def input_data(self, grid): + p_cc = random_field(grid, CellDim, KDim) + p_face = random_field(grid, CellDim, KDim, extend={KDim: 1}) + l_limit = zero_field(grid, CellDim, KDim, dtype=int32) + return dict(p_face=p_face, p_cc=p_cc, l_limit=l_limit) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_02.py new file mode 100644 index 0000000000..c79595b338 --- /dev/null +++ b/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_02.py @@ -0,0 +1,58 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.v_limit_prbl_sm_stencil_02 import v_limit_prbl_sm_stencil_02 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + random_mask, + zero_field, +) + + +class TestVLimitPrblSmStencil02(StencilTest): + PROGRAM = v_limit_prbl_sm_stencil_02 + OUTPUTS = ("p_face_up", "p_face_low") + + @staticmethod + def reference(grid, l_limit: np.array, p_face: np.array, p_cc: np.array, **kwargs): + q_face_up, q_face_low = np.where( + l_limit != int32(0), + np.where( + (p_cc < np.minimum(p_face[:, :-1], p_face[:, 1:])), + (p_cc, p_cc), + np.where( + p_face[:, :-1] > p_face[:, 1:], + (3.0 * p_cc - 2.0 * p_face[:, 1:], p_face[:, 1:]), + (p_face[:, :-1], 3.0 * p_cc - 2.0 * p_face[:, :-1]), + ), + ), + (p_face[:, :-1], p_face[:, 1:]), + ) + return dict(p_face_up=q_face_up, p_face_low=q_face_low) + + @pytest.fixture + def input_data(self, grid): + l_limit = random_mask(grid, CellDim, KDim, dtype=int32) + p_cc = random_field(grid, CellDim, KDim) + p_face = random_field(grid, CellDim, KDim, extend={KDim: 1}) + p_face_up = zero_field(grid, CellDim, KDim) + p_face_low = zero_field(grid, CellDim, KDim) + return dict( + l_limit=l_limit, p_cc=p_cc, p_face=p_face, p_face_up=p_face_up, p_face_low=p_face_low + ) diff --git a/model/common/src/icon4py/model/common/test_utils/helpers.py b/model/common/src/icon4py/model/common/test_utils/helpers.py index 1b6f757b41..470d30503a 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -239,7 +239,7 @@ class StencilTest: """ PROGRAM: ClassVar[Program] - OUTPUTS: ClassVar[tuple[str | Output, ...]] + OUTPUTS: ClassVar[tuple[str, ...]] def __init_subclass__(cls, **kwargs): # Add two methods for verification and benchmarking. In order to have names that diff --git a/model/tox.ini b/model/tox.ini index 39ce4ab666..2718b464e4 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -30,18 +30,17 @@ allowlist_externals = commands = pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/diffusion/tests/diffusion_stencil_tests --benchmark-skip -n auto {posargs} pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/dycore/tests/dycore_stencil_tests --benchmark-skip -n auto {posargs} - pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/advection/tests/advection_stencil_tests --benchmark-skip -n auto {posargs} [testenv:run_benchmarks] commands = pytest -s -m "not slow_tests" atmosphere/diffusion/tests/diffusion_stencil_tests --benchmark-only {posargs} pytest -s -m "not slow_tests" atmosphere/dycore/tests/dycore_stencil_tests --benchmark-only {posargs} - pytest -s -m "not slow_tests" atmosphere/advection/tests/advection_stencil_tests --benchmark-only {posargs} [testenv:run_model_tests] commands = pytest -v -s -m "not slow_tests" --datatest {posargs} + [testenv:dev] setenv = PIP_SRC = _external_src From b216479da50deb3c217edec6c6333a84dde866f8 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 15 May 2024 15:27:01 +0200 Subject: [PATCH 53/53] rm utest --- .../tests/advection_stencil_tests/__init__.py | 12 - .../test_btraj_dreg_stencil_01.py | 52 - .../test_btraj_dreg_stencil_02.py | 66 - .../test_btraj_dreg_stencil_03.py | 233 --- .../test_divide_flux_area_list_stencil_01.py | 1551 ----------------- .../test_divide_flux_area_list_stencil_02.py | 290 --- .../test_face_val_ppm_stencil_01.py | 78 - .../test_face_val_ppm_stencil_02.py | 76 - .../test_face_val_ppm_stencil_02a.py | 43 - .../test_face_val_ppm_stencil_02b.py | 35 - .../test_face_val_ppm_stencil_02c.py | 52 - .../test_face_val_ppm_stencil_05.py | 71 - .../test_hflux_ffsl_hybrid_stencil_01a.py | 184 -- .../test_hflux_ffsl_hybrid_stencil_02.py | 49 - .../test_hflx_limiter_mo_stencil_01a.py | 61 - .../test_hflx_limiter_mo_stencil_01b.py | 120 -- .../test_hflx_limiter_mo_stencil_02.py | 85 - .../test_hflx_limiter_mo_stencil_03.py | 101 -- .../test_hflx_limiter_mo_stencil_04.py | 50 - .../test_hflx_limiter_pd_stencil_01.py | 77 - .../test_hflx_limiter_pd_stencil_02.py | 44 - .../test_hor_adv_stencil_01.py | 74 - .../test_init_zero_c_k.py | 32 - ..._mo_advection_traj_btraj_compute_o1_dsl.py | 154 -- ...st_prep_gauss_quadrature_c_list_stencil.py | 636 ------- .../test_prep_gauss_quadrature_c_stencil.py | 591 ------- .../test_rbf_intp_edge_stencil_01.py | 37 - .../test_recon_lsq_cell_c_svd_stencil.py | 318 ---- .../test_recon_lsq_cell_l_svd_stencil.py | 77 - .../test_step_advection_stencil_01.py | 64 - .../test_step_advection_stencil_02.py | 61 - .../test_step_advection_stencil_03.py | 50 - .../test_step_advection_stencil_04.py | 49 - .../test_upwind_hflux_miura3_stencil_01.py | 194 --- ...test_upwind_hflux_miura_cycl_stencil_01.py | 92 - ...test_upwind_hflux_miura_cycl_stencil_02.py | 93 - ...est_upwind_hflux_miura_cycl_stencil_03a.py | 47 - ...est_upwind_hflux_miura_cycl_stencil_03b.py | 50 - .../test_upwind_hflux_miura_stencil_01.py | 92 - .../test_upwind_vflux_ppm_stencil_01.py | 45 - .../test_vert_adv_stencil_01.py | 87 - .../test_vlimit_prbl_sm_stencil_01.py | 39 - .../test_vlimit_prbl_sm_stencil_02.py | 58 - 43 files changed, 6270 deletions(-) delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/__init__.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_01.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_02.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_03.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_02.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_01.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02a.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02b.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02c.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_05.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01a.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01b.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_02.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_03.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_04.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_01.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_hor_adv_stencil_01.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_init_zero_c_k.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_stencil.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_rbf_intp_edge_stencil_01.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_c_svd_stencil.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_l_svd_stencil.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_01.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_02.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_03.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_04.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura3_stencil_01.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_stencil_01.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_vflux_ppm_stencil_01.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_vert_adv_stencil_01.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_01.py delete mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_02.py diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/__init__.py b/model/atmosphere/advection/tests/advection_stencil_tests/__init__.py deleted file mode 100644 index 15dfdb0098..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_01.py deleted file mode 100644 index 229ac10410..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_01.py +++ /dev/null @@ -1,52 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.btraj_dreg_stencil_01 import btraj_dreg_stencil_01 -from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestBtrajDregStencil01(StencilTest): - PROGRAM = btraj_dreg_stencil_01 - OUTPUTS = ("lvn_sys_pos",) - - @staticmethod - def reference( - grid, lcounterclock: bool, p_vn: np.array, tangent_orientation: np.array, **kwargs - ): - tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) - - tangent_orientation = np.broadcast_to(tangent_orientation, p_vn.shape) - - lvn_sys_pos_true = np.where(tangent_orientation * p_vn >= 0.0, True, False) - - mask_lcounterclock = np.broadcast_to(lcounterclock, p_vn.shape) - - lvn_sys_pos = np.where(mask_lcounterclock, lvn_sys_pos_true, False) - - return dict(lvn_sys_pos=lvn_sys_pos) - - @pytest.fixture - def input_data(self, grid): - lcounterclock = True - p_vn = random_field(grid, EdgeDim, KDim) - tangent_orientation = random_field(grid, EdgeDim) - lvn_sys_pos = zero_field(grid, EdgeDim, KDim, dtype=bool) - return dict( - lcounterclock=lcounterclock, - p_vn=p_vn, - tangent_orientation=tangent_orientation, - lvn_sys_pos=lvn_sys_pos, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_02.py deleted file mode 100644 index d602eae212..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_02.py +++ /dev/null @@ -1,66 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.btraj_dreg_stencil_02 import btraj_dreg_stencil_02 -from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - numpy_to_1D_sparse_field, - random_field, - zero_field, -) - - -class TestBtrajDregStencil02(StencilTest): - PROGRAM = btraj_dreg_stencil_02 - OUTPUTS = ("opt_famask_dsl",) - - @staticmethod - def reference(grid, p_vn: np.array, p_vt: np.array, p_dt: float, **kwargs): - lvn_pos = np.where(p_vn >= 0.0, True, False) - - traj_length = np.sqrt(p_vn**2 + p_vt**2) * p_dt - - edge_cell_length = np.expand_dims( - np.asarray(grid.connectivities[E2CDim], dtype=float), axis=-1 - ) - e2c_length = np.where(lvn_pos, edge_cell_length[:, 0], edge_cell_length[:, 1]) - - opt_famask_dsl = np.where( - traj_length > (1.25 * np.broadcast_to(e2c_length, p_vn.shape)), - int32(1), - int32(0), - ) - - return dict(opt_famask_dsl=opt_famask_dsl) - - @pytest.fixture - def input_data(self, grid): - p_vn = random_field(grid, EdgeDim, KDim) - p_vt = random_field(grid, EdgeDim, KDim) - edge_cell_length = np.asarray(grid.connectivities[E2CDim], dtype=float) - edge_cell_length_new = numpy_to_1D_sparse_field(edge_cell_length, ECDim) - p_dt = 1.0 - opt_famask_dsl = zero_field(grid, EdgeDim, KDim, dtype=int32) - - return dict( - p_vn=p_vn, - p_vt=p_vt, - edge_cell_length=edge_cell_length_new, - p_dt=p_dt, - opt_famask_dsl=opt_famask_dsl, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_03.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_03.py deleted file mode 100644 index 927bccd654..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_btraj_dreg_stencil_03.py +++ /dev/null @@ -1,233 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.btraj_dreg_stencil_03 import btraj_dreg_stencil_03 -from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - as_1D_sparse_field, - constant_field, - numpy_to_1D_sparse_field, - random_field, - reshape, -) - - -class TestBtrajDregStencil03(StencilTest): - PROGRAM = btraj_dreg_stencil_03 - OUTPUTS = ( - "p_cell_idx", - "p_cell_rel_idx_dsl", - "p_cell_blk", - "p_coords_dreg_v_1_lon_dsl", - "p_coords_dreg_v_2_lon_dsl", - "p_coords_dreg_v_3_lon_dsl", - "p_coords_dreg_v_4_lon_dsl", - "p_coords_dreg_v_1_lat_dsl", - "p_coords_dreg_v_2_lat_dsl", - "p_coords_dreg_v_3_lat_dsl", - "p_coords_dreg_v_4_lat_dsl", - ) - - @staticmethod - def reference( - grid, - p_vn: np.array, - p_vt: np.array, - cell_idx: np.array, - cell_blk: np.array, - edge_verts_1_x: np.array, - edge_verts_2_x: np.array, - edge_verts_1_y: np.array, - edge_verts_2_y: np.array, - pos_on_tplane_e_1_x: np.array, - pos_on_tplane_e_2_x: np.array, - pos_on_tplane_e_1_y: np.array, - pos_on_tplane_e_2_y: np.array, - primal_normal_cell_x: np.array, - primal_normal_cell_y: np.array, - dual_normal_cell_x: np.array, - dual_normal_cell_y: np.array, - lvn_sys_pos: np.array, - p_dt: float, - **kwargs, - ): - e2c_shape = grid.connectivities[E2CDim].shape - cell_idx = reshape(cell_idx, e2c_shape) - cell_blk = reshape(cell_blk, e2c_shape) - primal_normal_cell_x = reshape(primal_normal_cell_x, e2c_shape) - dual_normal_cell_x = reshape(dual_normal_cell_x, e2c_shape) - primal_normal_cell_y = reshape(primal_normal_cell_y, e2c_shape) - dual_normal_cell_y = reshape(dual_normal_cell_y, e2c_shape) - - lvn_pos = np.where(p_vn >= 0.0, True, False) - cell_idx = np.expand_dims(cell_idx, axis=-1) - cell_blk = np.expand_dims(cell_blk, axis=-1) - primal_normal_cell_x = np.expand_dims(primal_normal_cell_x, axis=-1) - dual_normal_cell_x = np.expand_dims(dual_normal_cell_x, axis=-1) - primal_normal_cell_y = np.expand_dims(primal_normal_cell_y, axis=-1) - dual_normal_cell_y = np.expand_dims(dual_normal_cell_y, axis=-1) - edge_verts_1_x = np.expand_dims(edge_verts_1_x, axis=-1) - edge_verts_1_y = np.expand_dims(edge_verts_1_y, axis=-1) - edge_verts_2_x = np.expand_dims(edge_verts_2_x, axis=-1) - edge_verts_2_y = np.expand_dims(edge_verts_2_y, axis=-1) - pos_on_tplane_e_1_x = np.expand_dims(pos_on_tplane_e_1_x, axis=-1) - pos_on_tplane_e_1_y = np.expand_dims(pos_on_tplane_e_1_y, axis=-1) - pos_on_tplane_e_2_x = np.expand_dims(pos_on_tplane_e_2_x, axis=-1) - pos_on_tplane_e_2_y = np.expand_dims(pos_on_tplane_e_2_y, axis=-1) - - p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) - p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) - p_cell_rel_idx_dsl = np.where(lvn_pos, int32(0), int32(1)) - - depart_pts_1_x = np.broadcast_to(edge_verts_1_x, p_vn.shape) - p_vn * p_dt - depart_pts_1_y = np.broadcast_to(edge_verts_1_y, p_vn.shape) - p_vt * p_dt - depart_pts_2_x = np.broadcast_to(edge_verts_2_x, p_vn.shape) - p_vn * p_dt - depart_pts_2_y = np.broadcast_to(edge_verts_2_y, p_vn.shape) - p_vt * p_dt - - pos_on_tplane_e_x = np.where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) - pos_on_tplane_e_y = np.where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) - - pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x - pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y - pos_dreg_vert_c_2_x = ( - np.where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x - ) - pos_dreg_vert_c_2_y = ( - np.where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y - ) - pos_dreg_vert_c_3_x = depart_pts_2_x - pos_on_tplane_e_x - pos_dreg_vert_c_3_y = depart_pts_2_y - pos_on_tplane_e_y - pos_dreg_vert_c_4_x = ( - np.where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x - ) - pos_dreg_vert_c_4_y = ( - np.where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y - ) - - pn_cell_1 = np.where(lvn_pos, primal_normal_cell_x[:, 0], primal_normal_cell_x[:, 1]) - pn_cell_2 = np.where(lvn_pos, primal_normal_cell_y[:, 0], primal_normal_cell_y[:, 1]) - dn_cell_1 = np.where(lvn_pos, dual_normal_cell_x[:, 0], dual_normal_cell_x[:, 1]) - dn_cell_2 = np.where(lvn_pos, dual_normal_cell_y[:, 0], dual_normal_cell_y[:, 1]) - - p_coords_dreg_v_1_lon_dsl = ( - pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 - ) - p_coords_dreg_v_2_lon_dsl = ( - pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 - ) - p_coords_dreg_v_3_lon_dsl = ( - pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 - ) - p_coords_dreg_v_4_lon_dsl = ( - pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 - ) - p_coords_dreg_v_1_lat_dsl = ( - pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 - ) - p_coords_dreg_v_2_lat_dsl = ( - pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 - ) - p_coords_dreg_v_3_lat_dsl = ( - pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 - ) - p_coords_dreg_v_4_lat_dsl = ( - pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 - ) - return { - "p_cell_idx": p_cell_idx, - "p_cell_rel_idx_dsl": p_cell_rel_idx_dsl, - "p_cell_blk": p_cell_blk, - "p_coords_dreg_v_1_lon_dsl": p_coords_dreg_v_1_lon_dsl, - "p_coords_dreg_v_2_lon_dsl": p_coords_dreg_v_2_lon_dsl, - "p_coords_dreg_v_3_lon_dsl": p_coords_dreg_v_3_lon_dsl, - "p_coords_dreg_v_4_lon_dsl": p_coords_dreg_v_4_lon_dsl, - "p_coords_dreg_v_1_lat_dsl": p_coords_dreg_v_1_lat_dsl, - "p_coords_dreg_v_2_lat_dsl": p_coords_dreg_v_2_lat_dsl, - "p_coords_dreg_v_3_lat_dsl": p_coords_dreg_v_3_lat_dsl, - "p_coords_dreg_v_4_lat_dsl": p_coords_dreg_v_4_lat_dsl, - } - - @pytest.fixture - def input_data(self, grid): - p_vn = random_field(grid, EdgeDim, KDim) - p_vt = random_field(grid, EdgeDim, KDim) - cell_idx = np.asarray(grid.connectivities[E2CDim], dtype=int32) - cell_idx_new = numpy_to_1D_sparse_field(cell_idx, ECDim) - cell_blk = constant_field(grid, 1, EdgeDim, E2CDim, dtype=int32) - cell_blk_new = as_1D_sparse_field(cell_blk, ECDim) - - edge_verts_1_x = random_field(grid, EdgeDim) - edge_verts_2_x = random_field(grid, EdgeDim) - edge_verts_1_y = random_field(grid, EdgeDim) - edge_verts_2_y = random_field(grid, EdgeDim) - pos_on_tplane_e_1_x = random_field(grid, EdgeDim) - pos_on_tplane_e_2_x = random_field(grid, EdgeDim) - pos_on_tplane_e_1_y = random_field(grid, EdgeDim) - pos_on_tplane_e_2_y = random_field(grid, EdgeDim) - primal_normal_cell_x = random_field(grid, EdgeDim, E2CDim) - primal_normal_cell_x_new = as_1D_sparse_field(primal_normal_cell_x, ECDim) - dual_normal_cell_x = random_field(grid, EdgeDim, E2CDim) - dual_normal_cell_x_new = as_1D_sparse_field(dual_normal_cell_x, ECDim) - primal_normal_cell_y = random_field(grid, EdgeDim, E2CDim) - primal_normal_cell_y_new = as_1D_sparse_field(primal_normal_cell_y, ECDim) - dual_normal_cell_y = random_field(grid, EdgeDim, E2CDim) - dual_normal_cell_y_new = as_1D_sparse_field(dual_normal_cell_y, ECDim) - lvn_sys_pos = constant_field(grid, True, EdgeDim, KDim, dtype=bool) - p_dt = 2.0 - p_cell_idx = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) - p_cell_rel_idx_dsl = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) - p_cell_blk = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) - p_coords_dreg_v_1_lon_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_2_lon_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_3_lon_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_4_lon_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_1_lat_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_2_lat_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_3_lat_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_4_lat_dsl = random_field(grid, EdgeDim, KDim) - - return dict( - p_vn=p_vn, - p_vt=p_vt, - cell_idx=cell_idx_new, - cell_blk=cell_blk_new, - edge_verts_1_x=edge_verts_1_x, - edge_verts_2_x=edge_verts_2_x, - edge_verts_1_y=edge_verts_1_y, - edge_verts_2_y=edge_verts_2_y, - pos_on_tplane_e_1_x=pos_on_tplane_e_1_x, - pos_on_tplane_e_2_x=pos_on_tplane_e_2_x, - pos_on_tplane_e_1_y=pos_on_tplane_e_1_y, - pos_on_tplane_e_2_y=pos_on_tplane_e_2_y, - primal_normal_cell_x=primal_normal_cell_x_new, - primal_normal_cell_y=primal_normal_cell_y_new, - dual_normal_cell_x=dual_normal_cell_x_new, - dual_normal_cell_y=dual_normal_cell_y_new, - lvn_sys_pos=lvn_sys_pos, - p_dt=p_dt, - p_cell_idx=p_cell_idx, - p_cell_rel_idx_dsl=p_cell_rel_idx_dsl, - p_cell_blk=p_cell_blk, - p_coords_dreg_v_1_lon_dsl=p_coords_dreg_v_1_lon_dsl, - p_coords_dreg_v_2_lon_dsl=p_coords_dreg_v_2_lon_dsl, - p_coords_dreg_v_3_lon_dsl=p_coords_dreg_v_3_lon_dsl, - p_coords_dreg_v_4_lon_dsl=p_coords_dreg_v_4_lon_dsl, - p_coords_dreg_v_1_lat_dsl=p_coords_dreg_v_1_lat_dsl, - p_coords_dreg_v_2_lat_dsl=p_coords_dreg_v_2_lat_dsl, - p_coords_dreg_v_3_lat_dsl=p_coords_dreg_v_3_lat_dsl, - p_coords_dreg_v_4_lat_dsl=p_coords_dreg_v_4_lat_dsl, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py deleted file mode 100644 index 034799fbb1..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_01.py +++ /dev/null @@ -1,1551 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.divide_flux_area_list_stencil_01 import ( - divide_flux_area_list_stencil_01, -) -from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - as_1D_sparse_field, - constant_field, - random_field, - random_mask, - reshape, - zero_field, -) - - -# Check whether lines inters. -def _ccw_numpy( - p0_lon, - p0_lat, - p1_lon, - p1_lat, - p2_lon, - p2_lat, -): - dx1 = p1_lon - p0_lon - dy1 = p1_lat - p0_lat - - dx2 = p2_lon - p0_lon - dy2 = p2_lat - p0_lat - - dx1dy2 = dx1 * dy2 - dy1dx2 = dy1 * dx2 - - lccw = np.where(dx1dy2 > dy1dx2, True, False) - ccw_out = np.where(lccw, int32(1), int32(-1)) # 1: clockwise, -1: counterclockwise - return ccw_out - - -# Check whether two lines intersect -def _lintersect_numpy( - line1_p1_lon, - line1_p1_lat, - line1_p2_lon, - line1_p2_lat, - line2_p1_lon, - line2_p1_lat, - line2_p2_lon, - line2_p2_lat, -): - intersect1 = _ccw_numpy( - line1_p1_lon, - line1_p1_lat, - line1_p2_lon, - line1_p2_lat, - line2_p1_lon, - line2_p1_lat, - ) * _ccw_numpy( - line1_p1_lon, - line1_p1_lat, - line1_p2_lon, - line1_p2_lat, - line2_p2_lon, - line2_p2_lat, - ) - intersect2 = _ccw_numpy( - line2_p1_lon, - line2_p1_lat, - line2_p2_lon, - line2_p2_lat, - line1_p1_lon, - line1_p1_lat, - ) * _ccw_numpy( - line2_p1_lon, - line2_p1_lat, - line2_p2_lon, - line2_p2_lat, - line1_p2_lon, - line1_p2_lat, - ) - lintersect_out = np.where((intersect1 + intersect2) == -2, True, False) - - return lintersect_out - - -# Compute intersection point of two lines in 2D -def _line_intersect_numpy( - line1_p1_lon, - line1_p1_lat, - line1_p2_lon, - line1_p2_lat, - line2_p1_lon, - line2_p1_lat, - line2_p2_lon, - line2_p2_lat, -): - d1 = line1_p2_lon - line1_p1_lon - d1 = np.where(d1 != float(0), d1, line1_p2_lon) - - d2 = line2_p2_lon - line2_p1_lon - d2 = np.where(d2 != float(0), d2, line2_p2_lon) - - m1 = (line1_p2_lat - line1_p1_lat) / d1 - m2 = (line2_p2_lat - line2_p1_lat) / d2 - - intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / (m1 - m2) - intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) - - return intersect_1, intersect_2 - - -@pytest.mark.slow_tests -class TestDivideFluxAreaListStencil01(StencilTest): - PROGRAM = divide_flux_area_list_stencil_01 - OUTPUTS = ( - "dreg_patch0_1_lon_dsl", - "dreg_patch0_1_lat_dsl", - "dreg_patch0_2_lon_dsl", - "dreg_patch0_2_lat_dsl", - "dreg_patch0_3_lon_dsl", - "dreg_patch0_3_lat_dsl", - "dreg_patch0_4_lon_dsl", - "dreg_patch0_4_lat_dsl", - "dreg_patch1_1_lon_vmask", - "dreg_patch1_1_lat_vmask", - "dreg_patch1_2_lon_vmask", - "dreg_patch1_2_lat_vmask", - "dreg_patch1_3_lon_vmask", - "dreg_patch1_3_lat_vmask", - "dreg_patch1_4_lon_vmask", - "dreg_patch1_4_lat_vmask", - "dreg_patch2_1_lon_vmask", - "dreg_patch2_1_lat_vmask", - "dreg_patch2_2_lon_vmask", - "dreg_patch2_2_lat_vmask", - "dreg_patch2_3_lon_vmask", - "dreg_patch2_3_lat_vmask", - "dreg_patch2_4_lon_vmask", - "dreg_patch2_4_lat_vmask", - ) - - @staticmethod - def _generate_flux_area_geometry( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - p_vn, - ptr_v3_lon_e, - ptr_v3_lat_e, - ): - arrival_pts_1_lon_dsl = dreg_patch0_1_lon_dsl - arrival_pts_1_lat_dsl = dreg_patch0_1_lat_dsl - arrival_pts_2_lon_dsl = dreg_patch0_2_lon_dsl - arrival_pts_2_lat_dsl = dreg_patch0_2_lat_dsl - depart_pts_1_lon_dsl = dreg_patch0_4_lon_dsl - depart_pts_1_lat_dsl = dreg_patch0_4_lat_dsl - depart_pts_2_lon_dsl = dreg_patch0_3_lon_dsl - depart_pts_2_lat_dsl = dreg_patch0_3_lat_dsl - - lvn_pos = np.where(p_vn >= 0.0, True, False) - - fl_line_p1_lon = depart_pts_1_lon_dsl - fl_line_p1_lat = depart_pts_1_lat_dsl - fl_line_p2_lon = depart_pts_2_lon_dsl - fl_line_p2_lat = depart_pts_2_lat_dsl - - tri_line1_p1_lon = arrival_pts_1_lon_dsl - tri_line1_p1_lat = arrival_pts_1_lat_dsl - tri_line1_p2_lon = np.where( - lvn_pos, - np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), - np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), - ) - tri_line1_p2_lat = np.where( - lvn_pos, - np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), - np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), - ) - - tri_line2_p1_lon = arrival_pts_2_lon_dsl - tri_line2_p1_lat = arrival_pts_2_lat_dsl - tri_line2_p2_lon = np.where( - lvn_pos, - np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), - np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), - ) - tri_line2_p2_lat = np.where( - lvn_pos, - np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), - np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), - ) - - return ( - fl_line_p1_lon, - fl_line_p1_lat, - fl_line_p2_lon, - fl_line_p2_lat, - tri_line1_p1_lon, - tri_line1_p1_lat, - tri_line1_p2_lon, - tri_line1_p2_lat, - tri_line2_p1_lon, - tri_line2_p1_lat, - tri_line2_p2_lon, - tri_line2_p2_lat, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - ) - - @staticmethod - def _apply_case1_patch0( - mask_case1, - lvn_sys_pos, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - ps1_x, - ps1_y, - ps2_x, - ps2_y, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - ): - dreg_patch0_1_lon_dsl = arrival_pts_1_lon_dsl - dreg_patch0_1_lat_dsl = arrival_pts_1_lat_dsl - dreg_patch0_2_lon_dsl = np.where( - mask_case1, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), - arrival_pts_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case1, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), - arrival_pts_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where(mask_case1, ps2_x, depart_pts_2_lon_dsl) - dreg_patch0_3_lat_dsl = np.where(mask_case1, ps2_y, depart_pts_2_lat_dsl) - dreg_patch0_4_lon_dsl = np.where( - mask_case1, - np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), - depart_pts_1_lon_dsl, - ) - dreg_patch0_4_lat_dsl = np.where( - mask_case1, - np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), - depart_pts_1_lat_dsl, - ) - - return ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) - - @staticmethod - def _apply_case1_patch1( - mask_case1, - lvn_sys_pos, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - ps1_x, - ps1_y, - ): - dreg_patch1_1_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) - dreg_patch1_1_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) - dreg_patch1_4_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) - dreg_patch1_4_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) - dreg_patch1_2_lon_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 - ) - dreg_patch1_2_lat_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 - ) - dreg_patch1_3_lon_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 - ) - dreg_patch1_3_lat_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 - ) - - return ( - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - ) - - @staticmethod - def _apply_case1_patch2( - mask_case1, - lvn_sys_pos, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - ps2_x, - ps2_y, - ): - # Case 1 - patch 2 - dreg_patch2_1_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) - dreg_patch2_1_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) - dreg_patch2_4_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) - dreg_patch2_4_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) - dreg_patch2_2_lon_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0 - ) - dreg_patch2_2_lat_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0 - ) - dreg_patch2_3_lon_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0 - ) - dreg_patch2_3_lat_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0 - ) - - return ( - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - ) - - @staticmethod - def _apply_case2a_patch0( - mask_case2a, - lvn_sys_pos, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - ps1_x, - ps1_y, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ): - dreg_patch0_1_lon_dsl = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = np.where( - mask_case2a, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), - dreg_patch0_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case2a, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), - dreg_patch0_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = np.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = np.where( - mask_case2a, - np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), - dreg_patch0_4_lon_dsl, - ) - dreg_patch0_4_lat_dsl = np.where( - mask_case2a, - np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), - dreg_patch0_4_lat_dsl, - ) - - return ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) - - @staticmethod - def _apply_case2a_patch1( - mask_case2a, - lvn_sys_pos, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - ps1_x, - ps1_y, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - ): - dreg_patch1_1_lon_vmask = np.where( - mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask - ) - dreg_patch1_1_lat_vmask = np.where( - mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask - ) - dreg_patch1_4_lon_vmask = np.where( - mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask - ) - dreg_patch1_4_lat_vmask = np.where( - mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask - ) - dreg_patch1_2_lon_vmask = np.where( - mask_case2a, - np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), - dreg_patch1_2_lon_vmask, - ) - dreg_patch1_2_lat_vmask = np.where( - mask_case2a, - np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), - dreg_patch1_2_lat_vmask, - ) - dreg_patch1_3_lon_vmask = np.where( - mask_case2a, - np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), - dreg_patch1_3_lon_vmask, - ) - dreg_patch1_3_lat_vmask = np.where( - mask_case2a, - np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), - dreg_patch1_3_lat_vmask, - ) - - return ( - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - ) - - @staticmethod - def _apply_case2b_patch0( - mask_case2b, - lvn_sys_pos, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - ps2_x, - ps2_y, - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ): - dreg_patch0_1_lon_dsl = np.where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = np.where( - mask_case2b, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), - dreg_patch0_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case2b, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), - dreg_patch0_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = np.where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = np.where( - mask_case2b, - np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), - dreg_patch0_4_lon_dsl, - ) - dreg_patch0_4_lat_dsl = np.where( - mask_case2b, - np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), - dreg_patch0_4_lat_dsl, - ) - - return ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) - - @staticmethod - def _apply_case2b_patch1( - mask_case2b, - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - ): - zeros_array = np.zeros_like(mask_case2b) - - dreg_patch1_1_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_1_lat_vmask) - dreg_patch1_2_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_2_lon_vmask) - dreg_patch1_2_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_2_lat_vmask) - dreg_patch1_3_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_3_lat_vmask) - dreg_patch1_4_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_4_lat_vmask) - - return ( - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - ) - - @staticmethod - def _apply_case2b_patch2( - mask_case2b, - lvn_sys_pos, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - ps2_x, - ps2_y, - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - ): - dreg_patch2_1_lon_vmask = np.where( - mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask - ) - dreg_patch2_1_lat_vmask = np.where( - mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask - ) - dreg_patch2_4_lon_vmask = np.where( - mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask - ) - dreg_patch2_4_lat_vmask = np.where( - mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask - ) - dreg_patch2_2_lon_vmask = np.where( - mask_case2b, - np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), - dreg_patch2_2_lon_vmask, - ) - dreg_patch2_2_lat_vmask = np.where( - mask_case2b, - np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), - dreg_patch2_2_lat_vmask, - ) - dreg_patch2_3_lon_vmask = np.where( - mask_case2b, - np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), - dreg_patch2_3_lon_vmask, - ) - dreg_patch2_3_lat_vmask = np.where( - mask_case2b, - np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), - dreg_patch2_3_lat_vmask, - ) - - return ( - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - ) - - @staticmethod - def _apply_case3a_patch0( - mask_case3a, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - lvn_sys_pos, - ps2_x, - ps2_y, - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ): - dreg_patch0_1_lon_dsl = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = np.where( - mask_case3a, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), - dreg_patch0_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case3a, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), - dreg_patch0_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = np.where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = np.where( - mask_case3a, - np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), - dreg_patch0_4_lon_dsl, - ) - dreg_patch0_4_lat_dsl = np.where( - mask_case3a, - np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), - dreg_patch0_4_lat_dsl, - ) - - return ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) - - @staticmethod - def _apply_case3a_patch1( - mask_case3a, - lvn_sys_pos, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - pi1_x, - pi1_y, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - ): - dreg_patch1_1_lon_vmask = np.where( - mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask - ) - dreg_patch1_1_lat_vmask = np.where( - mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask - ) - dreg_patch1_2_lon_vmask = np.where( - mask_case3a, - np.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), - dreg_patch1_2_lon_vmask, - ) - dreg_patch1_2_lat_vmask = np.where( - mask_case3a, - np.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), - dreg_patch1_2_lat_vmask, - ) - dreg_patch1_3_lon_vmask = np.where( - mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask - ) - dreg_patch1_3_lat_vmask = np.where( - mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask - ) - dreg_patch1_4_lon_vmask = np.where( - mask_case3a, - np.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), - dreg_patch1_4_lon_vmask, - ) - dreg_patch1_4_lat_vmask = np.where( - mask_case3a, - np.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), - dreg_patch1_4_lat_vmask, - ) - - return ( - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - ) - - @staticmethod - def _apply_case3b_patch0( - mask_case3b, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - pi2_x, - pi2_y, - lvn_sys_pos, - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - ): - dreg_patch0_1_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_4_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) - dreg_patch0_4_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) - dreg_patch0_2_lon_dsl = np.where( - mask_case3b, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), - dreg_patch0_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case3b, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), - dreg_patch0_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where( - mask_case3b, - np.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), - dreg_patch0_3_lon_dsl, - ) - dreg_patch0_3_lat_dsl = np.where( - mask_case3b, - np.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), - dreg_patch0_3_lat_dsl, - ) - - return ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - ) - - @staticmethod - def _apply_case3b_patch2( - mask_case3b, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - pi2_x, - pi2_y, - lvn_sys_pos, - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - ): - dreg_patch2_1_lon_vmask = np.where( - mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask - ) - dreg_patch2_1_lat_vmask = np.where( - mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask - ) - dreg_patch2_2_lon_vmask = np.where( - mask_case3b, - np.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), - dreg_patch2_2_lon_vmask, - ) - dreg_patch2_2_lat_vmask = np.where( - mask_case3b, - np.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), - dreg_patch2_2_lat_vmask, - ) - dreg_patch2_3_lon_vmask = np.where( - mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask - ) - dreg_patch2_3_lat_vmask = np.where( - mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask - ) - dreg_patch2_4_lon_vmask = np.where( - mask_case3b, - np.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), - dreg_patch2_4_lon_vmask, - ) - dreg_patch2_4_lat_vmask = np.where( - mask_case3b, - np.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), - dreg_patch2_4_lat_vmask, - ) - - return ( - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - ) - - @classmethod - def reference( - cls, - grid, - famask_int, - p_vn, - ptr_v3_lon, - ptr_v3_lat, - tangent_orientation_dsl, - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - **kwargs, - ): - e2c = grid.connectivities[E2CDim] - ptr_v3_lon = reshape(ptr_v3_lon, e2c.shape) - ptr_v3_lon_e = np.expand_dims(ptr_v3_lon, axis=-1) - ptr_v3_lat = reshape(ptr_v3_lat, e2c.shape) - ptr_v3_lat_e = np.expand_dims(ptr_v3_lat, axis=-1) - tangent_orientation_dsl = np.expand_dims(tangent_orientation_dsl, axis=-1) - - result_tuple = cls._generate_flux_area_geometry( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - p_vn, - ptr_v3_lon_e, - ptr_v3_lat_e, - ) - - ( - fl_line_p1_lon, - fl_line_p1_lat, - fl_line_p2_lon, - fl_line_p2_lat, - tri_line1_p1_lon, - tri_line1_p1_lat, - tri_line1_p2_lon, - tri_line1_p2_lat, - tri_line2_p1_lon, - tri_line2_p1_lat, - tri_line2_p2_lon, - tri_line2_p2_lat, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - ) = result_tuple - - # Create first mask does departure-line segment intersects with A1V3 - lintersect_line1 = _lintersect_numpy( - fl_line_p1_lon, - fl_line_p1_lat, - fl_line_p2_lon, - fl_line_p2_lat, - tri_line1_p1_lon, - tri_line1_p1_lat, - tri_line1_p2_lon, - tri_line1_p2_lat, - ) - # Create first mask does departure-line segment intersects with A2V3 - lintersect_line2 = _lintersect_numpy( - fl_line_p1_lon, - fl_line_p1_lat, - fl_line_p2_lon, - fl_line_p2_lat, - tri_line2_p1_lon, - tri_line2_p1_lat, - tri_line2_p2_lon, - tri_line2_p2_lat, - ) - - lvn_sys_pos = np.where( - (p_vn * np.broadcast_to(tangent_orientation_dsl, p_vn.shape)) >= 0.0, - True, - False, - ) - famask_bool = np.where(famask_int == int32(1), True, False) - mask_case1 = np.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) - ps1_x, ps1_y = _line_intersect_numpy( - fl_line_p1_lon, - fl_line_p1_lat, - fl_line_p2_lon, - fl_line_p2_lat, - tri_line1_p1_lon, - tri_line1_p1_lat, - tri_line1_p2_lon, - tri_line1_p2_lat, - ) - ps2_x, ps2_y = _line_intersect_numpy( - fl_line_p1_lon, - fl_line_p1_lat, - fl_line_p2_lon, - fl_line_p2_lat, - tri_line2_p1_lon, - tri_line2_p1_lat, - tri_line2_p2_lon, - tri_line2_p2_lat, - ) - - # ------------------------------------------------- Case 1 - # Case 1 - patch 0 - ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) = cls._apply_case1_patch0( - mask_case1, - lvn_sys_pos, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - ps1_x, - ps1_y, - ps2_x, - ps2_y, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - ) - # Case 1 - patch 1 - ( - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - ) = cls._apply_case1_patch1( - mask_case1, - lvn_sys_pos, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - ps1_x, - ps1_y, - ) - # Case 1 - patch 2 - ( - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - ) = cls._apply_case1_patch2( - mask_case1, - lvn_sys_pos, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - ps2_x, - ps2_y, - ) - - # ------------------------------------------------- Case 2a - mask_case2a = np.logical_and.reduce( - [lintersect_line1, np.logical_not(lintersect_line2), famask_bool] - ) - # Case 2a - patch 0 - result_tuple_patch0 = cls._apply_case2a_patch0( - mask_case2a, - lvn_sys_pos, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - ps1_x, - ps1_y, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) - - ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) = result_tuple_patch0 - # Case 2a - patch 1 - result_tuple_patch1 = cls._apply_case2a_patch1( - mask_case2a, - lvn_sys_pos, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - ps1_x, - ps1_y, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - ) - - ( - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - ) = result_tuple_patch1 - # Case 2a - patch 2 - dreg_patch2_1_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lat_vmask) - dreg_patch2_2_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lon_vmask) - dreg_patch2_2_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lat_vmask) - dreg_patch2_3_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lat_vmask) - dreg_patch2_4_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lon_vmask) - dreg_patch2_4_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lat_vmask) - - # -------------------------------------------------- Case 2b - mask_case2b = np.logical_and.reduce( - [lintersect_line2, np.logical_not(lintersect_line1), famask_bool] - ) - # Case 2b - patch 0 - result_tuple_patch0_case2b = cls._apply_case2b_patch0( - mask_case2b, - lvn_sys_pos, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - ps2_x, - ps2_y, - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) - - ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) = result_tuple_patch0_case2b - - # Case 2b - patch 1 - result_tuple_patch1_case2b = cls._apply_case2b_patch1( - mask_case2b, - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - ) - - ( - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - ) = result_tuple_patch1_case2b - - # Case 2b - patch 2 - result_tuple_patch2_case2b = cls._apply_case2b_patch2( - mask_case2b, - lvn_sys_pos, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - ps2_x, - ps2_y, - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - ) - - ( - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - ) = result_tuple_patch2_case2b - - # flux area edge 1 and 2 - fl_e1_p1_lon = arrival_pts_1_lon_dsl - fl_e1_p1_lat = arrival_pts_1_lat_dsl - fl_e1_p2_lon = depart_pts_1_lon_dsl - fl_e1_p2_lat = depart_pts_1_lat_dsl - fl_e2_p1_lon = arrival_pts_2_lon_dsl - fl_e2_p1_lat = arrival_pts_2_lat_dsl - fl_e2_p2_lon = depart_pts_2_lon_dsl - fl_e2_p2_lat = depart_pts_2_lat_dsl - - # ----------------------------------------------- Case 3a - # Check whether flux area edge 2 intersects with triangle edge 1 - lintersect_e2_line1 = _lintersect_numpy( - fl_e2_p1_lon, - fl_e2_p1_lat, - fl_e2_p2_lon, - fl_e2_p2_lat, - tri_line1_p1_lon, - tri_line1_p1_lat, - tri_line1_p2_lon, - tri_line1_p2_lat, - ) - mask_case3a = np.logical_and(lintersect_e2_line1, famask_bool) - pi1_x, pi1_y = _line_intersect_numpy( - fl_e2_p1_lon, - fl_e2_p1_lat, - fl_e2_p2_lon, - fl_e2_p2_lat, - tri_line1_p1_lon, - tri_line1_p1_lat, - tri_line1_p2_lon, - tri_line1_p2_lat, - ) - # Case 3a - patch 0 - result = cls._apply_case3a_patch0( - mask_case3a, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - lvn_sys_pos, - ps2_x, - ps2_y, - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) - - ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) = result - - # Case 3a - patch 1 - ( - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - ) = cls._apply_case3a_patch1( - mask_case3a, - lvn_sys_pos, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - pi1_x, - pi1_y, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - ) - # Case 3a - patch 2 - dreg_patch2_1_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lat_vmask) - dreg_patch2_2_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lon_vmask) - dreg_patch2_2_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lat_vmask) - dreg_patch2_3_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lat_vmask) - dreg_patch2_4_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lon_vmask) - dreg_patch2_4_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lat_vmask) - - # ------------------------------------------------ Case 3b - # Check whether flux area edge 1 intersects with triangle edge 2 - lintersect_e1_line2 = _lintersect_numpy( - fl_e1_p1_lon, - fl_e1_p1_lat, - fl_e1_p2_lon, - fl_e1_p2_lat, - tri_line2_p1_lon, - tri_line2_p1_lat, - tri_line2_p2_lon, - tri_line2_p2_lat, - ) - mask_case3b = lintersect_e1_line2 & famask_bool - pi2_x, pi2_y = _line_intersect_numpy( - fl_e1_p1_lon, - fl_e1_p1_lat, - fl_e1_p2_lon, - fl_e1_p2_lat, - tri_line2_p1_lon, - tri_line2_p1_lat, - tri_line2_p2_lon, - tri_line2_p2_lat, - ) - # Case 3b - patch 0 - ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - ) = cls._apply_case3b_patch0( - mask_case3b, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - pi2_x, - pi2_y, - lvn_sys_pos, - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - ) - # Case 3b - patch 1 - dreg_patch1_1_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lat_vmask) - dreg_patch1_2_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lon_vmask) - dreg_patch1_2_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lat_vmask) - dreg_patch1_3_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lat_vmask) - dreg_patch1_4_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) - - # Case 3b - patch 2 - ( - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - ) = cls._apply_case3b_patch2( - mask_case3b, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - pi2_x, - pi2_y, - lvn_sys_pos, - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - ) - # --------------------------------------------- Case 4 - # NB: Next line acts as the "ELSE IF", indices that already previously matched one of the above conditions - # can't be overwritten by this new condition. - indices_previously_matched = np.logical_or.reduce( - [mask_case3b, mask_case3a, mask_case2b, mask_case2a, mask_case1] - ) - # mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices - mask_case4 = np.logical_and.reduce( - [famask_bool, np.logical_not(indices_previously_matched)] - ) - # Case 4 - patch 0 - no change - # Case 4 - patch 1 - dreg_patch1_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lat_vmask) - dreg_patch1_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lon_vmask) - dreg_patch1_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lat_vmask) - dreg_patch1_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lat_vmask) - dreg_patch1_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lat_vmask) - # Case 4 - patch 2 - dreg_patch2_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lat_vmask) - dreg_patch2_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lon_vmask) - dreg_patch2_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lat_vmask) - dreg_patch2_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lat_vmask) - dreg_patch2_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_4_lon_vmask) - - return dict( - dreg_patch0_1_lon_dsl=dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl=dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl=dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl=dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl=dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl=dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl=dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl=dreg_patch0_4_lat_dsl, - dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, - dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, - dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, - dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, - dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, - ) - - @pytest.fixture - def input_data(self, grid): - famask_int = random_mask(grid, EdgeDim, KDim, dtype=int32) - p_vn = random_field(grid, EdgeDim, KDim) - ptr_v3_lon = random_field(grid, EdgeDim, E2CDim, low=0.1, high=1.0) - ptr_v3_lon_field = as_1D_sparse_field(ptr_v3_lon, ECDim) - ptr_v3_lat = random_field(grid, EdgeDim, E2CDim, low=0.1, high=1.0) - ptr_v3_lat_field = as_1D_sparse_field(ptr_v3_lat, ECDim) - tangent_orientation_dsl = random_field(grid, EdgeDim, low=0.1, high=1.0) - dreg_patch0_1_lon_dsl = constant_field(grid, 1.0, EdgeDim, KDim) - dreg_patch0_1_lat_dsl = constant_field(grid, 1.0, EdgeDim, KDim) - dreg_patch0_2_lon_dsl = constant_field(grid, 2.0, EdgeDim, KDim) - dreg_patch0_2_lat_dsl = constant_field(grid, 2.0, EdgeDim, KDim) - dreg_patch0_3_lon_dsl = constant_field(grid, 3.0, EdgeDim, KDim) - dreg_patch0_3_lat_dsl = constant_field(grid, 3.0, EdgeDim, KDim) - dreg_patch0_4_lon_dsl = constant_field(grid, 4.0, EdgeDim, KDim) - dreg_patch0_4_lat_dsl = constant_field(grid, 4.0, EdgeDim, KDim) - dreg_patch1_1_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_1_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_2_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_2_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_3_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_3_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_4_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_4_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_1_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_1_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_2_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_2_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_3_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_3_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_4_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_4_lat_vmask = zero_field(grid, EdgeDim, KDim) - return dict( - famask_int=famask_int, - p_vn=p_vn, - ptr_v3_lon=ptr_v3_lon_field, - ptr_v3_lat=ptr_v3_lat_field, - tangent_orientation_dsl=tangent_orientation_dsl, - dreg_patch0_1_lon_dsl=dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl=dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl=dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl=dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl=dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl=dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl=dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl=dreg_patch0_4_lat_dsl, - dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, - dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, - dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, - dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, - dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_02.py deleted file mode 100644 index 6b1daabc13..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_divide_flux_area_list_stencil_02.py +++ /dev/null @@ -1,290 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.divide_flux_area_list_stencil_02 import ( - divide_flux_area_list_stencil_02, -) -from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - as_1D_sparse_field, - random_field, - random_mask, - reshape, -) - - -class TestDivideFluxAreaListStencil02(StencilTest): - PROGRAM = divide_flux_area_list_stencil_02 - OUTPUTS = ( - "dreg_patch1_1_lon_vmask", - "dreg_patch1_1_lat_vmask", - "dreg_patch1_2_lon_vmask", - "dreg_patch1_2_lat_vmask", - "dreg_patch1_3_lon_vmask", - "dreg_patch1_3_lat_vmask", - "dreg_patch1_4_lon_vmask", - "dreg_patch1_4_lat_vmask", - "dreg_patch2_1_lon_vmask", - "dreg_patch2_1_lat_vmask", - "dreg_patch2_2_lon_vmask", - "dreg_patch2_2_lat_vmask", - "dreg_patch2_3_lon_vmask", - "dreg_patch2_3_lat_vmask", - "dreg_patch2_4_lon_vmask", - "dreg_patch2_4_lat_vmask", - "patch1_cell_idx_vmask", - "patch1_cell_blk_vmask", - "patch2_cell_idx_vmask", - "patch2_cell_blk_vmask", - ) - - @staticmethod - def reference( - grid, - famask_int: np.array, - p_vn: np.array, - bf_cc_patch1_lon: np.array, - bf_cc_patch1_lat: np.array, - bf_cc_patch2_lon: np.array, - bf_cc_patch2_lat: np.array, - butterfly_idx_patch1_vnpos: np.array, - butterfly_idx_patch1_vnneg: np.array, - butterfly_blk_patch1_vnpos: np.array, - butterfly_blk_patch1_vnneg: np.array, - butterfly_idx_patch2_vnpos: np.array, - butterfly_idx_patch2_vnneg: np.array, - butterfly_blk_patch2_vnpos: np.array, - butterfly_blk_patch2_vnneg: np.array, - dreg_patch1_1_lon_vmask: np.array, - dreg_patch1_1_lat_vmask: np.array, - dreg_patch1_2_lon_vmask: np.array, - dreg_patch1_2_lat_vmask: np.array, - dreg_patch1_3_lon_vmask: np.array, - dreg_patch1_3_lat_vmask: np.array, - dreg_patch1_4_lon_vmask: np.array, - dreg_patch1_4_lat_vmask: np.array, - dreg_patch2_1_lon_vmask: np.array, - dreg_patch2_1_lat_vmask: np.array, - dreg_patch2_2_lon_vmask: np.array, - dreg_patch2_2_lat_vmask: np.array, - dreg_patch2_3_lon_vmask: np.array, - dreg_patch2_3_lat_vmask: np.array, - dreg_patch2_4_lon_vmask: np.array, - dreg_patch2_4_lat_vmask: np.array, - **kwargs, - ): - e2c = grid.connectivities[E2CDim] - famask_bool = np.where(famask_int == int32(1), True, False) - lvn_pos = np.where(p_vn >= np.broadcast_to(0.0, p_vn.shape), True, False) - # Translation of patch 1 and patch 2 in system relative to respective cell - bf_cc_patch1_lon = reshape(bf_cc_patch1_lon, e2c.shape) - bf_cc_patch1_lon_e = np.expand_dims(bf_cc_patch1_lon, axis=-1) - bf_cc_patch1_lat = reshape(bf_cc_patch1_lat, e2c.shape) - bf_cc_patch1_lat_e = np.expand_dims(bf_cc_patch1_lat, axis=-1) - bf_cc_patch2_lon = reshape(bf_cc_patch2_lon, e2c.shape) - bf_cc_patch2_lon_e = np.expand_dims(bf_cc_patch2_lon, axis=-1) - bf_cc_patch2_lat = reshape(bf_cc_patch2_lat, e2c.shape) - bf_cc_patch2_lat_e = np.expand_dims(bf_cc_patch2_lat, axis=-1) - - bf_cc_patch1_lon = np.where( - famask_bool, - np.where(lvn_pos, bf_cc_patch1_lon_e[:, 0], bf_cc_patch1_lon_e[:, 1]), - 0.0, - ) - bf_cc_patch1_lat = np.where( - famask_bool, - np.where(lvn_pos, bf_cc_patch1_lat_e[:, 0], bf_cc_patch1_lat_e[:, 1]), - 0.0, - ) - bf_cc_patch2_lon = np.where( - famask_bool, - np.where(lvn_pos, bf_cc_patch2_lon_e[:, 0], bf_cc_patch2_lon_e[:, 1]), - 0.0, - ) - bf_cc_patch2_lat = np.where( - famask_bool, - np.where(lvn_pos, bf_cc_patch2_lat_e[:, 0], bf_cc_patch2_lat_e[:, 1]), - 0.0, - ) - - # patch1 in translated system - dreg_patch1_1_lon_vmask = dreg_patch1_1_lon_vmask - bf_cc_patch1_lon - dreg_patch1_1_lat_vmask = dreg_patch1_1_lat_vmask - bf_cc_patch1_lat - dreg_patch1_2_lon_vmask = dreg_patch1_2_lon_vmask - bf_cc_patch1_lon - dreg_patch1_2_lat_vmask = dreg_patch1_2_lat_vmask - bf_cc_patch1_lat - dreg_patch1_3_lon_vmask = dreg_patch1_3_lon_vmask - bf_cc_patch1_lon - dreg_patch1_3_lat_vmask = dreg_patch1_3_lat_vmask - bf_cc_patch1_lat - dreg_patch1_4_lon_vmask = dreg_patch1_4_lon_vmask - bf_cc_patch1_lon - dreg_patch1_4_lat_vmask = dreg_patch1_4_lat_vmask - bf_cc_patch1_lat - # patch2 in translated system - dreg_patch2_1_lon_vmask = dreg_patch2_1_lon_vmask - bf_cc_patch2_lon - dreg_patch2_1_lat_vmask = dreg_patch2_1_lat_vmask - bf_cc_patch2_lat - dreg_patch2_2_lon_vmask = dreg_patch2_2_lon_vmask - bf_cc_patch2_lon - dreg_patch2_2_lat_vmask = dreg_patch2_2_lat_vmask - bf_cc_patch2_lat - dreg_patch2_3_lon_vmask = dreg_patch2_3_lon_vmask - bf_cc_patch2_lon - dreg_patch2_3_lat_vmask = dreg_patch2_3_lat_vmask - bf_cc_patch2_lat - dreg_patch2_4_lon_vmask = dreg_patch2_4_lon_vmask - bf_cc_patch2_lon - dreg_patch2_4_lat_vmask = dreg_patch2_4_lat_vmask - bf_cc_patch2_lat - - # Store global index of the underlying grid cell - # Adapt dimensions to fit ofr multiple levels - butterfly_idx_patch1_vnpos_3d = np.broadcast_to( - np.expand_dims(butterfly_idx_patch1_vnpos, axis=-1), p_vn.shape - ) - butterfly_idx_patch1_vnneg_3d = np.broadcast_to( - np.expand_dims(butterfly_idx_patch1_vnneg, axis=-1), p_vn.shape - ) - butterfly_idx_patch2_vnpos_3d = np.broadcast_to( - np.expand_dims(butterfly_idx_patch2_vnpos, axis=-1), p_vn.shape - ) - butterfly_idx_patch2_vnneg_3d = np.broadcast_to( - np.expand_dims(butterfly_idx_patch2_vnneg, axis=-1), p_vn.shape - ) - butterfly_blk_patch1_vnpos_3d = np.broadcast_to( - np.expand_dims(butterfly_blk_patch1_vnpos, axis=-1), p_vn.shape - ) - butterfly_blk_patch1_vnneg_3d = np.broadcast_to( - np.expand_dims(butterfly_blk_patch1_vnneg, axis=-1), p_vn.shape - ) - butterfly_blk_patch2_vnpos_3d = np.broadcast_to( - np.expand_dims(butterfly_blk_patch2_vnpos, axis=-1), p_vn.shape - ) - butterfly_blk_patch2_vnneg_3d = np.broadcast_to( - np.expand_dims(butterfly_blk_patch2_vnneg, axis=-1), p_vn.shape - ) - patch1_cell_idx_vmask = np.where( - famask_bool, - np.where(lvn_pos, butterfly_idx_patch1_vnpos_3d, butterfly_idx_patch1_vnneg_3d), - int32(0), - ) - patch2_cell_idx_vmask = np.where( - famask_bool, - np.where(lvn_pos, butterfly_idx_patch2_vnpos_3d, butterfly_idx_patch2_vnneg_3d), - int32(0), - ) - patch1_cell_blk_vmask = np.where( - famask_bool, - np.where(lvn_pos, butterfly_blk_patch1_vnpos_3d, butterfly_blk_patch1_vnneg_3d), - int32(0), - ) - patch2_cell_blk_vmask = np.where( - famask_bool, - np.where(lvn_pos, butterfly_blk_patch2_vnpos_3d, butterfly_blk_patch2_vnneg_3d), - int32(0), - ) - - return dict( - dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, - dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, - dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, - dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, - dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, - patch1_cell_idx_vmask=patch1_cell_idx_vmask, - patch1_cell_blk_vmask=patch1_cell_blk_vmask, - patch2_cell_idx_vmask=patch2_cell_idx_vmask, - patch2_cell_blk_vmask=patch2_cell_blk_vmask, - ) - - @pytest.fixture - def input_data(self, grid): - famask_int = random_mask(grid, EdgeDim, KDim, dtype=int32) - p_vn = random_field(grid, EdgeDim, KDim) - bf_cc_patch1_lon = random_field(grid, EdgeDim, E2CDim) - bf_cc_patch1_lon_field = as_1D_sparse_field(bf_cc_patch1_lon, ECDim) - bf_cc_patch1_lat = random_field(grid, EdgeDim, E2CDim) - bf_cc_patch1_lat_field = as_1D_sparse_field(bf_cc_patch1_lat, ECDim) - bf_cc_patch2_lon = random_field(grid, EdgeDim, E2CDim) - bf_cc_patch2_lon_field = as_1D_sparse_field(bf_cc_patch2_lon, ECDim) - bf_cc_patch2_lat = random_field(grid, EdgeDim, E2CDim) - bf_cc_patch2_lat_field = as_1D_sparse_field(bf_cc_patch2_lat, ECDim) - butterfly_idx_patch1_vnpos = random_mask(grid, EdgeDim, dtype=int32) - butterfly_idx_patch1_vnneg = random_mask(grid, EdgeDim, dtype=int32) - butterfly_blk_patch1_vnpos = random_mask(grid, EdgeDim, dtype=int32) - butterfly_blk_patch1_vnneg = random_mask(grid, EdgeDim, dtype=int32) - butterfly_idx_patch2_vnpos = random_mask(grid, EdgeDim, dtype=int32) - butterfly_idx_patch2_vnneg = random_mask(grid, EdgeDim, dtype=int32) - butterfly_blk_patch2_vnpos = random_mask(grid, EdgeDim, dtype=int32) - butterfly_blk_patch2_vnneg = random_mask(grid, EdgeDim, dtype=int32) - dreg_patch1_1_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_1_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_2_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_2_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_3_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_3_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_4_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_4_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_1_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_1_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_2_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_2_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_3_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_3_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_4_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_4_lat_vmask = random_field(grid, EdgeDim, KDim) - patch1_cell_idx_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) - patch1_cell_blk_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) - patch2_cell_idx_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) - patch2_cell_blk_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) - - return dict( - famask_int=famask_int, - p_vn=p_vn, - bf_cc_patch1_lon=bf_cc_patch1_lon_field, - bf_cc_patch1_lat=bf_cc_patch1_lat_field, - bf_cc_patch2_lon=bf_cc_patch2_lon_field, - bf_cc_patch2_lat=bf_cc_patch2_lat_field, - butterfly_idx_patch1_vnpos=butterfly_idx_patch1_vnpos, - butterfly_idx_patch1_vnneg=butterfly_idx_patch1_vnneg, - butterfly_blk_patch1_vnpos=butterfly_blk_patch1_vnpos, - butterfly_blk_patch1_vnneg=butterfly_blk_patch1_vnneg, - butterfly_idx_patch2_vnpos=butterfly_idx_patch2_vnpos, - butterfly_idx_patch2_vnneg=butterfly_idx_patch2_vnneg, - butterfly_blk_patch2_vnpos=butterfly_blk_patch2_vnpos, - butterfly_blk_patch2_vnneg=butterfly_blk_patch2_vnneg, - dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, - dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, - dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, - dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, - dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, - patch1_cell_idx_vmask=patch1_cell_idx_vmask, - patch1_cell_blk_vmask=patch1_cell_blk_vmask, - patch2_cell_idx_vmask=patch2_cell_idx_vmask, - patch2_cell_blk_vmask=patch2_cell_blk_vmask, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_01.py deleted file mode 100644 index 222a53e035..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_01.py +++ /dev/null @@ -1,78 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np -import pytest -from gt4py.next import as_field -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - Output, - StencilTest, - _shape, - random_field, - zero_field, -) - - -class TestFaceValPpmStencil01(StencilTest): - PROGRAM = face_val_ppm_stencil_01 - OUTPUTS = ( - Output( - "z_slope", refslice=(slice(None), slice(None, -1)), gtslice=(slice(None), slice(1, -1)) - ), - ) - - @staticmethod - def reference( - grid, p_cc: np.array, p_cellhgt_mc_now: np.array, k: np.array, elev: int32, **kwargs - ): - zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( - p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] - ) - zfac = (p_cc[:, 2:] - p_cc[:, 1:-1]) / (p_cellhgt_mc_now[:, 2:] + p_cellhgt_mc_now[:, 1:-1]) - z_slope_a = ( - p_cellhgt_mc_now[:, 1:-1] - / (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 2:]) - ) * ( - (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac - + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 2:]) * zfac_m1 - ) - - zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( - p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] - ) - zfac = (p_cc[:, 1:-1] - p_cc[:, 1:-1]) / ( - p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1] - ) - z_slope_b = ( - p_cellhgt_mc_now[:, 1:-1] - / (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1]) - ) * ( - (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac - + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 1:-1]) * zfac_m1 - ) - - z_slope = np.where(k[1:-1] < elev, z_slope_a, z_slope_b) - return dict(z_slope=z_slope) - - @pytest.fixture - def input_data(self, grid): - z_slope = zero_field(grid, CellDim, KDim) - p_cc = random_field(grid, CellDim, KDim, extend={KDim: 1}) - p_cellhgt_mc_now = random_field(grid, CellDim, KDim, extend={KDim: 1}) - k = as_field((KDim,), np.arange(0, _shape(grid, KDim, extend={KDim: 1})[0], dtype=int32)) - elev = k[-2].as_scalar() - - return dict(p_cc=p_cc, p_cellhgt_mc_now=p_cellhgt_mc_now, k=k, elev=elev, z_slope=z_slope) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02.py deleted file mode 100644 index b9f27fb7fe..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02.py +++ /dev/null @@ -1,76 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next import as_field -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02 import face_val_ppm_stencil_02 -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, _shape, random_field, zero_field - - -class TestFaceValPpmStencil02(StencilTest): - PROGRAM = face_val_ppm_stencil_02 - OUTPUTS = ("p_face",) - - @staticmethod - def reference( - grid, - p_cc: np.array, - p_cellhgt_mc_now: np.array, - p_face_in: np.array, - k: np.array, - slev: int32, - elev: int32, - slevp1: int32, - elevp1: int32, - **kwargs, - ): - p_face_a = p_face_in - p_face_a[:, 1:] = p_cc[:, 1:] * ( - 1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) - ) + (p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:])) * ( - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1] - ) - - p_face = np.where((k == slevp1) | (k == elev), p_face_a, p_face_in) - p_face = np.where((k == slev), p_cc, p_face) - p_face[:, 1:] = np.where((k[1:] == elevp1), p_cc[:, :-1], p_face[:, 1:]) - return dict(p_face=p_face) - - @pytest.fixture - def input_data(self, grid): - p_cc = random_field(grid, CellDim, KDim) - p_cellhgt_mc_now = random_field(grid, CellDim, KDim) - p_face_in = random_field(grid, CellDim, KDim) - p_face = zero_field(grid, CellDim, KDim) - - k = as_field((KDim,), np.arange(0, _shape(grid, KDim)[0], dtype=int32)) - slev = int32(1) - slevp1 = slev + int32(1) - elev = int32(k[-3].as_scalar()) - elevp1 = elev + int32(1) - - return dict( - p_cc=p_cc, - p_cellhgt_mc_now=p_cellhgt_mc_now, - p_face_in=p_face_in, - k=k, - slev=slev, - elev=elev, - slevp1=slevp1, - elevp1=elevp1, - p_face=p_face, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02a.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02a.py deleted file mode 100644 index d1d8d5d900..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02a.py +++ /dev/null @@ -1,43 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02a import face_val_ppm_stencil_02a -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field - - -outslice = (slice(None), slice(1, None)) - - -class TestFaceValPpmStencil02a(StencilTest): - PROGRAM = face_val_ppm_stencil_02a - OUTPUTS = (Output("p_face", refslice=outslice, gtslice=outslice),) - - @staticmethod - def reference(grid, p_cc: np.array, p_cellhgt_mc_now: np.array, **kwargs): - p_face = p_cc.copy() - p_face[:, 1:] = p_cc[:, 1:] * ( - 1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) - ) + (p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:])) * ( - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1] - ) - return dict(p_face=p_face) - - @pytest.fixture - def input_data(self, grid): - p_face = random_field(grid, CellDim, KDim) - p_cc = random_field(grid, CellDim, KDim) - p_cellhgt_mc_now = random_field(grid, CellDim, KDim) - return dict(p_cc=p_cc, p_cellhgt_mc_now=p_cellhgt_mc_now, p_face=p_face) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02b.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02b.py deleted file mode 100644 index 1a0a540106..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02b.py +++ /dev/null @@ -1,35 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02b import face_val_ppm_stencil_02b -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field - - -class TestFaceValPpmStencil02b(StencilTest): - PROGRAM = face_val_ppm_stencil_02b - OUTPUTS = ("p_face",) - - @staticmethod - def reference(grid, p_cc: np.array, **kwargs): - p_face = p_cc.copy() - return dict(p_face=p_face) - - @pytest.fixture - def input_data(self, grid): - p_cc = random_field(grid, CellDim, KDim) - p_face = random_field(grid, CellDim, KDim) - return dict(p_cc=p_cc, p_face=p_face) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02c.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02c.py deleted file mode 100644 index bdc5529eaa..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_02c.py +++ /dev/null @@ -1,52 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02c import face_val_ppm_stencil_02c -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field - - -class TestFaceValPpmStencil02c(StencilTest): - PROGRAM = face_val_ppm_stencil_02c - OUTPUTS = ("p_face",) - - @staticmethod - def reference( - grid, - p_cc: np.array, - p_face: np.array, - horizontal_start: int32, - horizontal_end: int32, - vertical_start: int32, - vertical_end: int32, - ): - subset = (slice(horizontal_start, horizontal_end), slice(vertical_start, vertical_end)) - p_face = p_face.copy() - p_face[subset] = np.roll(p_cc, shift=1, axis=1)[subset] - return dict(p_face=p_face) - - @pytest.fixture - def input_data(self, grid): - p_cc = random_field(grid, CellDim, KDim) - p_face = random_field(grid, CellDim, KDim) - return dict( - p_cc=p_cc, - p_face=p_face, - horizontal_start=int32(0), - horizontal_end=int32(grid.num_cells), - vertical_start=int32(1), - vertical_end=int32(grid.num_levels), - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_05.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_05.py deleted file mode 100644 index 6d7115b942..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_face_val_ppm_stencil_05.py +++ /dev/null @@ -1,71 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.face_val_ppm_stencil_05 import face_val_ppm_stencil_05 -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field, zero_field - - -class TestFaceValPpmStencil05(StencilTest): - PROGRAM = face_val_ppm_stencil_05 - OUTPUTS = (Output("p_face", gtslice=(slice(None), slice(2, None))),) - - @staticmethod - def reference(grid, p_cc: np.array, p_cellhgt_mc_now: np.array, z_slope: np.array, **kwargs): - p_cellhgt_mc_now_k_minus_1 = p_cellhgt_mc_now[:, 1:-2] - p_cellhgt_mc_now_k_minus_2 = p_cellhgt_mc_now[:, 0:-3] - p_cellhgt_mc_now_k_plus_1 = p_cellhgt_mc_now[:, 3:] - p_cellhgt_mc_now = p_cellhgt_mc_now[:, 2:-1] - - p_cc_k_minus_1 = p_cc[:, 1:-1] - p_cc = p_cc[:, 2:] - z_slope_k_minus_1 = z_slope[:, 1:-1] - z_slope = z_slope[:, 2:] - - zgeo1 = p_cellhgt_mc_now_k_minus_1 / (p_cellhgt_mc_now_k_minus_1 + p_cellhgt_mc_now) - zgeo2 = 1.0 / ( - p_cellhgt_mc_now_k_minus_2 - + p_cellhgt_mc_now_k_minus_1 - + p_cellhgt_mc_now - + p_cellhgt_mc_now_k_plus_1 - ) - zgeo3 = (p_cellhgt_mc_now_k_minus_2 + p_cellhgt_mc_now_k_minus_1) / ( - 2.0 * p_cellhgt_mc_now_k_minus_1 + p_cellhgt_mc_now - ) - zgeo4 = (p_cellhgt_mc_now_k_plus_1 + p_cellhgt_mc_now) / ( - 2 * p_cellhgt_mc_now + p_cellhgt_mc_now_k_minus_1 - ) - - p_face = ( - p_cc_k_minus_1 - + zgeo1 * (p_cc - p_cc_k_minus_1) - + zgeo2 - * ( - (2 * p_cellhgt_mc_now * zgeo1) * (zgeo3 - zgeo4) * (p_cc - p_cc_k_minus_1) - - zgeo3 * p_cellhgt_mc_now_k_minus_1 * z_slope - + zgeo4 * p_cellhgt_mc_now * z_slope_k_minus_1 - ) - ) - return dict(p_face=p_face) - - @pytest.fixture - def input_data(self, grid): - p_cc = random_field(grid, CellDim, KDim) - p_cellhgt_mc_now = random_field(grid, CellDim, KDim, extend={KDim: 1}) - z_slope = random_field(grid, CellDim, KDim) - p_face = zero_field(grid, CellDim, KDim) - - return dict(p_cc=p_cc, p_cellhgt_mc_now=p_cellhgt_mc_now, z_slope=z_slope, p_face=p_face) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py deleted file mode 100644 index acfa579ff3..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py +++ /dev/null @@ -1,184 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_01a import ( - hflux_ffsl_hybrid_stencil_01a, -) -from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - constant_field, - random_field, - zero_field, -) - - -class TestHfluxFfslHybridStencil01a(StencilTest): - PROGRAM = hflux_ffsl_hybrid_stencil_01a - OUTPUTS = ("p_out_e_hybrid_1a",) - - @staticmethod - def reference( - grid, - z_lsq_coeff_1: np.array, - z_lsq_coeff_2: np.array, - z_lsq_coeff_3: np.array, - z_lsq_coeff_4: np.array, - z_lsq_coeff_5: np.array, - z_lsq_coeff_6: np.array, - z_lsq_coeff_7: np.array, - z_lsq_coeff_8: np.array, - z_lsq_coeff_9: np.array, - z_lsq_coeff_10: np.array, - z_quad_vector_sum0_1: np.array, - z_quad_vector_sum0_2: np.array, - z_quad_vector_sum0_3: np.array, - z_quad_vector_sum0_4: np.array, - z_quad_vector_sum0_5: np.array, - z_quad_vector_sum0_6: np.array, - z_quad_vector_sum0_7: np.array, - z_quad_vector_sum0_8: np.array, - z_quad_vector_sum0_9: np.array, - z_quad_vector_sum0_10: np.array, - patch0_cell_rel_idx_dsl: np.array, - **kwargs, - ): - e2c = grid.connectivities[E2CDim] - z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] - z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] - z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] - z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] - z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] - z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] - z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] - z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] - z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] - z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] - - p_out_e_hybrid_1a = ( - np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_1_e2c[:, 1], - z_lsq_coeff_1_e2c[:, 0], - ) - * z_quad_vector_sum0_1 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_2_e2c[:, 1], - z_lsq_coeff_2_e2c[:, 0], - ) - * z_quad_vector_sum0_2 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_3_e2c[:, 1], - z_lsq_coeff_3_e2c[:, 0], - ) - * z_quad_vector_sum0_3 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_4_e2c[:, 1], - z_lsq_coeff_4_e2c[:, 0], - ) - * z_quad_vector_sum0_4 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_5_e2c[:, 1], - z_lsq_coeff_5_e2c[:, 0], - ) - * z_quad_vector_sum0_5 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_6_e2c[:, 1], - z_lsq_coeff_6_e2c[:, 0], - ) - * z_quad_vector_sum0_6 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_7_e2c[:, 1], - z_lsq_coeff_7_e2c[:, 0], - ) - * z_quad_vector_sum0_7 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_8_e2c[:, 1], - z_lsq_coeff_8_e2c[:, 0], - ) - * z_quad_vector_sum0_8 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_9_e2c[:, 1], - z_lsq_coeff_9_e2c[:, 0], - ) - * z_quad_vector_sum0_9 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_10_e2c[:, 1], - z_lsq_coeff_10_e2c[:, 0], - ) - * z_quad_vector_sum0_10 - ) - - return dict(p_out_e_hybrid_1a=p_out_e_hybrid_1a) - - @pytest.fixture - def input_data(self, grid): - z_lsq_coeff_1 = random_field(grid, CellDim, KDim) - z_lsq_coeff_2 = random_field(grid, CellDim, KDim) - z_lsq_coeff_3 = random_field(grid, CellDim, KDim) - z_lsq_coeff_4 = random_field(grid, CellDim, KDim) - z_lsq_coeff_5 = random_field(grid, CellDim, KDim) - z_lsq_coeff_6 = random_field(grid, CellDim, KDim) - z_lsq_coeff_7 = random_field(grid, CellDim, KDim) - z_lsq_coeff_8 = random_field(grid, CellDim, KDim) - z_lsq_coeff_9 = random_field(grid, CellDim, KDim) - z_lsq_coeff_10 = random_field(grid, CellDim, KDim) - z_quad_vector_sum0_1 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_2 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_3 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_4 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_5 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_6 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_7 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_8 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_9 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_10 = random_field(grid, EdgeDim, KDim) - patch0_cell_rel_idx_dsl = constant_field(grid, 1, EdgeDim, KDim, dtype=int32) - p_out_e_hybrid_1a = zero_field(grid, EdgeDim, KDim) - return dict( - z_lsq_coeff_1=z_lsq_coeff_1, - z_lsq_coeff_2=z_lsq_coeff_2, - z_lsq_coeff_3=z_lsq_coeff_3, - z_lsq_coeff_4=z_lsq_coeff_4, - z_lsq_coeff_5=z_lsq_coeff_5, - z_lsq_coeff_6=z_lsq_coeff_6, - z_lsq_coeff_7=z_lsq_coeff_7, - z_lsq_coeff_8=z_lsq_coeff_8, - z_lsq_coeff_9=z_lsq_coeff_9, - z_lsq_coeff_10=z_lsq_coeff_10, - z_quad_vector_sum0_1=z_quad_vector_sum0_1, - z_quad_vector_sum0_2=z_quad_vector_sum0_2, - z_quad_vector_sum0_3=z_quad_vector_sum0_3, - z_quad_vector_sum0_4=z_quad_vector_sum0_4, - z_quad_vector_sum0_5=z_quad_vector_sum0_5, - z_quad_vector_sum0_6=z_quad_vector_sum0_6, - z_quad_vector_sum0_7=z_quad_vector_sum0_7, - z_quad_vector_sum0_8=z_quad_vector_sum0_8, - z_quad_vector_sum0_9=z_quad_vector_sum0_9, - z_quad_vector_sum0_10=z_quad_vector_sum0_10, - patch0_cell_rel_idx_dsl=patch0_cell_rel_idx_dsl, - p_out_e_hybrid_1a=p_out_e_hybrid_1a, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py deleted file mode 100644 index 5c57bf79b9..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py +++ /dev/null @@ -1,49 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_02 import ( - hflux_ffsl_hybrid_stencil_02, -) -from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field - - -class TestHfluxFfslHybridStencil02(StencilTest): - PROGRAM = hflux_ffsl_hybrid_stencil_02 - OUTPUTS = ("p_out_e_hybrid_2",) - - @staticmethod - def reference( - grid, - p_out_e_hybrid_2: np.array, - p_mass_flx_e: np.array, - z_dreg_area: np.array, - **kwargs, - ): - p_out_e_hybrid_2 = p_mass_flx_e * p_out_e_hybrid_2 / z_dreg_area - - return dict(p_out_e_hybrid_2=p_out_e_hybrid_2) - - @pytest.fixture - def input_data(self, grid): - p_out_e_hybrid_2 = random_field(grid, EdgeDim, KDim) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - z_dreg_area = random_field(grid, EdgeDim, KDim) - return dict( - p_mass_flx_e=p_mass_flx_e, - z_dreg_area=z_dreg_area, - p_out_e_hybrid_2=p_out_e_hybrid_2, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01a.py deleted file mode 100644 index a885048535..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01a.py +++ /dev/null @@ -1,61 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01a import ( - hflx_limiter_mo_stencil_01a, -) -from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestHflxLimiterMoStencil01a(StencilTest): - PROGRAM = hflx_limiter_mo_stencil_01a - OUTPUTS = ( - "z_mflx_low", - "z_anti", - ) - - @staticmethod - def reference( - grid, p_mflx_tracer_h: np.ndarray, p_mass_flx_e: np.ndarray, p_cc: np.ndarray, **kwargs - ): - e2c = grid.connectivities[E2CDim] - p_cc_e2c = p_cc[e2c] - - z_mflx_low = 0.5 * ( - p_mass_flx_e * (p_cc_e2c[:, 0] + p_cc_e2c[:, 1]) - # - np.absolute(p_mass_flx_e) * (p_cc_e2c[:, 1] - p_cc_e2c[:, 0]) - ) - - z_anti = p_mflx_tracer_h - z_mflx_low - - return dict(z_mflx_low=z_mflx_low, z_anti=z_anti) - - @pytest.fixture - def input_data(self, grid): - p_mflx_tracer_h = random_field(grid, EdgeDim, KDim) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - p_cc = random_field(grid, CellDim, KDim) - z_mflx_low = zero_field(grid, EdgeDim, KDim) - z_anti = zero_field(grid, EdgeDim, KDim) - - return dict( - p_mflx_tracer_h=p_mflx_tracer_h, - p_mass_flx_e=p_mass_flx_e, - p_cc=p_cc, - z_mflx_low=z_mflx_low, - z_anti=z_anti, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01b.py deleted file mode 100644 index 70181f6105..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_01b.py +++ /dev/null @@ -1,120 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01b import ( - hflx_limiter_mo_stencil_01b, -) -from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - as_1D_sparse_field, - random_field, - reshape, - zero_field, -) - - -class TestHflxLimiterMoStencil01b(StencilTest): - PROGRAM = hflx_limiter_mo_stencil_01b - OUTPUTS = ( - "z_mflx_anti_in", - "z_mflx_anti_out", - "z_tracer_new_low", - "z_tracer_max", - "z_tracer_min", - ) - - @staticmethod - def reference( - grid, - geofac_div: np.ndarray, - p_rhodz_now: np.ndarray, - p_rhodz_new: np.ndarray, - z_mflx_low: np.ndarray, - z_anti: np.ndarray, - p_cc: np.ndarray, - p_dtime: float, - **kwargs, - ): - c2e = grid.connectivities[C2EDim] - z_anti_c2e = z_anti[c2e] - - geofac_div = reshape(geofac_div, c2e.shape) - geofac_div = np.expand_dims(geofac_div, axis=-1) - - zero_array = np.zeros(p_rhodz_now.shape) - - z_mflx_anti_1 = p_dtime * geofac_div[:, 0] / p_rhodz_new * z_anti_c2e[:, 0] - z_mflx_anti_2 = p_dtime * geofac_div[:, 1] / p_rhodz_new * z_anti_c2e[:, 1] - z_mflx_anti_3 = p_dtime * geofac_div[:, 2] / p_rhodz_new * z_anti_c2e[:, 2] - - z_mflx_anti_in = -1.0 * ( - np.minimum(zero_array, z_mflx_anti_1) - + np.minimum(zero_array, z_mflx_anti_2) - + np.minimum(zero_array, z_mflx_anti_3) - ) - - z_mflx_anti_out = ( - np.maximum(zero_array, z_mflx_anti_1) - + np.maximum(zero_array, z_mflx_anti_2) - + np.maximum(zero_array, z_mflx_anti_3) - ) - - z_fluxdiv_c = np.sum(z_mflx_low[c2e] * geofac_div, axis=1) - - z_tracer_new_low = (p_cc * p_rhodz_now - p_dtime * z_fluxdiv_c) / p_rhodz_new - z_tracer_max = np.maximum(p_cc, z_tracer_new_low) - z_tracer_min = np.minimum(p_cc, z_tracer_new_low) - - return dict( - z_mflx_anti_in=z_mflx_anti_in, - z_mflx_anti_out=z_mflx_anti_out, - z_tracer_new_low=z_tracer_new_low, - z_tracer_max=z_tracer_max, - z_tracer_min=z_tracer_min, - ) - - @pytest.fixture - def input_data(self, grid): - geofac_div = random_field(grid, CellDim, C2EDim) - geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) - p_rhodz_now = random_field(grid, CellDim, KDim) - p_rhodz_new = random_field(grid, CellDim, KDim) - z_mflx_low = random_field(grid, EdgeDim, KDim) - z_anti = random_field(grid, EdgeDim, KDim) - p_cc = random_field(grid, CellDim, KDim) - p_dtime = 5.0 - - z_mflx_anti_in = zero_field(grid, CellDim, KDim) - z_mflx_anti_out = zero_field(grid, CellDim, KDim) - z_tracer_new_low = zero_field(grid, CellDim, KDim) - z_tracer_max = zero_field(grid, CellDim, KDim) - z_tracer_min = zero_field(grid, CellDim, KDim) - - return dict( - geofac_div=geofac_div_new, - p_rhodz_now=p_rhodz_now, - p_rhodz_new=p_rhodz_new, - z_mflx_low=z_mflx_low, - z_anti=z_anti, - p_cc=p_cc, - p_dtime=p_dtime, - z_mflx_anti_in=z_mflx_anti_in, - z_mflx_anti_out=z_mflx_anti_out, - z_tracer_new_low=z_tracer_new_low, - z_tracer_max=z_tracer_max, - z_tracer_min=z_tracer_min, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_02.py deleted file mode 100644 index 4925d83987..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_02.py +++ /dev/null @@ -1,85 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np -import pytest -from numpy import int32 - -from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_02 import hflx_limiter_mo_stencil_02 -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - constant_field, - random_field, - zero_field, -) - - -class TestHflxLimiterMoStencil02(StencilTest): - PROGRAM = hflx_limiter_mo_stencil_02 - OUTPUTS = ("z_tracer_new_low", "z_tracer_max", "z_tracer_min") - - @staticmethod - def reference( - grid, - refin_ctrl: np.ndarray, - p_cc: np.ndarray, - z_tracer_new_low: np.ndarray, - z_tracer_max: np.ndarray, - z_tracer_min: np.ndarray, - lo_bound: float, - hi_bound: float, - **kwargs, - ): - refin_ctrl = np.expand_dims(refin_ctrl, axis=1) - condition = np.logical_or( - np.equal(refin_ctrl, lo_bound * np.ones(refin_ctrl.shape, dtype=int32)), - np.equal(refin_ctrl, hi_bound * np.ones(refin_ctrl.shape, dtype=int32)), - ) - z_tracer_new_out = np.where( - condition, - np.minimum(1.1 * p_cc, np.maximum(0.9 * p_cc, z_tracer_new_low)), - z_tracer_new_low, - ) - z_tracer_max_out = np.where(condition, np.maximum(p_cc, z_tracer_new_out), z_tracer_max) - z_tracer_min_out = np.where(condition, np.minimum(p_cc, z_tracer_new_out), z_tracer_min) - return dict( - z_tracer_new_low=z_tracer_new_out, - z_tracer_max=z_tracer_max_out, - z_tracer_min=z_tracer_min_out, - ) - - @pytest.fixture() - def input_data(self, grid): - hi_bound, lo_bound = 3, 1 - refin_ctrl = constant_field(grid, 2, CellDim, dtype=int32) - p_cc = random_field(grid, CellDim, KDim) - z_tracer_new_low_in = random_field(grid, CellDim, KDim) - z_tracer_max_in = random_field(grid, CellDim, KDim) - z_tracer_min_in = random_field(grid, CellDim, KDim) - - z_tracer_new_low_out = zero_field(grid, CellDim, KDim) - z_tracer_max_out = zero_field(grid, CellDim, KDim) - z_tracer_min_out = zero_field(grid, CellDim, KDim) - - return dict( - refin_ctrl=refin_ctrl, - p_cc=p_cc, - z_tracer_new_low=z_tracer_new_low_in, - z_tracer_max=z_tracer_max_in, - z_tracer_min=z_tracer_min_in, - lo_bound=lo_bound, - hi_bound=hi_bound, - z_tracer_new_low_out=z_tracer_new_low_out, - z_tracer_max_out=z_tracer_max_out, - z_tracer_min_out=z_tracer_min_out, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_03.py deleted file mode 100644 index f8819ee83e..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_03.py +++ /dev/null @@ -1,101 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_03 import ( - hflx_limiter_mo_stencil_03, - hflx_limiter_mo_stencil_03_min_max, -) -from icon4py.model.common.dimension import C2E2CDim, CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestHflxLimiterMoStencil03MinMax(StencilTest): - PROGRAM = hflx_limiter_mo_stencil_03_min_max - OUTPUTS = ("z_max", "z_min") - - @staticmethod - def reference(grid, z_tracer_max, z_tracer_min, beta_fct, r_beta_fct, **kwargs): - c2e2c = grid.connectivities[C2E2CDim] - z_max = beta_fct * np.maximum(np.max(z_tracer_max[c2e2c], axis=1), z_tracer_max) - z_min = r_beta_fct * np.minimum(np.min(z_tracer_min[c2e2c], axis=1), z_tracer_min) - return dict(z_max=z_max, z_min=z_min) - - @pytest.fixture - def input_data(self, grid): - z_tracer_max = random_field(grid, CellDim, KDim) - z_tracer_min = random_field(grid, CellDim, KDim) - beta_fct = 0.9 - r_beta_fct = 0.3 - z_max = zero_field(grid, CellDim, KDim) - z_min = zero_field(grid, CellDim, KDim) - return dict( - z_tracer_max=z_tracer_max, - z_tracer_min=z_tracer_min, - beta_fct=beta_fct, - r_beta_fct=r_beta_fct, - z_max=z_max, - z_min=z_min, - ) - - -class TestHflxLimiterMoStencil03(StencilTest): - PROGRAM = hflx_limiter_mo_stencil_03 - OUTPUTS = ("r_p", "r_m") - - @staticmethod - def reference( - grid, - z_tracer_max, - z_tracer_min, - beta_fct, - r_beta_fct, - z_mflx_anti_in, - z_mflx_anti_out, - z_tracer_new_low, - dbl_eps, - **kwargs, - ): - res = TestHflxLimiterMoStencil03MinMax.reference( - grid, z_tracer_max, z_tracer_min, beta_fct, r_beta_fct - ) - z_max, z_min = res["z_max"], res["z_min"] - r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) - r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out + dbl_eps) - return dict(r_p=r_p, r_m=r_m) - - @pytest.fixture - def input_data(self, grid): - z_tracer_max = random_field(grid, CellDim, KDim) - z_tracer_min = random_field(grid, CellDim, KDim) - beta_fct = 0.4 - r_beta_fct = 0.6 - z_mflx_anti_in = random_field(grid, CellDim, KDim) - z_mflx_anti_out = random_field(grid, CellDim, KDim) - z_tracer_new_low = random_field(grid, CellDim, KDim) - dbl_eps = 1e-5 - r_p = zero_field(grid, CellDim, KDim) - r_m = zero_field(grid, CellDim, KDim) - return dict( - z_tracer_max=z_tracer_max, - z_tracer_min=z_tracer_min, - beta_fct=beta_fct, - r_beta_fct=r_beta_fct, - z_mflx_anti_in=z_mflx_anti_in, - z_mflx_anti_out=z_mflx_anti_out, - z_tracer_new_low=z_tracer_new_low, - dbl_eps=dbl_eps, - r_p=r_p, - r_m=r_m, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_04.py deleted file mode 100644 index d1b5e3244d..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_mo_stencil_04.py +++ /dev/null @@ -1,50 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_04 import hflx_limiter_mo_stencil_04 -from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestHflxLimiterMoStencil04(StencilTest): - PROGRAM = hflx_limiter_mo_stencil_04 - OUTPUTS = ("p_mflx_tracer_h",) - - @staticmethod - def reference( - grid, z_anti: np.ndarray, r_m: np.ndarray, r_p: np.ndarray, z_mflx_low: np.ndarray, **kwargs - ): - r_frac = np.where( - z_anti >= 0, - np.minimum( - r_m[grid.connectivities[E2CDim][:, 0]], r_p[grid.connectivities[E2CDim][:, 1]] - ), - np.minimum( - r_m[grid.connectivities[E2CDim][:, 1]], r_p[grid.connectivities[E2CDim][:, 0]] - ), - ) - return dict(p_mflx_tracer_h=z_mflx_low + np.minimum(1.0, r_frac) * z_anti) - - @pytest.fixture - def input_data(self, grid): - z_anti = random_field(grid, EdgeDim, KDim, low=-2.0, high=2.0) - r_m = random_field(grid, CellDim, KDim) - r_p = random_field(grid, CellDim, KDim) - z_mflx_low = random_field(grid, EdgeDim, KDim) - p_mflx_tracer_h = zero_field(grid, EdgeDim, KDim) - return dict( - z_anti=z_anti, r_m=r_m, r_p=r_p, z_mflx_low=z_mflx_low, p_mflx_tracer_h=p_mflx_tracer_h - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_01.py deleted file mode 100644 index 9ebfe4692a..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_01.py +++ /dev/null @@ -1,77 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_01 import hflx_limiter_pd_stencil_01 -from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - as_1D_sparse_field, - random_field, - reshape, - zero_field, -) - - -class TestHflxLimiterPdStencil01(StencilTest): - PROGRAM = hflx_limiter_pd_stencil_01 - OUTPUTS = ("r_m",) - - @staticmethod - def reference( - grid, - geofac_div: np.ndarray, - p_cc: np.ndarray, - p_rhodz_now: np.ndarray, - p_mflx_tracer_h: np.ndarray, - p_dtime, - dbl_eps, - **kwargs, - ): - geofac_div = reshape(geofac_div, grid.connectivities[C2EDim].shape) - geofac_div = np.expand_dims(geofac_div, axis=-1) - p_m_0 = np.maximum( - 0.0, p_mflx_tracer_h[grid.connectivities[C2EDim][:, 0]] * geofac_div[:, 0] * p_dtime - ) - p_m_1 = np.maximum( - 0.0, p_mflx_tracer_h[grid.connectivities[C2EDim][:, 1]] * geofac_div[:, 1] * p_dtime - ) - p_m_2 = np.maximum( - 0.0, p_mflx_tracer_h[grid.connectivities[C2EDim][:, 2]] * geofac_div[:, 2] * p_dtime - ) - - p_m = p_m_0 + p_m_1 + p_m_2 - r_m = np.minimum(1.0, p_cc * p_rhodz_now / (p_m + dbl_eps)) - - return dict(r_m=r_m) - - @pytest.fixture - def input_data(self, grid): - geofac_div = random_field(grid, CellDim, C2EDim) - geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) - p_cc = random_field(grid, CellDim, KDim) - p_rhodz_now = random_field(grid, CellDim, KDim) - p_mflx_tracer_h = random_field(grid, EdgeDim, KDim) - r_m = zero_field(grid, CellDim, KDim) - p_dtime = np.float64(5) - dbl_eps = np.float64(1e-9) - return dict( - geofac_div=geofac_div_new, - p_cc=p_cc, - p_rhodz_now=p_rhodz_now, - p_mflx_tracer_h=p_mflx_tracer_h, - p_dtime=p_dtime, - dbl_eps=dbl_eps, - r_m=r_m, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py deleted file mode 100644 index 166f376e41..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_hflx_limiter_pd_stencil_02.py +++ /dev/null @@ -1,44 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_02 import hflx_limiter_pd_stencil_02 -from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field - - -class TestHflxLimiterPdStencil02(StencilTest): - PROGRAM = hflx_limiter_pd_stencil_02 - OUTPUTS = ("p_mflx_tracer_h",) - - @staticmethod - def reference(grid, r_m, p_mflx_tracer_h, **kwargs): - e2c = grid.connectivities[E2CDim] - r_m_e2c = r_m[e2c] - p_mflx_tracer_h_out = np.where( - p_mflx_tracer_h >= 0, - p_mflx_tracer_h * r_m_e2c[:, 0], - p_mflx_tracer_h * r_m_e2c[:, 1], - ) - return dict(p_mflx_tracer_h=p_mflx_tracer_h_out) - - @pytest.fixture(params=[("no_match", 4), ("everywhere_match", 7), ("partly_match", 4)]) - def input_data(self, request, grid): - r_m = random_field(grid, CellDim, KDim) - p_mflx_tracer_h_in = random_field(grid, EdgeDim, KDim) - - return dict( - r_m=r_m, - p_mflx_tracer_h=p_mflx_tracer_h_in, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_hor_adv_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_hor_adv_stencil_01.py deleted file mode 100644 index b517a76c3c..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_hor_adv_stencil_01.py +++ /dev/null @@ -1,74 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.hor_adv_stencil_01 import hor_adv_stencil_01 -from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - as_1D_sparse_field, - random_field, - reshape, - zero_field, -) - - -class TestHorAdvStencil01(StencilTest): - PROGRAM = hor_adv_stencil_01 - OUTPUTS = ("tracer_new_hor",) - - @staticmethod - def reference( - grid, - p_mflx_tracer_h: np.array, - deepatmo_divh: np.array, - tracer_now: np.array, - rhodz_now: np.array, - rhodz_new: np.array, - geofac_div: np.array, - p_dtime, - **kwargs, - ) -> np.array: - geofac_div = reshape(geofac_div, grid.connectivities[C2EDim].shape) - geofac_div = np.expand_dims(geofac_div, axis=-1) - tracer_new_hor = ( - tracer_now * rhodz_now - - p_dtime - * deepatmo_divh - * np.sum(p_mflx_tracer_h[grid.connectivities[C2EDim]] * geofac_div, axis=1) - ) / rhodz_new - return dict(tracer_new_hor=tracer_new_hor) - - @pytest.fixture - def input_data(self, grid): - p_mflx_tracer_h = random_field(grid, EdgeDim, KDim) - deepatmo_divh = random_field(grid, KDim) - tracer_now = random_field(grid, CellDim, KDim) - rhodz_now = random_field(grid, CellDim, KDim) - rhodz_new = random_field(grid, CellDim, KDim) - geofac_div = random_field(grid, CellDim, C2EDim) - geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) - p_dtime = np.float64(5.0) - tracer_new_hor = zero_field(grid, CellDim, KDim) - return dict( - p_mflx_tracer_h=p_mflx_tracer_h, - deepatmo_divh=deepatmo_divh, - tracer_now=tracer_now, - rhodz_now=rhodz_now, - rhodz_new=rhodz_new, - geofac_div=geofac_div_new, - p_dtime=p_dtime, - tracer_new_hor=tracer_new_hor, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_init_zero_c_k.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_init_zero_c_k.py deleted file mode 100644 index 0f0e3f5d74..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_init_zero_c_k.py +++ /dev/null @@ -1,32 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import pytest - -from icon4py.model.atmosphere.advection.init_zero_c_k import init_zero_c_k -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestInitZeroCK(StencilTest): - PROGRAM = init_zero_c_k - OUTPUTS = ("field",) - - @staticmethod - def reference(grid, **kwargs): - return dict(field=zero_field(grid, CellDim, KDim).asnumpy()) - - @pytest.fixture - def input_data(self, grid): - field = random_field(grid, CellDim, KDim) - return dict(field=field) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py deleted file mode 100644 index 0be7c47b80..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py +++ /dev/null @@ -1,154 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.mo_advection_traj_btraj_compute_o1_dsl import ( - mo_advection_traj_btraj_compute_o1_dsl, -) -from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - as_1D_sparse_field, - constant_field, - numpy_to_1D_sparse_field, - random_field, -) - - -class TestMoAdvectionTrajBtrajComputeO1Dsl(StencilTest): - PROGRAM = mo_advection_traj_btraj_compute_o1_dsl - OUTPUTS = ("p_cell_idx", "p_cell_blk", "p_distv_bary_1", "p_distv_bary_2") - - @staticmethod - def reference( - grid, - p_vn: np.array, - p_vt: np.array, - cell_idx: np.array, - cell_blk: np.array, - pos_on_tplane_e_1: np.array, - pos_on_tplane_e_2: np.array, - primal_normal_cell_1: np.array, - dual_normal_cell_1: np.array, - primal_normal_cell_2: np.array, - dual_normal_cell_2: np.array, - p_dthalf: float, - **kwargs, - ) -> dict: - e2c = grid.connectivities[E2CDim] - cell_idx = cell_idx.reshape(e2c.shape) - cell_blk = cell_blk.reshape(e2c.shape) - pos_on_tplane_e_1 = pos_on_tplane_e_1.reshape(e2c.shape) - pos_on_tplane_e_2 = pos_on_tplane_e_2.reshape(e2c.shape) - primal_normal_cell_1 = primal_normal_cell_1.reshape(e2c.shape) - primal_normal_cell_2 = primal_normal_cell_2.reshape(e2c.shape) - dual_normal_cell_1 = dual_normal_cell_1.reshape(e2c.shape) - dual_normal_cell_2 = dual_normal_cell_2.reshape(e2c.shape) - - lvn_pos = np.where(p_vn > 0.0, True, False) - cell_idx = np.expand_dims(cell_idx, axis=-1) - cell_blk = np.expand_dims(cell_blk, axis=-1) - pos_on_tplane_e_1 = np.expand_dims(pos_on_tplane_e_1, axis=-1) - pos_on_tplane_e_2 = np.expand_dims(pos_on_tplane_e_2, axis=-1) - primal_normal_cell_1 = np.expand_dims(primal_normal_cell_1, axis=-1) - dual_normal_cell_1 = np.expand_dims(dual_normal_cell_1, axis=-1) - primal_normal_cell_2 = np.expand_dims(primal_normal_cell_2, axis=-1) - dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) - - p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) - p_cell_rel_idx_dsl = np.where(lvn_pos, int32(0), int32(1)) - p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) - - z_ntdistv_bary_1 = -( - p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) - ) - z_ntdistv_bary_2 = -( - p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) - ) - - p_distv_bary_1 = np.where( - lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] - + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] - + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], - ) - - p_distv_bary_2 = np.where( - lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] - + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] - + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], - ) - - return dict( - p_cell_idx=p_cell_idx, - p_cell_rel_idx_dsl=p_cell_rel_idx_dsl, - p_cell_blk=p_cell_blk, - p_distv_bary_1=p_distv_bary_1, - p_distv_bary_2=p_distv_bary_2, - ) - - @pytest.fixture - def input_data(self, grid): - p_vn = random_field(grid, EdgeDim, KDim) - p_vt = random_field(grid, EdgeDim, KDim) - cell_idx = np.asarray(grid.connectivities[E2CDim], dtype=int32) - cell_idx_new = numpy_to_1D_sparse_field(cell_idx, ECDim) - cell_blk = constant_field(grid, 1, EdgeDim, E2CDim, dtype=int32) - cell_blk_new = as_1D_sparse_field(cell_blk, ECDim) - pos_on_tplane_e_1 = random_field(grid, EdgeDim, E2CDim) - pos_on_tplane_e_1_new = as_1D_sparse_field(pos_on_tplane_e_1, ECDim) - pos_on_tplane_e_2 = random_field(grid, EdgeDim, E2CDim) - pos_on_tplane_e_2_new = as_1D_sparse_field(pos_on_tplane_e_2, ECDim) - primal_normal_cell_1 = random_field(grid, EdgeDim, E2CDim) - primal_normal_cell_1_new = as_1D_sparse_field(primal_normal_cell_1, ECDim) - dual_normal_cell_1 = random_field(grid, EdgeDim, E2CDim) - dual_normal_cell_1_new = as_1D_sparse_field(dual_normal_cell_1, ECDim) - primal_normal_cell_2 = random_field(grid, EdgeDim, E2CDim) - primal_normal_cell_2_new = as_1D_sparse_field(primal_normal_cell_2, ECDim) - dual_normal_cell_2 = random_field(grid, EdgeDim, E2CDim) - dual_normal_cell_2_new = as_1D_sparse_field(dual_normal_cell_2, ECDim) - p_cell_idx = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) - p_cell_rel_idx_dsl = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) - p_cell_blk = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) - p_distv_bary_1 = random_field(grid, EdgeDim, KDim) - p_distv_bary_2 = random_field(grid, EdgeDim, KDim) - p_dthalf = 2.0 - - return dict( - p_vn=p_vn, - p_vt=p_vt, - cell_idx=cell_idx_new, - cell_blk=cell_blk_new, - pos_on_tplane_e_1=pos_on_tplane_e_1_new, - pos_on_tplane_e_2=pos_on_tplane_e_2_new, - primal_normal_cell_1=primal_normal_cell_1_new, - dual_normal_cell_1=dual_normal_cell_1_new, - primal_normal_cell_2=primal_normal_cell_2_new, - dual_normal_cell_2=dual_normal_cell_2_new, - p_cell_idx=p_cell_idx, - p_cell_rel_idx_dsl=p_cell_rel_idx_dsl, - p_cell_blk=p_cell_blk, - p_distv_bary_1=p_distv_bary_1, - p_distv_bary_2=p_distv_bary_2, - p_dthalf=p_dthalf, - horizontal_start=int32(0), - horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), - vertical_end=int32(grid.num_levels), - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py deleted file mode 100644 index f883693ab1..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py +++ /dev/null @@ -1,636 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.prep_gauss_quadrature_c_list_stencil import ( - prep_gauss_quadrature_c_list_stencil, -) -from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - constant_field, - random_field, - zero_field, -) - - -@pytest.mark.slow_tests -class TestPrepGaussQuadratureCListStencil(StencilTest): - PROGRAM = prep_gauss_quadrature_c_list_stencil - OUTPUTS = ( - "p_quad_vector_sum_1", - "p_quad_vector_sum_2", - "p_quad_vector_sum_3", - "p_quad_vector_sum_4", - "p_quad_vector_sum_5", - "p_quad_vector_sum_6", - "p_quad_vector_sum_7", - "p_quad_vector_sum_8", - "p_quad_vector_sum_9", - "p_quad_vector_sum_10", - "p_dreg_area", - ) - - @staticmethod - def _compute_wgt_t_detjac( - wgt_zeta_1, - wgt_eta_1, - wgt_zeta_2, - wgt_eta_2, - eta_1, - eta_2, - eta_3, - eta_4, - zeta_1, - zeta_2, - zeta_3, - zeta_4, - famask_int, - p_coords_dreg_v_1_x, - p_coords_dreg_v_2_x, - p_coords_dreg_v_3_x, - p_coords_dreg_v_4_x, - p_coords_dreg_v_1_y, - p_coords_dreg_v_2_y, - p_coords_dreg_v_3_y, - p_coords_dreg_v_4_y, - dbl_eps, - ): - z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 - z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 - z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 - z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 - - z_eta_1_1, z_eta_2_1, z_eta_3_1, z_eta_4_1 = ( - 1.0 - eta_1, - 1.0 - eta_2, - 1.0 - eta_3, - 1.0 - eta_4, - ) - z_eta_1_2, z_eta_2_2, z_eta_3_2, z_eta_4_2 = ( - 1.0 + eta_1, - 1.0 + eta_2, - 1.0 + eta_3, - 1.0 + eta_4, - ) - z_eta_1_3, z_eta_2_3, z_eta_3_3, z_eta_4_3 = ( - 1.0 - zeta_1, - 1.0 - zeta_2, - 1.0 - zeta_3, - 1.0 - zeta_4, - ) - z_eta_1_4, z_eta_2_4, z_eta_3_4, z_eta_4_4 = ( - 1.0 + zeta_1, - 1.0 + zeta_2, - 1.0 + zeta_3, - 1.0 + zeta_4, - ) - - famask_bool = np.where(famask_int == np.int32(1), True, False) - - p_coords_dreg_v_1_x = np.where(famask_bool, p_coords_dreg_v_1_x, 0.0) - p_coords_dreg_v_2_x = np.where(famask_bool, p_coords_dreg_v_2_x, 0.0) - p_coords_dreg_v_3_x = np.where(famask_bool, p_coords_dreg_v_3_x, 0.0) - p_coords_dreg_v_4_x = np.where(famask_bool, p_coords_dreg_v_4_x, 0.0) - p_coords_dreg_v_1_y = np.where(famask_bool, p_coords_dreg_v_1_y, 0.0) - p_coords_dreg_v_2_y = np.where(famask_bool, p_coords_dreg_v_2_y, 0.0) - p_coords_dreg_v_3_y = np.where(famask_bool, p_coords_dreg_v_3_y, 0.0) - p_coords_dreg_v_4_y = np.where(famask_bool, p_coords_dreg_v_4_y, 0.0) - - wgt_t_detjac_1 = np.where( - famask_bool, - dbl_eps - + z_wgt_1 - * ( - ( - z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ), - 0.0, - ) - - wgt_t_detjac_2 = np.where( - famask_bool, - dbl_eps - + z_wgt_2 - * ( - ( - z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ), - 0.0, - ) - - wgt_t_detjac_3 = np.where( - famask_bool, - dbl_eps - + z_wgt_3 - * ( - ( - z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ), - 0.0, - ) - wgt_t_detjac_4 = np.where( - famask_bool, - dbl_eps - + z_wgt_4 - * ( - ( - z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ), - 0.0, - ) - - return wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 - - @staticmethod - def _compute_z_gauss_points( - p_coords_dreg_v_1_x, - p_coords_dreg_v_2_x, - p_coords_dreg_v_3_x, - p_coords_dreg_v_4_x, - p_coords_dreg_v_1_y, - p_coords_dreg_v_2_y, - p_coords_dreg_v_3_y, - p_coords_dreg_v_4_y, - shape_func_1_1, - shape_func_2_1, - shape_func_3_1, - shape_func_4_1, - shape_func_1_2, - shape_func_2_2, - shape_func_3_2, - shape_func_4_2, - shape_func_1_3, - shape_func_2_3, - shape_func_3_3, - shape_func_4_3, - shape_func_1_4, - shape_func_2_4, - shape_func_3_4, - shape_func_4_4, - ): - z_gauss_pts_1_x = ( - shape_func_1_1 * p_coords_dreg_v_1_x - + shape_func_2_1 * p_coords_dreg_v_2_x - + shape_func_3_1 * p_coords_dreg_v_3_x - + shape_func_4_1 * p_coords_dreg_v_4_x - ) - z_gauss_pts_1_y = ( - shape_func_1_1 * p_coords_dreg_v_1_y - + shape_func_2_1 * p_coords_dreg_v_2_y - + shape_func_3_1 * p_coords_dreg_v_3_y - + shape_func_4_1 * p_coords_dreg_v_4_y - ) - z_gauss_pts_2_x = ( - shape_func_1_2 * p_coords_dreg_v_1_x - + shape_func_2_2 * p_coords_dreg_v_2_x - + shape_func_3_2 * p_coords_dreg_v_3_x - + shape_func_4_2 * p_coords_dreg_v_4_x - ) - z_gauss_pts_2_y = ( - shape_func_1_2 * p_coords_dreg_v_1_y - + shape_func_2_2 * p_coords_dreg_v_2_y - + shape_func_3_2 * p_coords_dreg_v_3_y - + shape_func_4_2 * p_coords_dreg_v_4_y - ) - z_gauss_pts_3_x = ( - shape_func_1_3 * p_coords_dreg_v_1_x - + shape_func_2_3 * p_coords_dreg_v_2_x - + shape_func_3_3 * p_coords_dreg_v_3_x - + shape_func_4_3 * p_coords_dreg_v_4_x - ) - z_gauss_pts_3_y = ( - shape_func_1_3 * p_coords_dreg_v_1_y - + shape_func_2_3 * p_coords_dreg_v_2_y - + shape_func_3_3 * p_coords_dreg_v_3_y - + shape_func_4_3 * p_coords_dreg_v_4_y - ) - z_gauss_pts_4_x = ( - shape_func_1_4 * p_coords_dreg_v_1_x - + shape_func_2_4 * p_coords_dreg_v_2_x - + shape_func_3_4 * p_coords_dreg_v_3_x - + shape_func_4_4 * p_coords_dreg_v_4_x - ) - z_gauss_pts_4_y = ( - shape_func_1_4 * p_coords_dreg_v_1_y - + shape_func_2_4 * p_coords_dreg_v_2_y - + shape_func_3_4 * p_coords_dreg_v_3_y - + shape_func_4_4 * p_coords_dreg_v_4_y - ) - return ( - z_gauss_pts_1_x, - z_gauss_pts_1_y, - z_gauss_pts_2_x, - z_gauss_pts_2_y, - z_gauss_pts_3_x, - z_gauss_pts_3_y, - z_gauss_pts_4_x, - z_gauss_pts_4_y, - ) - - @staticmethod - def _compute_vector_sums( - wgt_t_detjac_1, - wgt_t_detjac_2, - wgt_t_detjac_3, - wgt_t_detjac_4, - z_gauss_pts_1_x, - z_gauss_pts_1_y, - z_gauss_pts_2_x, - z_gauss_pts_2_y, - z_gauss_pts_3_x, - z_gauss_pts_3_y, - z_gauss_pts_4_x, - z_gauss_pts_4_y, - ): - p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 - p_quad_vector_sum_2 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x - ) - p_quad_vector_sum_3 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y - ) - p_quad_vector_sum_4 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x - ) - p_quad_vector_sum_5 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - p_quad_vector_sum_6 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y - ) - p_quad_vector_sum_7 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x - ) - p_quad_vector_sum_8 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - p_quad_vector_sum_9 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y - ) - p_quad_vector_sum_10 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - return ( - p_quad_vector_sum_1, - p_quad_vector_sum_2, - p_quad_vector_sum_3, - p_quad_vector_sum_4, - p_quad_vector_sum_5, - p_quad_vector_sum_6, - p_quad_vector_sum_7, - p_quad_vector_sum_8, - p_quad_vector_sum_9, - p_quad_vector_sum_10, - ) - - @classmethod - def reference( - cls, - grid, - famask_int: np.array, - p_coords_dreg_v_1_x: np.array, - p_coords_dreg_v_2_x: np.array, - p_coords_dreg_v_3_x: np.array, - p_coords_dreg_v_4_x: np.array, - p_coords_dreg_v_1_y: np.array, - p_coords_dreg_v_2_y: np.array, - p_coords_dreg_v_3_y: np.array, - p_coords_dreg_v_4_y: np.array, - shape_func_1_1: float, - shape_func_2_1: float, - shape_func_3_1: float, - shape_func_4_1: float, - shape_func_1_2: float, - shape_func_2_2: float, - shape_func_3_2: float, - shape_func_4_2: float, - shape_func_1_3: float, - shape_func_2_3: float, - shape_func_3_3: float, - shape_func_4_3: float, - shape_func_1_4: float, - shape_func_2_4: float, - shape_func_3_4: float, - shape_func_4_4: float, - zeta_1: float, - zeta_2: float, - zeta_3: float, - zeta_4: float, - eta_1: float, - eta_2: float, - eta_3: float, - eta_4: float, - wgt_zeta_1: float, - wgt_zeta_2: float, - wgt_eta_1: float, - wgt_eta_2: float, - dbl_eps: float, - eps: float, - p_dreg_area_in: np.array, - **kwargs, - ): - wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 = cls._compute_wgt_t_detjac( - wgt_zeta_1, - wgt_eta_1, - wgt_zeta_2, - wgt_eta_2, - eta_1, - eta_2, - eta_3, - eta_4, - zeta_1, - zeta_2, - zeta_3, - zeta_4, - famask_int, - p_coords_dreg_v_1_x, - p_coords_dreg_v_2_x, - p_coords_dreg_v_3_x, - p_coords_dreg_v_4_x, - p_coords_dreg_v_1_y, - p_coords_dreg_v_2_y, - p_coords_dreg_v_3_y, - p_coords_dreg_v_4_y, - dbl_eps, - ) - - ( - z_gauss_pts_1_x, - z_gauss_pts_1_y, - z_gauss_pts_2_x, - z_gauss_pts_2_y, - z_gauss_pts_3_x, - z_gauss_pts_3_y, - z_gauss_pts_4_x, - z_gauss_pts_4_y, - ) = cls._compute_z_gauss_points( - p_coords_dreg_v_1_x, - p_coords_dreg_v_2_x, - p_coords_dreg_v_3_x, - p_coords_dreg_v_4_x, - p_coords_dreg_v_1_y, - p_coords_dreg_v_2_y, - p_coords_dreg_v_3_y, - p_coords_dreg_v_4_y, - shape_func_1_1, - shape_func_2_1, - shape_func_3_1, - shape_func_4_1, - shape_func_1_2, - shape_func_2_2, - shape_func_3_2, - shape_func_4_2, - shape_func_1_3, - shape_func_2_3, - shape_func_3_3, - shape_func_4_3, - shape_func_1_4, - shape_func_2_4, - shape_func_3_4, - shape_func_4_4, - ) - - ( - p_quad_vector_sum_1, - p_quad_vector_sum_2, - p_quad_vector_sum_3, - p_quad_vector_sum_4, - p_quad_vector_sum_5, - p_quad_vector_sum_6, - p_quad_vector_sum_7, - p_quad_vector_sum_8, - p_quad_vector_sum_9, - p_quad_vector_sum_10, - ) = cls._compute_vector_sums( - wgt_t_detjac_1, - wgt_t_detjac_2, - wgt_t_detjac_3, - wgt_t_detjac_4, - z_gauss_pts_1_x, - z_gauss_pts_1_y, - z_gauss_pts_2_x, - z_gauss_pts_2_y, - z_gauss_pts_3_x, - z_gauss_pts_3_y, - z_gauss_pts_4_x, - z_gauss_pts_4_y, - ) - - p_dreg_area = p_dreg_area_in + p_quad_vector_sum_1 - return dict( - p_quad_vector_sum_1=p_quad_vector_sum_1, - p_quad_vector_sum_2=p_quad_vector_sum_2, - p_quad_vector_sum_3=p_quad_vector_sum_3, - p_quad_vector_sum_4=p_quad_vector_sum_4, - p_quad_vector_sum_5=p_quad_vector_sum_5, - p_quad_vector_sum_6=p_quad_vector_sum_6, - p_quad_vector_sum_7=p_quad_vector_sum_7, - p_quad_vector_sum_8=p_quad_vector_sum_8, - p_quad_vector_sum_9=p_quad_vector_sum_9, - p_quad_vector_sum_10=p_quad_vector_sum_10, - p_dreg_area=p_dreg_area, - ) - - @pytest.fixture - def input_data(self, grid): - famask_int = constant_field(grid, 1, EdgeDim, KDim, dtype=int32) - p_coords_dreg_v_1_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_2_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_3_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_4_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_1_y = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_2_y = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_3_y = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_4_y = random_field(grid, EdgeDim, KDim) - shape_func_1_1 = 0.001 - shape_func_2_1 = 0.001 - shape_func_3_1 = 0.001 - shape_func_4_1 = 0.001 - shape_func_1_2 = 0.001 - shape_func_2_2 = 0.001 - shape_func_3_2 = 0.001 - shape_func_4_2 = 0.001 - shape_func_1_3 = 0.001 - shape_func_2_3 = 0.001 - shape_func_3_3 = 0.001 - shape_func_4_3 = 0.001 - shape_func_1_4 = 0.001 - shape_func_2_4 = 0.001 - shape_func_3_4 = 0.001 - shape_func_4_4 = 0.001 - zeta_1 = 0.002 - zeta_2 = 0.002 - zeta_3 = 0.002 - zeta_4 = 0.002 - eta_1 = 0.5 - eta_2 = 0.5 - eta_3 = 0.5 - eta_4 = 0.5 - wgt_zeta_1 = 0.003 - wgt_zeta_2 = 0.003 - wgt_eta_1 = 0.002 - wgt_eta_2 = 0.007 - dbl_eps = np.float64(0.1) - eps = 0.1 - p_dreg_area_in = random_field(grid, EdgeDim, KDim) - p_quad_vector_sum_1 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_2 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_3 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_4 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_5 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_6 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_7 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_8 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_9 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_10 = zero_field(grid, EdgeDim, KDim) - p_dreg_area = zero_field(grid, EdgeDim, KDim) - return dict( - famask_int=famask_int, - p_coords_dreg_v_1_x=p_coords_dreg_v_1_x, - p_coords_dreg_v_2_x=p_coords_dreg_v_2_x, - p_coords_dreg_v_3_x=p_coords_dreg_v_3_x, - p_coords_dreg_v_4_x=p_coords_dreg_v_4_x, - p_coords_dreg_v_1_y=p_coords_dreg_v_1_y, - p_coords_dreg_v_2_y=p_coords_dreg_v_2_y, - p_coords_dreg_v_3_y=p_coords_dreg_v_3_y, - p_coords_dreg_v_4_y=p_coords_dreg_v_4_y, - shape_func_1_1=shape_func_1_1, - shape_func_2_1=shape_func_2_1, - shape_func_3_1=shape_func_3_1, - shape_func_4_1=shape_func_4_1, - shape_func_1_2=shape_func_1_2, - shape_func_2_2=shape_func_2_2, - shape_func_3_2=shape_func_3_2, - shape_func_4_2=shape_func_4_2, - shape_func_1_3=shape_func_1_3, - shape_func_2_3=shape_func_2_3, - shape_func_3_3=shape_func_3_3, - shape_func_4_3=shape_func_4_3, - shape_func_1_4=shape_func_1_4, - shape_func_2_4=shape_func_2_4, - shape_func_3_4=shape_func_3_4, - shape_func_4_4=shape_func_4_4, - zeta_1=zeta_1, - zeta_2=zeta_2, - zeta_3=zeta_3, - zeta_4=zeta_4, - eta_1=eta_1, - eta_2=eta_2, - eta_3=eta_3, - eta_4=eta_4, - wgt_zeta_1=wgt_zeta_1, - wgt_zeta_2=wgt_zeta_2, - wgt_eta_1=wgt_eta_1, - wgt_eta_2=wgt_eta_2, - dbl_eps=dbl_eps, - eps=eps, - p_dreg_area_in=p_dreg_area_in, - p_quad_vector_sum_1=p_quad_vector_sum_1, - p_quad_vector_sum_2=p_quad_vector_sum_2, - p_quad_vector_sum_3=p_quad_vector_sum_3, - p_quad_vector_sum_4=p_quad_vector_sum_4, - p_quad_vector_sum_5=p_quad_vector_sum_5, - p_quad_vector_sum_6=p_quad_vector_sum_6, - p_quad_vector_sum_7=p_quad_vector_sum_7, - p_quad_vector_sum_8=p_quad_vector_sum_8, - p_quad_vector_sum_9=p_quad_vector_sum_9, - p_quad_vector_sum_10=p_quad_vector_sum_10, - p_dreg_area=p_dreg_area, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_stencil.py deleted file mode 100644 index a6b7b6b732..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_prep_gauss_quadrature_c_stencil.py +++ /dev/null @@ -1,591 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.prep_gauss_quadrature_c_stencil import ( - prep_gauss_quadrature_c_stencil, -) -from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -@pytest.mark.slow_tests -class TestPrepGaussQuadratureCStencil(StencilTest): - PROGRAM = prep_gauss_quadrature_c_stencil - OUTPUTS = ( - "p_quad_vector_sum_1", - "p_quad_vector_sum_2", - "p_quad_vector_sum_3", - "p_quad_vector_sum_4", - "p_quad_vector_sum_5", - "p_quad_vector_sum_6", - "p_quad_vector_sum_7", - "p_quad_vector_sum_8", - "p_quad_vector_sum_9", - "p_quad_vector_sum_10", - "p_dreg_area_out", - ) - - @staticmethod - def _compute_wgt_t_detjac( - wgt_zeta_1, - wgt_zeta_2, - wgt_eta_1, - wgt_eta_2, - dbl_eps, - p_coords_dreg_v_1_x, - p_coords_dreg_v_2_x, - p_coords_dreg_v_3_x, - p_coords_dreg_v_4_x, - p_coords_dreg_v_1_y, - p_coords_dreg_v_2_y, - p_coords_dreg_v_3_y, - p_coords_dreg_v_4_y, - zeta_1, - zeta_2, - zeta_3, - zeta_4, - eta_1, - eta_2, - eta_3, - eta_4, - ): - # Compute z_wgt - z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 - z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 - z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 - z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 - - # Compute z_eta - z_eta_1_1, z_eta_2_1, z_eta_3_1, z_eta_4_1 = ( - 1.0 - eta_1, - 1.0 - eta_2, - 1.0 - eta_3, - 1.0 - eta_4, - ) - z_eta_1_2, z_eta_2_2, z_eta_3_2, z_eta_4_2 = ( - 1.0 + eta_1, - 1.0 + eta_2, - 1.0 + eta_3, - 1.0 + eta_4, - ) - z_eta_1_3, z_eta_2_3, z_eta_3_3, z_eta_4_3 = ( - 1.0 - zeta_1, - 1.0 - zeta_2, - 1.0 - zeta_3, - 1.0 - zeta_4, - ) - z_eta_1_4, z_eta_2_4, z_eta_3_4, z_eta_4_4 = ( - 1.0 + zeta_1, - 1.0 + zeta_2, - 1.0 + zeta_3, - 1.0 + zeta_4, - ) - - wgt_t_detjac_1 = dbl_eps + z_wgt_1 * ( - ( - z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ) - wgt_t_detjac_2 = dbl_eps + z_wgt_2 * ( - ( - z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ) - wgt_t_detjac_3 = dbl_eps + z_wgt_3 * ( - ( - z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ) - wgt_t_detjac_4 = dbl_eps + z_wgt_4 * ( - ( - z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ) - return wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 - - @staticmethod - def _compute_z_gauss_points( - p_coords_dreg_v_1_x, - p_coords_dreg_v_2_x, - p_coords_dreg_v_3_x, - p_coords_dreg_v_4_x, - p_coords_dreg_v_1_y, - p_coords_dreg_v_2_y, - p_coords_dreg_v_3_y, - p_coords_dreg_v_4_y, - shape_func_1_1, - shape_func_2_1, - shape_func_3_1, - shape_func_4_1, - shape_func_1_2, - shape_func_2_2, - shape_func_3_2, - shape_func_4_2, - shape_func_1_3, - shape_func_2_3, - shape_func_3_3, - shape_func_4_3, - shape_func_1_4, - shape_func_2_4, - shape_func_3_4, - shape_func_4_4, - ): - z_gauss_pts_1_x = ( - shape_func_1_1 * p_coords_dreg_v_1_x - + shape_func_2_1 * p_coords_dreg_v_2_x - + shape_func_3_1 * p_coords_dreg_v_3_x - + shape_func_4_1 * p_coords_dreg_v_4_x - ) - z_gauss_pts_1_y = ( - shape_func_1_1 * p_coords_dreg_v_1_y - + shape_func_2_1 * p_coords_dreg_v_2_y - + shape_func_3_1 * p_coords_dreg_v_3_y - + shape_func_4_1 * p_coords_dreg_v_4_y - ) - z_gauss_pts_2_x = ( - shape_func_1_2 * p_coords_dreg_v_1_x - + shape_func_2_2 * p_coords_dreg_v_2_x - + shape_func_3_2 * p_coords_dreg_v_3_x - + shape_func_4_2 * p_coords_dreg_v_4_x - ) - z_gauss_pts_2_y = ( - shape_func_1_2 * p_coords_dreg_v_1_y - + shape_func_2_2 * p_coords_dreg_v_2_y - + shape_func_3_2 * p_coords_dreg_v_3_y - + shape_func_4_2 * p_coords_dreg_v_4_y - ) - z_gauss_pts_3_x = ( - shape_func_1_3 * p_coords_dreg_v_1_x - + shape_func_2_3 * p_coords_dreg_v_2_x - + shape_func_3_3 * p_coords_dreg_v_3_x - + shape_func_4_3 * p_coords_dreg_v_4_x - ) - z_gauss_pts_3_y = ( - shape_func_1_3 * p_coords_dreg_v_1_y - + shape_func_2_3 * p_coords_dreg_v_2_y - + shape_func_3_3 * p_coords_dreg_v_3_y - + shape_func_4_3 * p_coords_dreg_v_4_y - ) - z_gauss_pts_4_x = ( - shape_func_1_4 * p_coords_dreg_v_1_x - + shape_func_2_4 * p_coords_dreg_v_2_x - + shape_func_3_4 * p_coords_dreg_v_3_x - + shape_func_4_4 * p_coords_dreg_v_4_x - ) - z_gauss_pts_4_y = ( - shape_func_1_4 * p_coords_dreg_v_1_y - + shape_func_2_4 * p_coords_dreg_v_2_y - + shape_func_3_4 * p_coords_dreg_v_3_y - + shape_func_4_4 * p_coords_dreg_v_4_y - ) - return ( - z_gauss_pts_1_x, - z_gauss_pts_1_y, - z_gauss_pts_2_x, - z_gauss_pts_2_y, - z_gauss_pts_3_x, - z_gauss_pts_3_y, - z_gauss_pts_4_x, - z_gauss_pts_4_y, - ) - - @staticmethod - def _compute_vector_sums( - wgt_t_detjac_1, - wgt_t_detjac_2, - wgt_t_detjac_3, - wgt_t_detjac_4, - z_gauss_pts_1_x, - z_gauss_pts_1_y, - z_gauss_pts_2_x, - z_gauss_pts_2_y, - z_gauss_pts_3_x, - z_gauss_pts_3_y, - z_gauss_pts_4_x, - z_gauss_pts_4_y, - ): - p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 - p_quad_vector_sum_2 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x - ) - p_quad_vector_sum_3 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y - ) - p_quad_vector_sum_4 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x - ) - p_quad_vector_sum_5 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - p_quad_vector_sum_6 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y - ) - p_quad_vector_sum_7 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x - ) - p_quad_vector_sum_8 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - p_quad_vector_sum_9 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y - ) - p_quad_vector_sum_10 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - return ( - p_quad_vector_sum_1, - p_quad_vector_sum_2, - p_quad_vector_sum_3, - p_quad_vector_sum_4, - p_quad_vector_sum_5, - p_quad_vector_sum_6, - p_quad_vector_sum_7, - p_quad_vector_sum_8, - p_quad_vector_sum_9, - p_quad_vector_sum_10, - ) - - @classmethod - def reference( - cls, - grid, - p_coords_dreg_v_1_x: np.array, - p_coords_dreg_v_2_x: np.array, - p_coords_dreg_v_3_x: np.array, - p_coords_dreg_v_4_x: np.array, - p_coords_dreg_v_1_y: np.array, - p_coords_dreg_v_2_y: np.array, - p_coords_dreg_v_3_y: np.array, - p_coords_dreg_v_4_y: np.array, - shape_func_1_1: float, - shape_func_2_1: float, - shape_func_3_1: float, - shape_func_4_1: float, - shape_func_1_2: float, - shape_func_2_2: float, - shape_func_3_2: float, - shape_func_4_2: float, - shape_func_1_3: float, - shape_func_2_3: float, - shape_func_3_3: float, - shape_func_4_3: float, - shape_func_1_4: float, - shape_func_2_4: float, - shape_func_3_4: float, - shape_func_4_4: float, - zeta_1: float, - zeta_2: float, - zeta_3: float, - zeta_4: float, - eta_1: float, - eta_2: float, - eta_3: float, - eta_4: float, - wgt_zeta_1: float, - wgt_zeta_2: float, - wgt_eta_1: float, - wgt_eta_2: float, - dbl_eps: float, - eps: float, - **kwargs, - ): - wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 = cls._compute_wgt_t_detjac( - wgt_zeta_1, - wgt_zeta_2, - wgt_eta_1, - wgt_eta_2, - dbl_eps, - p_coords_dreg_v_1_x, - p_coords_dreg_v_2_x, - p_coords_dreg_v_3_x, - p_coords_dreg_v_4_x, - p_coords_dreg_v_1_y, - p_coords_dreg_v_2_y, - p_coords_dreg_v_3_y, - p_coords_dreg_v_4_y, - zeta_1, - zeta_2, - zeta_3, - zeta_4, - eta_1, - eta_2, - eta_3, - eta_4, - ) - - ( - z_gauss_pts_1_x, - z_gauss_pts_1_y, - z_gauss_pts_2_x, - z_gauss_pts_2_y, - z_gauss_pts_3_x, - z_gauss_pts_3_y, - z_gauss_pts_4_x, - z_gauss_pts_4_y, - ) = cls._compute_z_gauss_points( - p_coords_dreg_v_1_x, - p_coords_dreg_v_2_x, - p_coords_dreg_v_3_x, - p_coords_dreg_v_4_x, - p_coords_dreg_v_1_y, - p_coords_dreg_v_2_y, - p_coords_dreg_v_3_y, - p_coords_dreg_v_4_y, - shape_func_1_1, - shape_func_2_1, - shape_func_3_1, - shape_func_4_1, - shape_func_1_2, - shape_func_2_2, - shape_func_3_2, - shape_func_4_2, - shape_func_1_3, - shape_func_2_3, - shape_func_3_3, - shape_func_4_3, - shape_func_1_4, - shape_func_2_4, - shape_func_3_4, - shape_func_4_4, - ) - - ( - p_quad_vector_sum_1, - p_quad_vector_sum_2, - p_quad_vector_sum_3, - p_quad_vector_sum_4, - p_quad_vector_sum_5, - p_quad_vector_sum_6, - p_quad_vector_sum_7, - p_quad_vector_sum_8, - p_quad_vector_sum_9, - p_quad_vector_sum_10, - ) = cls._compute_vector_sums( - wgt_t_detjac_1, - wgt_t_detjac_2, - wgt_t_detjac_3, - wgt_t_detjac_4, - z_gauss_pts_1_x, - z_gauss_pts_1_y, - z_gauss_pts_2_x, - z_gauss_pts_2_y, - z_gauss_pts_3_x, - z_gauss_pts_3_y, - z_gauss_pts_4_x, - z_gauss_pts_4_y, - ) - - z_area = p_quad_vector_sum_1 - p_dreg_area_out = np.where( - z_area >= 0.0, - np.maximum(eps, np.absolute(z_area)), - -np.maximum(eps, np.absolute(z_area)), - ) - return dict( - p_quad_vector_sum_1=p_quad_vector_sum_1, - p_quad_vector_sum_2=p_quad_vector_sum_2, - p_quad_vector_sum_3=p_quad_vector_sum_3, - p_quad_vector_sum_4=p_quad_vector_sum_4, - p_quad_vector_sum_5=p_quad_vector_sum_5, - p_quad_vector_sum_6=p_quad_vector_sum_6, - p_quad_vector_sum_7=p_quad_vector_sum_7, - p_quad_vector_sum_8=p_quad_vector_sum_8, - p_quad_vector_sum_9=p_quad_vector_sum_9, - p_quad_vector_sum_10=p_quad_vector_sum_10, - p_dreg_area_out=p_dreg_area_out, - ) - - @pytest.fixture - def input_data(self, grid): - p_coords_dreg_v_1_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_2_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_3_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_4_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_1_y = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_2_y = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_3_y = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_4_y = random_field(grid, EdgeDim, KDim) - shape_func_1_1 = 0.001 - shape_func_2_1 = 0.001 - shape_func_3_1 = 0.001 - shape_func_4_1 = 0.001 - shape_func_1_2 = 0.001 - shape_func_2_2 = 0.001 - shape_func_3_2 = 0.001 - shape_func_4_2 = 0.001 - shape_func_1_3 = 0.001 - shape_func_2_3 = 0.001 - shape_func_3_3 = 0.001 - shape_func_4_3 = 0.001 - shape_func_1_4 = 0.001 - shape_func_2_4 = 0.001 - shape_func_3_4 = 0.001 - shape_func_4_4 = 0.001 - zeta_1 = 0.002 - zeta_2 = 0.002 - zeta_3 = 0.002 - zeta_4 = 0.002 - eta_1 = 0.5 - eta_2 = 0.5 - eta_3 = 0.5 - eta_4 = 0.5 - wgt_zeta_1 = 0.003 - wgt_zeta_2 = 0.003 - wgt_eta_1 = 0.002 - wgt_eta_2 = 0.007 - dbl_eps = np.float64(0.1) - eps = 0.1 - p_quad_vector_sum_1 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_2 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_3 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_4 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_5 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_6 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_7 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_8 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_9 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_10 = zero_field(grid, EdgeDim, KDim) - p_dreg_area_out = zero_field(grid, EdgeDim, KDim) - return dict( - p_coords_dreg_v_1_x=p_coords_dreg_v_1_x, - p_coords_dreg_v_2_x=p_coords_dreg_v_2_x, - p_coords_dreg_v_3_x=p_coords_dreg_v_3_x, - p_coords_dreg_v_4_x=p_coords_dreg_v_4_x, - p_coords_dreg_v_1_y=p_coords_dreg_v_1_y, - p_coords_dreg_v_2_y=p_coords_dreg_v_2_y, - p_coords_dreg_v_3_y=p_coords_dreg_v_3_y, - p_coords_dreg_v_4_y=p_coords_dreg_v_4_y, - shape_func_1_1=shape_func_1_1, - shape_func_2_1=shape_func_2_1, - shape_func_3_1=shape_func_3_1, - shape_func_4_1=shape_func_4_1, - shape_func_1_2=shape_func_1_2, - shape_func_2_2=shape_func_2_2, - shape_func_3_2=shape_func_3_2, - shape_func_4_2=shape_func_4_2, - shape_func_1_3=shape_func_1_3, - shape_func_2_3=shape_func_2_3, - shape_func_3_3=shape_func_3_3, - shape_func_4_3=shape_func_4_3, - shape_func_1_4=shape_func_1_4, - shape_func_2_4=shape_func_2_4, - shape_func_3_4=shape_func_3_4, - shape_func_4_4=shape_func_4_4, - zeta_1=zeta_1, - zeta_2=zeta_2, - zeta_3=zeta_3, - zeta_4=zeta_4, - eta_1=eta_1, - eta_2=eta_2, - eta_3=eta_3, - eta_4=eta_4, - wgt_zeta_1=wgt_zeta_1, - wgt_zeta_2=wgt_zeta_2, - wgt_eta_1=wgt_eta_1, - wgt_eta_2=wgt_eta_2, - dbl_eps=dbl_eps, - eps=eps, - p_quad_vector_sum_1=p_quad_vector_sum_1, - p_quad_vector_sum_2=p_quad_vector_sum_2, - p_quad_vector_sum_3=p_quad_vector_sum_3, - p_quad_vector_sum_4=p_quad_vector_sum_4, - p_quad_vector_sum_5=p_quad_vector_sum_5, - p_quad_vector_sum_6=p_quad_vector_sum_6, - p_quad_vector_sum_7=p_quad_vector_sum_7, - p_quad_vector_sum_8=p_quad_vector_sum_8, - p_quad_vector_sum_9=p_quad_vector_sum_9, - p_quad_vector_sum_10=p_quad_vector_sum_10, - p_dreg_area_out=p_dreg_area_out, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_rbf_intp_edge_stencil_01.py deleted file mode 100644 index 47882226d9..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_rbf_intp_edge_stencil_01.py +++ /dev/null @@ -1,37 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.rbf_intp_edge_stencil_01 import rbf_intp_edge_stencil_01 -from icon4py.model.common.dimension import E2C2EDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestRbfIntpEdgeStencil01(StencilTest): - PROGRAM = rbf_intp_edge_stencil_01 - OUTPUTS = ("p_vt_out",) - - @staticmethod - def reference(grid, p_vn_in: np.array, ptr_coeff: np.array, **kwargs) -> np.array: - ptr_coeff = np.expand_dims(ptr_coeff, axis=-1) - p_vt_out = np.sum(p_vn_in[grid.connectivities[E2C2EDim]] * ptr_coeff, axis=1) - return dict(p_vt_out=p_vt_out) - - @pytest.fixture - def input_data(self, grid): - p_vn_in = random_field(grid, EdgeDim, KDim) - ptr_coeff = random_field(grid, EdgeDim, E2C2EDim) - p_vt_out = zero_field(grid, EdgeDim, KDim) - return dict(p_vn_in=p_vn_in, ptr_coeff=ptr_coeff, p_vt_out=p_vt_out) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_c_svd_stencil.py deleted file mode 100644 index c3c5bfec8c..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_c_svd_stencil.py +++ /dev/null @@ -1,318 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.recon_lsq_cell_c_svd_stencil import ( - recon_lsq_cell_c_svd_stencil, -) -from icon4py.model.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - as_1D_sparse_field, - random_field, - reshape, - zero_field, -) - - -@pytest.mark.slow_tests -class TestReconLsqCellCSvdStencil(StencilTest): - PROGRAM = recon_lsq_cell_c_svd_stencil - OUTPUTS = ( - "p_coeff_1_dsl", - "p_coeff_2_dsl", - "p_coeff_3_dsl", - "p_coeff_4_dsl", - "p_coeff_5_dsl", - "p_coeff_6_dsl", - "p_coeff_7_dsl", - "p_coeff_8_dsl", - "p_coeff_9_dsl", - "p_coeff_10_dsl", - ) - - @staticmethod - def reference( - grid, - p_cc: np.array, - lsq_pseudoinv_1: np.array, - lsq_pseudoinv_2: np.array, - lsq_pseudoinv_3: np.array, - lsq_pseudoinv_4: np.array, - lsq_pseudoinv_5: np.array, - lsq_pseudoinv_6: np.array, - lsq_pseudoinv_7: np.array, - lsq_pseudoinv_8: np.array, - lsq_pseudoinv_9: np.array, - lsq_moments_1: np.array, - lsq_moments_2: np.array, - lsq_moments_3: np.array, - lsq_moments_4: np.array, - lsq_moments_5: np.array, - lsq_moments_6: np.array, - lsq_moments_7: np.array, - lsq_moments_8: np.array, - lsq_moments_9: np.array, - **kwargs, - ): - c2e2c2e2c = grid.connectivities[C2E2C2E2CDim] - lsq_moments_1 = np.expand_dims(lsq_moments_1, axis=-1) - lsq_moments_2 = np.expand_dims(lsq_moments_2, axis=-1) - lsq_moments_3 = np.expand_dims(lsq_moments_3, axis=-1) - lsq_moments_4 = np.expand_dims(lsq_moments_4, axis=-1) - lsq_moments_5 = np.expand_dims(lsq_moments_5, axis=-1) - lsq_moments_6 = np.expand_dims(lsq_moments_6, axis=-1) - lsq_moments_7 = np.expand_dims(lsq_moments_7, axis=-1) - lsq_moments_8 = np.expand_dims(lsq_moments_8, axis=-1) - lsq_moments_9 = np.expand_dims(lsq_moments_9, axis=-1) - lsq_moments_1 = np.broadcast_to(lsq_moments_1, p_cc.shape) - lsq_moments_2 = np.broadcast_to(lsq_moments_2, p_cc.shape) - lsq_moments_3 = np.broadcast_to(lsq_moments_3, p_cc.shape) - lsq_moments_4 = np.broadcast_to(lsq_moments_4, p_cc.shape) - lsq_moments_5 = np.broadcast_to(lsq_moments_5, p_cc.shape) - lsq_moments_6 = np.broadcast_to(lsq_moments_6, p_cc.shape) - lsq_moments_7 = np.broadcast_to(lsq_moments_7, p_cc.shape) - lsq_moments_8 = np.broadcast_to(lsq_moments_8, p_cc.shape) - lsq_pseudoinv_9 = reshape(lsq_pseudoinv_9, c2e2c2e2c.shape) - lsq_pseudoinv_9 = np.expand_dims(lsq_pseudoinv_9, axis=-1) - lsq_pseudoinv_8 = reshape(lsq_pseudoinv_8, c2e2c2e2c.shape) - lsq_pseudoinv_8 = np.expand_dims(lsq_pseudoinv_8, axis=-1) - lsq_pseudoinv_7 = reshape(lsq_pseudoinv_7, c2e2c2e2c.shape) - lsq_pseudoinv_7 = np.expand_dims(lsq_pseudoinv_7, axis=-1) - lsq_pseudoinv_6 = reshape(lsq_pseudoinv_6, c2e2c2e2c.shape) - lsq_pseudoinv_6 = np.expand_dims(lsq_pseudoinv_6, axis=-1) - lsq_pseudoinv_5 = reshape(lsq_pseudoinv_5, c2e2c2e2c.shape) - lsq_pseudoinv_5 = np.expand_dims(lsq_pseudoinv_5, axis=-1) - lsq_pseudoinv_4 = reshape(lsq_pseudoinv_4, c2e2c2e2c.shape) - lsq_pseudoinv_4 = np.expand_dims(lsq_pseudoinv_4, axis=-1) - lsq_pseudoinv_3 = reshape(lsq_pseudoinv_3, c2e2c2e2c.shape) - lsq_pseudoinv_3 = np.expand_dims(lsq_pseudoinv_3, axis=-1) - lsq_pseudoinv_2 = reshape(lsq_pseudoinv_2, c2e2c2e2c.shape) - lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) - lsq_pseudoinv_1 = reshape(lsq_pseudoinv_1, c2e2c2e2c.shape) - lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) - - p_coeff_10_dsl = ( - lsq_pseudoinv_9[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_9[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_9[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_9[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_9[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_9[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_9[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_9[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_9[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) - - p_coeff_9_dsl = ( - lsq_pseudoinv_8[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_8[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_8[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_8[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_8[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_8[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_8[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_8[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_8[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) - - p_coeff_8_dsl = ( - lsq_pseudoinv_7[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_7[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_7[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_7[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_7[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_7[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_7[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_7[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_7[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) - - p_coeff_7_dsl = ( - lsq_pseudoinv_6[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_6[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_6[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_6[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_6[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_6[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_6[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_6[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_6[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) - - p_coeff_6_dsl = ( - lsq_pseudoinv_5[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_5[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_5[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_5[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_5[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_5[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_5[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_5[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_5[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) - - p_coeff_5_dsl = ( - lsq_pseudoinv_4[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_4[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_4[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_4[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_4[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_4[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_4[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_4[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_4[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) - - p_coeff_4_dsl = ( - lsq_pseudoinv_3[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_3[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_3[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_3[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_3[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_3[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_3[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_3[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_3[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) - - p_coeff_3_dsl = ( - lsq_pseudoinv_2[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_2[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_2[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_2[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_2[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_2[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_2[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_2[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_2[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) - - p_coeff_2_dsl = ( - lsq_pseudoinv_1[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_1[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_1[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_1[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_1[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_1[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_1[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_1[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_1[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) - - p_coeff_1_dsl = p_cc - ( - p_coeff_2_dsl * lsq_moments_1 - + p_coeff_3_dsl * lsq_moments_2 - + p_coeff_4_dsl * lsq_moments_3 - + p_coeff_5_dsl * lsq_moments_4 - + p_coeff_6_dsl * lsq_moments_5 - + p_coeff_7_dsl * lsq_moments_6 - + p_coeff_8_dsl * lsq_moments_7 - + p_coeff_9_dsl * lsq_moments_8 - + p_coeff_10_dsl * lsq_moments_9 - ) - return dict( - p_coeff_1_dsl=p_coeff_1_dsl, - p_coeff_2_dsl=p_coeff_2_dsl, - p_coeff_3_dsl=p_coeff_3_dsl, - p_coeff_4_dsl=p_coeff_4_dsl, - p_coeff_5_dsl=p_coeff_5_dsl, - p_coeff_6_dsl=p_coeff_6_dsl, - p_coeff_7_dsl=p_coeff_7_dsl, - p_coeff_8_dsl=p_coeff_8_dsl, - p_coeff_9_dsl=p_coeff_9_dsl, - p_coeff_10_dsl=p_coeff_10_dsl, - ) - - @pytest.fixture - def input_data(self, grid): - p_cc = random_field(grid, CellDim, KDim) - lsq_pseudoinv_1_field = as_1D_sparse_field( - random_field(grid, CellDim, C2E2C2E2CDim), CECECDim - ) - lsq_pseudoinv_2_field = as_1D_sparse_field( - random_field(grid, CellDim, C2E2C2E2CDim), CECECDim - ) - lsq_pseudoinv_3_field = as_1D_sparse_field( - random_field(grid, CellDim, C2E2C2E2CDim), CECECDim - ) - lsq_pseudoinv_4_field = as_1D_sparse_field( - random_field(grid, CellDim, C2E2C2E2CDim), CECECDim - ) - lsq_pseudoinv_5_field = as_1D_sparse_field( - random_field(grid, CellDim, C2E2C2E2CDim), CECECDim - ) - lsq_pseudoinv_6_field = as_1D_sparse_field( - random_field(grid, CellDim, C2E2C2E2CDim), CECECDim - ) - lsq_pseudoinv_7_field = as_1D_sparse_field( - random_field(grid, CellDim, C2E2C2E2CDim), CECECDim - ) - lsq_pseudoinv_8_field = as_1D_sparse_field( - random_field(grid, CellDim, C2E2C2E2CDim), CECECDim - ) - lsq_pseudoinv_9_field = as_1D_sparse_field( - random_field(grid, CellDim, C2E2C2E2CDim), CECECDim - ) - lsq_moments_1 = random_field(grid, CellDim) - lsq_moments_2 = random_field(grid, CellDim) - lsq_moments_3 = random_field(grid, CellDim) - lsq_moments_4 = random_field(grid, CellDim) - lsq_moments_5 = random_field(grid, CellDim) - lsq_moments_6 = random_field(grid, CellDim) - lsq_moments_7 = random_field(grid, CellDim) - lsq_moments_8 = random_field(grid, CellDim) - lsq_moments_9 = random_field(grid, CellDim) - p_coeff_1_dsl = zero_field(grid, CellDim, KDim) - p_coeff_2_dsl = zero_field(grid, CellDim, KDim) - p_coeff_3_dsl = zero_field(grid, CellDim, KDim) - p_coeff_4_dsl = zero_field(grid, CellDim, KDim) - p_coeff_5_dsl = zero_field(grid, CellDim, KDim) - p_coeff_6_dsl = zero_field(grid, CellDim, KDim) - p_coeff_7_dsl = zero_field(grid, CellDim, KDim) - p_coeff_8_dsl = zero_field(grid, CellDim, KDim) - p_coeff_9_dsl = zero_field(grid, CellDim, KDim) - p_coeff_10_dsl = zero_field(grid, CellDim, KDim) - return dict( - p_cc=p_cc, - lsq_pseudoinv_1=lsq_pseudoinv_1_field, - lsq_pseudoinv_2=lsq_pseudoinv_2_field, - lsq_pseudoinv_3=lsq_pseudoinv_3_field, - lsq_pseudoinv_4=lsq_pseudoinv_4_field, - lsq_pseudoinv_5=lsq_pseudoinv_5_field, - lsq_pseudoinv_6=lsq_pseudoinv_6_field, - lsq_pseudoinv_7=lsq_pseudoinv_7_field, - lsq_pseudoinv_8=lsq_pseudoinv_8_field, - lsq_pseudoinv_9=lsq_pseudoinv_9_field, - lsq_moments_1=lsq_moments_1, - lsq_moments_2=lsq_moments_2, - lsq_moments_3=lsq_moments_3, - lsq_moments_4=lsq_moments_4, - lsq_moments_5=lsq_moments_5, - lsq_moments_6=lsq_moments_6, - lsq_moments_7=lsq_moments_7, - lsq_moments_8=lsq_moments_8, - lsq_moments_9=lsq_moments_9, - p_coeff_1_dsl=p_coeff_1_dsl, - p_coeff_2_dsl=p_coeff_2_dsl, - p_coeff_3_dsl=p_coeff_3_dsl, - p_coeff_4_dsl=p_coeff_4_dsl, - p_coeff_5_dsl=p_coeff_5_dsl, - p_coeff_6_dsl=p_coeff_6_dsl, - p_coeff_7_dsl=p_coeff_7_dsl, - p_coeff_8_dsl=p_coeff_8_dsl, - p_coeff_9_dsl=p_coeff_9_dsl, - p_coeff_10_dsl=p_coeff_10_dsl, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_l_svd_stencil.py deleted file mode 100644 index 298b6e71f3..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_recon_lsq_cell_l_svd_stencil.py +++ /dev/null @@ -1,77 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.recon_lsq_cell_l_svd_stencil import ( - recon_lsq_cell_l_svd_stencil, -) -from icon4py.model.common.dimension import C2E2CDim, CECDim, CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - as_1D_sparse_field, - random_field, - reshape, - zero_field, -) - - -class TestReconLsqCellLSvdStencil(StencilTest): - PROGRAM = recon_lsq_cell_l_svd_stencil - OUTPUTS = ( - "p_coeff_1_dsl", - "p_coeff_2_dsl", - "p_coeff_3_dsl", - ) - - @staticmethod - def reference( - grid, - p_cc: np.array, - lsq_pseudoinv_1: np.array, - lsq_pseudoinv_2: np.array, - **kwargs, - ): - c2e2c = grid.connectivities[C2E2CDim] - p_cc_e = np.expand_dims(p_cc, axis=1) - n_diff = p_cc[c2e2c] - p_cc_e - lsq_pseudoinv_1 = reshape(lsq_pseudoinv_1, c2e2c.shape) - lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) - lsq_pseudoinv_2 = reshape(lsq_pseudoinv_2, c2e2c.shape) - lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) - p_coeff_1_dsl = p_cc - p_coeff_2_dsl = np.sum(lsq_pseudoinv_1 * n_diff, axis=1) - p_coeff_3_dsl = np.sum(lsq_pseudoinv_2 * n_diff, axis=1) - return dict( - p_coeff_1_dsl=p_coeff_1_dsl, - p_coeff_2_dsl=p_coeff_2_dsl, - p_coeff_3_dsl=p_coeff_3_dsl, - ) - - @pytest.fixture - def input_data(self, grid): - p_cc = random_field(grid, CellDim, KDim) - lsq_pseudoinv_1_field = as_1D_sparse_field(random_field(grid, CellDim, C2E2CDim), CECDim) - lsq_pseudoinv_2_field = as_1D_sparse_field(random_field(grid, CellDim, C2E2CDim), CECDim) - p_coeff_1_dsl = zero_field(grid, CellDim, KDim) - p_coeff_2_dsl = zero_field(grid, CellDim, KDim) - p_coeff_3_dsl = zero_field(grid, CellDim, KDim) - return dict( - p_cc=p_cc, - lsq_pseudoinv_1=lsq_pseudoinv_1_field, - lsq_pseudoinv_2=lsq_pseudoinv_2_field, - p_coeff_1_dsl=p_coeff_1_dsl, - p_coeff_2_dsl=p_coeff_2_dsl, - p_coeff_3_dsl=p_coeff_3_dsl, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_01.py deleted file mode 100644 index f45130e824..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_01.py +++ /dev/null @@ -1,64 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.step_advection_stencil_01 import step_advection_stencil_01 -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field, zero_field - - -class TestStepAdvectionStencil01(StencilTest): - PROGRAM = step_advection_stencil_01 - OUTPUTS = ( - Output( - "rhodz_ast2", - refslice=(slice(None), slice(None, -1)), - gtslice=(slice(None), slice(None, -1)), - ), - ) - - @staticmethod - def reference( - grid, - rhodz_ast: np.array, - p_mflx_contra_v: np.array, - deepatmo_divzl: np.array, - deepatmo_divzu: np.array, - p_dtime, - **kwargs, - ): - tmp = p_dtime * ( - p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu - ) - rhodz_ast2 = rhodz_ast + tmp - - return dict(rhodz_ast2=rhodz_ast2) - - @pytest.fixture - def input_data(self, grid): - rhodz_ast = random_field(grid, CellDim, KDim) - p_mflx_contra_v = random_field(grid, CellDim, KDim, extend={KDim: 1}) - deepatmo_divzl = random_field(grid, KDim) - deepatmo_divzu = random_field(grid, KDim) - rhodz_ast2 = zero_field(grid, CellDim, KDim) - p_dtime = 0.1 - return dict( - rhodz_ast=rhodz_ast, - p_mflx_contra_v=p_mflx_contra_v, - deepatmo_divzl=deepatmo_divzl, - deepatmo_divzu=deepatmo_divzu, - p_dtime=p_dtime, - rhodz_ast2=rhodz_ast2, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_02.py deleted file mode 100644 index c34cd6f037..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_02.py +++ /dev/null @@ -1,61 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.step_advection_stencil_02 import step_advection_stencil_02 -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field, zero_field - - -class TestStepAdvectionStencil02(StencilTest): - PROGRAM = step_advection_stencil_02 - OUTPUTS = ( - Output( - "rhodz_ast2", - refslice=(slice(None), slice(None, -1)), - gtslice=(slice(None), slice(None, -1)), - ), - ) - - @staticmethod - def reference( - grid, - p_rhodz_new: np.array, - p_mflx_contra_v: np.array, - deepatmo_divzl: np.array, - deepatmo_divzu: np.array, - p_dtime: float, - **kwargs, - ): - tmp = p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu - rhodz_ast2 = np.maximum(0.1 * p_rhodz_new, p_rhodz_new) - p_dtime * tmp - return dict(rhodz_ast2=rhodz_ast2) - - @pytest.fixture - def input_data(self, grid): - p_rhodz_new = random_field(grid, CellDim, KDim) - p_mflx_contra_v = random_field(grid, CellDim, KDim, extend={KDim: 1}) - deepatmo_divzl = random_field(grid, KDim) - deepatmo_divzu = random_field(grid, KDim) - p_dtime = 0.1 - rhodz_ast2 = zero_field(grid, CellDim, KDim) - return dict( - p_rhodz_new=p_rhodz_new, - p_mflx_contra_v=p_mflx_contra_v, - deepatmo_divzl=deepatmo_divzl, - deepatmo_divzu=deepatmo_divzu, - p_dtime=p_dtime, - rhodz_ast2=rhodz_ast2, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_03.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_03.py deleted file mode 100644 index 2e1f83ab08..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_03.py +++ /dev/null @@ -1,50 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.step_advection_stencil_03 import step_advection_stencil_03 -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field - - -class TestStepAdvectionStencil03(StencilTest): - PROGRAM = step_advection_stencil_03 - OUTPUTS = ("p_tracer_new",) - - @staticmethod - def reference( - grid, - p_tracer_now: np.array, - p_grf_tend_tracer: np.array, - p_dtime, - **kwargs, - ): - p_tracer_new = p_tracer_now + p_dtime * p_grf_tend_tracer - p_tracer_new = np.where(p_tracer_new < 0.0, 0.0, p_tracer_new) - - return dict(p_tracer_new=p_tracer_new) - - @pytest.fixture - def input_data(self, grid): - p_tracer_now = random_field(grid, CellDim, KDim) - p_grf_tend_tracer = random_field(grid, CellDim, KDim) - p_tracer_new = random_field(grid, CellDim, KDim) - p_dtime = np.float64(5.0) - return dict( - p_tracer_now=p_tracer_now, - p_grf_tend_tracer=p_grf_tend_tracer, - p_dtime=p_dtime, - p_tracer_new=p_tracer_new, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_04.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_04.py deleted file mode 100644 index 6541824d3a..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_step_advection_stencil_04.py +++ /dev/null @@ -1,49 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.step_advection_stencil_04 import step_advection_stencil_04 -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestStepAdvectionStencil04(StencilTest): - PROGRAM = step_advection_stencil_04 - OUTPUTS = ("opt_ddt_tracer_adv",) - - @staticmethod - def reference( - grid, - p_tracer_now: np.array, - p_tracer_new: np.array, - p_dtime, - **kwargs, - ): - opt_ddt_tracer_adv = (p_tracer_new - p_tracer_now) / p_dtime - - return dict(opt_ddt_tracer_adv=opt_ddt_tracer_adv) - - @pytest.fixture - def input_data(self, grid): - p_tracer_now = random_field(grid, CellDim, KDim) - p_tracer_new = random_field(grid, CellDim, KDim) - opt_ddt_tracer_adv = zero_field(grid, CellDim, KDim) - p_dtime = np.float64(5.0) - return dict( - p_tracer_now=p_tracer_now, - p_tracer_new=p_tracer_new, - p_dtime=p_dtime, - opt_ddt_tracer_adv=opt_ddt_tracer_adv, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura3_stencil_01.py deleted file mode 100644 index a538479e1e..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura3_stencil_01.py +++ /dev/null @@ -1,194 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.upwind_hflux_miura3_stencil_01 import ( - upwind_hflux_miura3_stencil_01, -) -from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - random_field, - random_mask, - zero_field, -) - - -class TestUpwindHfluxMiura3Stencil01(StencilTest): - PROGRAM = upwind_hflux_miura3_stencil_01 - OUTPUTS = ("p_out_e_miura3",) - - @staticmethod - def reference( - grid, - z_lsq_coeff_1: np.array, - z_lsq_coeff_2: np.array, - z_lsq_coeff_3: np.array, - z_lsq_coeff_4: np.array, - z_lsq_coeff_5: np.array, - z_lsq_coeff_6: np.array, - z_lsq_coeff_7: np.array, - z_lsq_coeff_8: np.array, - z_lsq_coeff_9: np.array, - z_lsq_coeff_10: np.array, - z_quad_vector_sum_1: np.array, - z_quad_vector_sum_2: np.array, - z_quad_vector_sum_3: np.array, - z_quad_vector_sum_4: np.array, - z_quad_vector_sum_5: np.array, - z_quad_vector_sum_6: np.array, - z_quad_vector_sum_7: np.array, - z_quad_vector_sum_8: np.array, - z_quad_vector_sum_9: np.array, - z_quad_vector_sum_10: np.array, - z_dreg_area: np.array, - p_mass_flx_e: np.array, - cell_rel_idx_dsl: np.array, - **kwargs, - ): - e2c = grid.connectivities[E2CDim] - z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] - z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] - z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] - z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] - z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] - z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] - z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] - z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] - z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] - z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] - - p_out_e_miura3 = ( - ( - np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_1_e2c[:, 1], - z_lsq_coeff_1_e2c[:, 0], - ) - * z_quad_vector_sum_1 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_2_e2c[:, 1], - z_lsq_coeff_2_e2c[:, 0], - ) - * z_quad_vector_sum_2 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_3_e2c[:, 1], - z_lsq_coeff_3_e2c[:, 0], - ) - * z_quad_vector_sum_3 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_4_e2c[:, 1], - z_lsq_coeff_4_e2c[:, 0], - ) - * z_quad_vector_sum_4 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_5_e2c[:, 1], - z_lsq_coeff_5_e2c[:, 0], - ) - * z_quad_vector_sum_5 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_6_e2c[:, 1], - z_lsq_coeff_6_e2c[:, 0], - ) - * z_quad_vector_sum_6 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_7_e2c[:, 1], - z_lsq_coeff_7_e2c[:, 0], - ) - * z_quad_vector_sum_7 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_8_e2c[:, 1], - z_lsq_coeff_8_e2c[:, 0], - ) - * z_quad_vector_sum_8 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_9_e2c[:, 1], - z_lsq_coeff_9_e2c[:, 0], - ) - * z_quad_vector_sum_9 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_10_e2c[:, 1], - z_lsq_coeff_10_e2c[:, 0], - ) - * z_quad_vector_sum_10 - ) - / z_dreg_area - * p_mass_flx_e - ) - - return dict(p_out_e_miura3=p_out_e_miura3) - - @pytest.fixture - def input_data(self, grid): - z_lsq_coeff_1 = random_field(grid, CellDim, KDim) - z_lsq_coeff_2 = random_field(grid, CellDim, KDim) - z_lsq_coeff_3 = random_field(grid, CellDim, KDim) - z_lsq_coeff_4 = random_field(grid, CellDim, KDim) - z_lsq_coeff_5 = random_field(grid, CellDim, KDim) - z_lsq_coeff_6 = random_field(grid, CellDim, KDim) - z_lsq_coeff_7 = random_field(grid, CellDim, KDim) - z_lsq_coeff_8 = random_field(grid, CellDim, KDim) - z_lsq_coeff_9 = random_field(grid, CellDim, KDim) - z_lsq_coeff_10 = random_field(grid, CellDim, KDim) - z_quad_vector_sum_1 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_2 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_3 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_4 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_5 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_6 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_7 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_8 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_9 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_10 = random_field(grid, EdgeDim, KDim) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - z_dreg_area = random_field(grid, EdgeDim, KDim) - cell_rel_idx_dsl = random_mask(grid, EdgeDim, KDim, dtype=int32) - p_out_e_miura3 = zero_field(grid, EdgeDim, KDim) - return dict( - z_lsq_coeff_1=z_lsq_coeff_1, - z_lsq_coeff_2=z_lsq_coeff_2, - z_lsq_coeff_3=z_lsq_coeff_3, - z_lsq_coeff_4=z_lsq_coeff_4, - z_lsq_coeff_5=z_lsq_coeff_5, - z_lsq_coeff_6=z_lsq_coeff_6, - z_lsq_coeff_7=z_lsq_coeff_7, - z_lsq_coeff_8=z_lsq_coeff_8, - z_lsq_coeff_9=z_lsq_coeff_9, - z_lsq_coeff_10=z_lsq_coeff_10, - z_quad_vector_sum_1=z_quad_vector_sum_1, - z_quad_vector_sum_2=z_quad_vector_sum_2, - z_quad_vector_sum_3=z_quad_vector_sum_3, - z_quad_vector_sum_4=z_quad_vector_sum_4, - z_quad_vector_sum_5=z_quad_vector_sum_5, - z_quad_vector_sum_6=z_quad_vector_sum_6, - z_quad_vector_sum_7=z_quad_vector_sum_7, - z_quad_vector_sum_8=z_quad_vector_sum_8, - z_quad_vector_sum_9=z_quad_vector_sum_9, - z_quad_vector_sum_10=z_quad_vector_sum_10, - z_dreg_area=z_dreg_area, - p_mass_flx_e=p_mass_flx_e, - cell_rel_idx_dsl=cell_rel_idx_dsl, - p_out_e_miura3=p_out_e_miura3, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py deleted file mode 100644 index 1e33e6a7a7..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py +++ /dev/null @@ -1,92 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_01 import ( - upwind_hflux_miura_cycl_stencil_01, -) -from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - random_field, - random_mask, - zero_field, -) - - -class TestUpwindHfluxMiuraCyclStencil01(StencilTest): - PROGRAM = upwind_hflux_miura_cycl_stencil_01 - OUTPUTS = ("z_tracer_mflx_dsl",) - - @staticmethod - def reference( - grid, - z_lsq_coeff_1_dsl: np.array, - z_lsq_coeff_2_dsl: np.array, - z_lsq_coeff_3_dsl: np.array, - distv_bary_1: np.array, - distv_bary_2: np.array, - p_mass_flx_e: np.array, - cell_rel_idx_dsl: np.array, - **kwargs, - ): - e2c = grid.connectivities[E2CDim] - z_lsq_coeff_1_dsl_e2c = z_lsq_coeff_1_dsl[e2c] - z_lsq_coeff_2_dsl_e2c = z_lsq_coeff_2_dsl[e2c] - z_lsq_coeff_3_dsl_e2c = z_lsq_coeff_3_dsl[e2c] - - z_tracer_mflx_dsl = ( - np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_1_dsl_e2c[:, 1], - z_lsq_coeff_1_dsl_e2c[:, 0], - ) - + distv_bary_1 - * np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_2_dsl_e2c[:, 1], - z_lsq_coeff_2_dsl_e2c[:, 0], - ) - + distv_bary_2 - * np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_3_dsl_e2c[:, 1], - z_lsq_coeff_3_dsl_e2c[:, 0], - ) - ) * p_mass_flx_e - - return dict(z_tracer_mflx_dsl=z_tracer_mflx_dsl) - - @pytest.fixture - def input_data(self, grid): - z_lsq_coeff_1_dsl = random_field(grid, CellDim, KDim) - z_lsq_coeff_2_dsl = random_field(grid, CellDim, KDim) - z_lsq_coeff_3_dsl = random_field(grid, CellDim, KDim) - distv_bary_1 = random_field(grid, EdgeDim, KDim) - distv_bary_2 = random_field(grid, EdgeDim, KDim) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - cell_rel_idx_dsl = random_mask(grid, EdgeDim, KDim, dtype=int32) - z_tracer_mflx_dsl = zero_field(grid, EdgeDim, KDim) - return dict( - z_lsq_coeff_1_dsl=z_lsq_coeff_1_dsl, - z_lsq_coeff_2_dsl=z_lsq_coeff_2_dsl, - z_lsq_coeff_3_dsl=z_lsq_coeff_3_dsl, - distv_bary_1=distv_bary_1, - distv_bary_2=distv_bary_2, - p_mass_flx_e=p_mass_flx_e, - cell_rel_idx_dsl=cell_rel_idx_dsl, - z_tracer_mflx_dsl=z_tracer_mflx_dsl, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py deleted file mode 100644 index dbb64d6807..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py +++ /dev/null @@ -1,93 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_02 import ( - upwind_hflux_miura_cycl_stencil_02, -) -from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestUpwindHfluxMiuraCyclStencil02(StencilTest): - PROGRAM = upwind_hflux_miura_cycl_stencil_02 - OUTPUTS = ( - "z_rhofluxdiv_c_out", - "z_fluxdiv_c_dsl", - "z_rho_new_dsl", - "z_tracer_new_dsl", - ) - - @staticmethod - def reference( - grid, - nsub: int32, - p_mass_flx_e: np.array, - geofac_div: np.array, - z_rhofluxdiv_c: np.array, - z_tracer_mflx: np.array, - z_rho_now: np.array, - z_tracer_now: np.array, - z_dtsub: float, - **kwargs, - ): - c2e = grid.connectivities[C2EDim] - p_mass_flx_e_c2e = p_mass_flx_e[c2e] - geofac_div = np.expand_dims(geofac_div, axis=-1) - z_tracer_mflx_c2e = z_tracer_mflx[c2e] - - z_rhofluxdiv_c_out = ( - np.sum(p_mass_flx_e_c2e * geofac_div, axis=1) if nsub == int32(1) else z_rhofluxdiv_c - ) - z_fluxdiv_c_dsl = np.sum(z_tracer_mflx_c2e * geofac_div, axis=1) - z_rho_new_dsl = z_rho_now - z_dtsub * z_rhofluxdiv_c_out - z_tracer_new_dsl = (z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl) / z_rho_new_dsl - - return dict( - z_rhofluxdiv_c_out=z_rhofluxdiv_c_out, - z_fluxdiv_c_dsl=z_fluxdiv_c_dsl, - z_rho_new_dsl=z_rho_new_dsl, - z_tracer_new_dsl=z_tracer_new_dsl, - ) - - @pytest.fixture - def input_data(self, grid): - nsub = int32(1) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - geofac_div = random_field(grid, CellDim, C2EDim) - z_rhofluxdiv_c = random_field(grid, CellDim, KDim) - z_tracer_mflx = random_field(grid, EdgeDim, KDim) - z_rho_now = random_field(grid, CellDim, KDim) - z_tracer_now = random_field(grid, CellDim, KDim) - z_dtsub = 0.5 - z_rhofluxdiv_c_out = zero_field(grid, CellDim, KDim) - z_fluxdiv_c_dsl = zero_field(grid, CellDim, KDim) - z_rho_new_dsl = zero_field(grid, CellDim, KDim) - z_tracer_new_dsl = zero_field(grid, CellDim, KDim) - return dict( - nsub=nsub, - p_mass_flx_e=p_mass_flx_e, - geofac_div=geofac_div, - z_rhofluxdiv_c=z_rhofluxdiv_c, - z_tracer_mflx=z_tracer_mflx, - z_rho_now=z_rho_now, - z_tracer_now=z_tracer_now, - z_dtsub=z_dtsub, - z_rhofluxdiv_c_out=z_rhofluxdiv_c_out, - z_fluxdiv_c_dsl=z_fluxdiv_c_dsl, - z_rho_new_dsl=z_rho_new_dsl, - z_tracer_new_dsl=z_tracer_new_dsl, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py deleted file mode 100644 index 8c1a99e509..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py +++ /dev/null @@ -1,47 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03a import ( - upwind_hflux_miura_cycl_stencil_03a, -) -from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestUpwindHfluxMiuraCyclStencil03a(StencilTest): - PROGRAM = upwind_hflux_miura_cycl_stencil_03a - OUTPUTS = ("p_out_e",) - - @staticmethod - def reference( - grid, - z_tracer_mflx_1_dsl: np.array, - z_tracer_mflx_2_dsl: np.array, - **kwargs, - ): - p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl) / float(2) - return dict(p_out_e=p_out_e) - - @pytest.fixture - def input_data(self, grid): - z_tracer_mflx_1_dsl = random_field(grid, EdgeDim, KDim) - z_tracer_mflx_2_dsl = random_field(grid, EdgeDim, KDim) - p_out_e = zero_field(grid, EdgeDim, KDim) - return dict( - z_tracer_mflx_1_dsl=z_tracer_mflx_1_dsl, - z_tracer_mflx_2_dsl=z_tracer_mflx_2_dsl, - p_out_e=p_out_e, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py deleted file mode 100644 index 481fd93c84..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py +++ /dev/null @@ -1,50 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03b import ( - upwind_hflux_miura_cycl_stencil_03b, -) -from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestUpwindHfluxMiuraCyclStencil03b(StencilTest): - PROGRAM = upwind_hflux_miura_cycl_stencil_03b - OUTPUTS = ("p_out_e",) - - @staticmethod - def reference( - grid, - z_tracer_mflx_1_dsl: np.array, - z_tracer_mflx_2_dsl: np.array, - z_tracer_mflx_3_dsl: np.array, - **kwargs, - ): - p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float(3) - return dict(p_out_e=p_out_e) - - @pytest.fixture - def input_data(self, grid): - z_tracer_mflx_1_dsl = random_field(grid, EdgeDim, KDim) - z_tracer_mflx_2_dsl = random_field(grid, EdgeDim, KDim) - z_tracer_mflx_3_dsl = random_field(grid, EdgeDim, KDim) - p_out_e = zero_field(grid, EdgeDim, KDim) - return dict( - z_tracer_mflx_1_dsl=z_tracer_mflx_1_dsl, - z_tracer_mflx_2_dsl=z_tracer_mflx_2_dsl, - z_tracer_mflx_3_dsl=z_tracer_mflx_3_dsl, - p_out_e=p_out_e, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_stencil_01.py deleted file mode 100644 index 59e3c026fe..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_hflux_miura_stencil_01.py +++ /dev/null @@ -1,92 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.upwind_hflux_miura_stencil_01 import ( - upwind_hflux_miura_stencil_01, -) -from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - constant_field, - random_field, - zero_field, -) - - -class TestUpwindHfluxMiuraStencil01(StencilTest): - PROGRAM = upwind_hflux_miura_stencil_01 - OUTPUTS = ("p_out_e",) - - @staticmethod - def reference( - grid, - z_lsq_coeff_1: np.array, - z_lsq_coeff_2: np.array, - z_lsq_coeff_3: np.array, - distv_bary_1: np.array, - distv_bary_2: np.array, - p_mass_flx_e: np.array, - cell_rel_idx_dsl: np.array, - **kwargs, - ): - e2c = grid.connectivities[E2CDim] - z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] - z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] - z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] - - p_out_e = ( - np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_1_e2c[:, 1], - z_lsq_coeff_1_e2c[:, 0], - ) - + distv_bary_1 - * np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_2_e2c[:, 1], - z_lsq_coeff_2_e2c[:, 0], - ) - + distv_bary_2 - * np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_3_e2c[:, 1], - z_lsq_coeff_3_e2c[:, 0], - ) - ) * p_mass_flx_e - - return dict(p_out_e=p_out_e) - - @pytest.fixture - def input_data(self, grid): - z_lsq_coeff_1 = random_field(grid, CellDim, KDim) - z_lsq_coeff_2 = random_field(grid, CellDim, KDim) - z_lsq_coeff_3 = random_field(grid, CellDim, KDim) - distv_bary_1 = random_field(grid, EdgeDim, KDim) - distv_bary_2 = random_field(grid, EdgeDim, KDim) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - cell_rel_idx_dsl = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) - p_out_e = zero_field(grid, EdgeDim, KDim) - return dict( - z_lsq_coeff_1=z_lsq_coeff_1, - z_lsq_coeff_2=z_lsq_coeff_2, - z_lsq_coeff_3=z_lsq_coeff_3, - distv_bary_1=distv_bary_1, - distv_bary_2=distv_bary_2, - p_mass_flx_e=p_mass_flx_e, - cell_rel_idx_dsl=cell_rel_idx_dsl, - p_out_e=p_out_e, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_vflux_ppm_stencil_01.py deleted file mode 100644 index a8f9c3303b..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_upwind_vflux_ppm_stencil_01.py +++ /dev/null @@ -1,45 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest - -from icon4py.model.atmosphere.advection.upwind_vflux_ppm_stencil_01 import ( - upwind_vflux_ppm_stencil_01, -) -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestUpwindVfluxPpmStencil01(StencilTest): - PROGRAM = upwind_vflux_ppm_stencil_01 - OUTPUTS = ("z_delta_q", "z_a1") - - @staticmethod - def reference( - grid, z_face_up: np.ndarray, z_face_low: np.ndarray, p_cc: np.ndarray, **kwargs - ) -> tuple[np.ndarray]: - z_delta_q = 0.5 * (z_face_up - z_face_low) - z_a1 = p_cc - 0.5 * (z_face_up + z_face_low) - return dict(z_delta_q=z_delta_q, z_a1=z_a1) - - @pytest.fixture - def input_data(self, grid): - z_face_up = random_field(grid, CellDim, KDim) - z_face_low = random_field(grid, CellDim, KDim) - p_cc = random_field(grid, CellDim, KDim) - z_delta_q = zero_field(grid, CellDim, KDim) - z_a1 = zero_field(grid, CellDim, KDim) - return dict( - z_face_up=z_face_up, z_face_low=z_face_low, p_cc=p_cc, z_delta_q=z_delta_q, z_a1=z_a1 - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_vert_adv_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_vert_adv_stencil_01.py deleted file mode 100644 index 14e64e37cd..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_vert_adv_stencil_01.py +++ /dev/null @@ -1,87 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next import as_field -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.vert_adv_stencil_01 import vert_adv_stencil_01 -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, _shape, random_field, zero_field - - -class TestVertAdvStencil01(StencilTest): - PROGRAM = vert_adv_stencil_01 - OUTPUTS = ("tracer_new",) - - @staticmethod - def reference( - grid, - tracer_now: np.array, - rhodz_now: np.array, - p_mflx_tracer_v: np.array, - deepatmo_divzl: np.array, - deepatmo_divzu: np.array, - rhodz_new: np.array, - k: np.array, - ivadv_tracer: int32, - iadv_slev_jt: int32, - p_dtime: np.float64, - **kwargs, - ) -> np.array: - if ivadv_tracer != int32(0): - tracer_new = np.where( - (iadv_slev_jt <= k), - ( - tracer_now * rhodz_now - + p_dtime - * ( - p_mflx_tracer_v[:, 1:] * deepatmo_divzl - - p_mflx_tracer_v[:, :-1] * deepatmo_divzu - ) - ) - / rhodz_new, - tracer_now, - ) - else: - tracer_new = tracer_now - - return dict(tracer_new=tracer_new) - - @pytest.fixture - def input_data(self, grid): - tracer_now = random_field(grid, CellDim, KDim) - rhodz_now = random_field(grid, CellDim, KDim) - p_mflx_tracer_v = random_field(grid, CellDim, KDim, extend={KDim: 1}) - deepatmo_divzl = random_field(grid, KDim) - deepatmo_divzu = random_field(grid, KDim) - rhodz_new = random_field(grid, CellDim, KDim) - k = as_field((KDim,), np.arange(0, _shape(grid, KDim)[0], dtype=int32)) - p_dtime = np.float64(5.0) - ivadv_tracer = int32(1) - iadv_slev_jt = int32(4) - tracer_new = zero_field(grid, CellDim, KDim) - return dict( - tracer_now=tracer_now, - rhodz_now=rhodz_now, - p_mflx_tracer_v=p_mflx_tracer_v, - deepatmo_divzl=deepatmo_divzl, - deepatmo_divzu=deepatmo_divzu, - rhodz_new=rhodz_new, - k=k, - p_dtime=p_dtime, - ivadv_tracer=ivadv_tracer, - iadv_slev_jt=iadv_slev_jt, - tracer_new=tracer_new, - ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_01.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_01.py deleted file mode 100644 index f6fc3c2ff8..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_01.py +++ /dev/null @@ -1,39 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.v_limit_prbl_sm_stencil_01 import v_limit_prbl_sm_stencil_01 -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - - -class TestVLimitPrblSmStencil01(StencilTest): - PROGRAM = v_limit_prbl_sm_stencil_01 - OUTPUTS = ("l_limit",) - - @staticmethod - def reference(grid, p_face: np.array, p_cc: np.array, **kwargs): - z_delta = p_face[:, :-1] - p_face[:, 1:] - z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:, :-1] + p_face[:, 1:])) - l_limit = np.where(np.abs(z_delta) < -1 * z_a6i, int32(1), int32(0)) - return dict(l_limit=l_limit) - - @pytest.fixture - def input_data(self, grid): - p_cc = random_field(grid, CellDim, KDim) - p_face = random_field(grid, CellDim, KDim, extend={KDim: 1}) - l_limit = zero_field(grid, CellDim, KDim, dtype=int32) - return dict(p_face=p_face, p_cc=p_cc, l_limit=l_limit) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_02.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_02.py deleted file mode 100644 index c79595b338..0000000000 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_vlimit_prbl_sm_stencil_02.py +++ /dev/null @@ -1,58 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.advection.v_limit_prbl_sm_stencil_02 import v_limit_prbl_sm_stencil_02 -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - random_field, - random_mask, - zero_field, -) - - -class TestVLimitPrblSmStencil02(StencilTest): - PROGRAM = v_limit_prbl_sm_stencil_02 - OUTPUTS = ("p_face_up", "p_face_low") - - @staticmethod - def reference(grid, l_limit: np.array, p_face: np.array, p_cc: np.array, **kwargs): - q_face_up, q_face_low = np.where( - l_limit != int32(0), - np.where( - (p_cc < np.minimum(p_face[:, :-1], p_face[:, 1:])), - (p_cc, p_cc), - np.where( - p_face[:, :-1] > p_face[:, 1:], - (3.0 * p_cc - 2.0 * p_face[:, 1:], p_face[:, 1:]), - (p_face[:, :-1], 3.0 * p_cc - 2.0 * p_face[:, :-1]), - ), - ), - (p_face[:, :-1], p_face[:, 1:]), - ) - return dict(p_face_up=q_face_up, p_face_low=q_face_low) - - @pytest.fixture - def input_data(self, grid): - l_limit = random_mask(grid, CellDim, KDim, dtype=int32) - p_cc = random_field(grid, CellDim, KDim) - p_face = random_field(grid, CellDim, KDim, extend={KDim: 1}) - p_face_up = zero_field(grid, CellDim, KDim) - p_face_low = zero_field(grid, CellDim, KDim) - return dict( - l_limit=l_limit, p_cc=p_cc, p_face=p_face, p_face_up=p_face_up, p_face_low=p_face_low - )