| # 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. |
| |
| |
| import numpy |
| from nexusproto.serialization import from_shaped_array, to_shaped_array |
| |
| from sdap.processors import NexusTileProcessor |
| |
| |
| def calculate_speed_direction(wind_u, wind_v): |
| speed = numpy.sqrt(numpy.add(numpy.multiply(wind_u, wind_u), numpy.multiply(wind_v, wind_v))) |
| direction = numpy.degrees(numpy.arctan2(-wind_u, -wind_v)) % 360 |
| return speed, direction |
| |
| |
| class ComputeSpeedDirFromUV(NexusTileProcessor): |
| |
| def __init__(self, wind_u_var_name, wind_v_var_name, *args, **kwargs): |
| super().__init__(*args, **kwargs) |
| |
| self.wind_u_var_name = wind_u_var_name |
| self.wind_v_var_name = wind_v_var_name |
| |
| 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) |
| |
| # Either wind_u or wind_v are in meta. Whichever is not in meta is in variable_data |
| try: |
| wind_v = next(meta for meta in the_tile_data.meta_data if meta.name == self.wind_v_var_name).meta_data |
| wind_u = the_tile_data.variable_data |
| except StopIteration: |
| try: |
| wind_u = next(meta for meta in the_tile_data.meta_data if meta.name == self.wind_u_var_name).meta_data |
| wind_v = the_tile_data.variable_data |
| except StopIteration: |
| if hasattr(nexus_tile, "summary"): |
| raise RuntimeError( |
| "Neither wind_u nor wind_v were found in the meta data for granule %s slice %s." |
| " Cannot compute wind speed or direction." % ( |
| getattr(nexus_tile.summary, "granule", "unknown"), |
| getattr(nexus_tile.summary, "section_spec", "unknown"))) |
| else: |
| raise RuntimeError( |
| "Neither wind_u nor wind_v were found in the meta data. Cannot compute wind speed or direction.") |
| |
| wind_u = from_shaped_array(wind_u) |
| wind_v = from_shaped_array(wind_v) |
| |
| assert wind_u.shape == wind_v.shape |
| |
| # Do calculation |
| wind_speed_data, wind_dir_data = calculate_speed_direction(wind_u, wind_v) |
| |
| # Add wind_speed to 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_data)) |
| |
| # Add wind_dir to meta data |
| wind_dir_meta = the_tile_data.meta_data.add() |
| wind_dir_meta.name = 'wind_dir' |
| wind_dir_meta.meta_data.CopyFrom(to_shaped_array(wind_dir_data)) |
| |
| yield nexus_tile |