| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| # the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| from math import cos |
| from math import radians |
| from math import sin |
| |
| import numpy |
| from nexusproto.serialization import from_shaped_array, to_shaped_array |
| |
| from sdap.processors import NexusTileProcessor |
| |
| |
| def enum(**enums): |
| return type('Enum', (), enums) |
| |
| |
| U_OR_V_ENUM = enum(U='u', V='v') |
| |
| |
| def calculate_u_component_value(direction, speed): |
| if direction is numpy.ma.masked or speed is numpy.ma.masked: |
| return numpy.ma.masked |
| |
| return speed * sin(direction) |
| |
| |
| def calculate_v_component_value(direction, speed): |
| if direction is numpy.ma.masked or speed is numpy.ma.masked: |
| return numpy.ma.masked |
| |
| return speed * cos(direction) |
| |
| |
| class WindDirSpeedToUV(NexusTileProcessor): |
| |
| def __init__(self, u_or_v, *args, **kwargs): |
| super().__init__(*args, **kwargs) |
| |
| self.u_or_v = u_or_v.lower() |
| |
| def process_nexus_tile(self, nexus_tile): |
| the_tile_type = nexus_tile.tile.WhichOneof("tile_type") |
| |
| the_tile_data = getattr(nexus_tile.tile, the_tile_type) |
| |
| wind_speed = from_shaped_array(the_tile_data.variable_data) |
| |
| wind_dir = from_shaped_array( |
| next(meta for meta in the_tile_data.meta_data if meta.name == 'wind_dir').meta_data) |
| |
| assert wind_speed.shape == wind_dir.shape |
| |
| wind_u_component = numpy.ma.empty(wind_speed.shape, dtype=float) |
| wind_v_component = numpy.ma.empty(wind_speed.shape, dtype=float) |
| wind_speed_iter = numpy.nditer(wind_speed, flags=['multi_index']) |
| while not wind_speed_iter.finished: |
| speed = wind_speed_iter[0] |
| current_index = wind_speed_iter.multi_index |
| direction = wind_dir[current_index] |
| |
| # Convert degrees to radians |
| direction = radians(direction) |
| |
| # Calculate component values |
| wind_u_component[current_index] = calculate_u_component_value(direction, speed) |
| wind_v_component[current_index] = calculate_v_component_value(direction, speed) |
| |
| wind_speed_iter.iternext() |
| |
| # Stick the original data into the meta data |
| wind_speed_meta = the_tile_data.meta_data.add() |
| wind_speed_meta.name = 'wind_speed' |
| wind_speed_meta.meta_data.CopyFrom(to_shaped_array(wind_speed)) |
| |
| # The u_or_v variable specifies which component variable is the 'data variable' for this tile |
| # Replace data with the appropriate component value and put the other component in metadata |
| if self.u_or_v == U_OR_V_ENUM.U: |
| the_tile_data.variable_data.CopyFrom(to_shaped_array(wind_u_component)) |
| wind_component_meta = the_tile_data.meta_data.add() |
| wind_component_meta.name = 'wind_v' |
| wind_component_meta.meta_data.CopyFrom(to_shaped_array(wind_v_component)) |
| elif self.u_or_v == U_OR_V_ENUM.V: |
| the_tile_data.variable_data.CopyFrom(to_shaped_array(wind_v_component)) |
| wind_component_meta = the_tile_data.meta_data.add() |
| wind_component_meta.name = 'wind_u' |
| wind_component_meta.meta_data.CopyFrom(to_shaped_array(wind_u_component)) |
| |
| yield nexus_tile |