diff --git a/src/roguewave/interpolate/nd_interp.py b/src/roguewave/interpolate/nd_interp.py index 12db621..3473092 100644 --- a/src/roguewave/interpolate/nd_interp.py +++ b/src/roguewave/interpolate/nd_interp.py @@ -1,9 +1,32 @@ +import itertools from typing import Tuple, Callable, List import numpy from roguewave.tools.math import wrapped_difference from roguewave.tools.grid import enclosing_points_1d from roguewave.interpolate.general import interpolation_weights_1d +_EARTH_RADIUS_KILOMETERS = 6371.0 + + +def _haversine_distance_kilometers( + latitude_a: numpy.ndarray, + longitude_a: numpy.ndarray, + latitude_b: numpy.ndarray, + longitude_b: numpy.ndarray, +) -> numpy.ndarray: + latitude_a_radians = numpy.radians(latitude_a) + latitude_b_radians = numpy.radians(latitude_b) + delta_latitude_radians = numpy.radians(latitude_b - latitude_a) + delta_longitude_radians = numpy.radians(longitude_b - longitude_a) + + haversine_term = ( + numpy.sin(delta_latitude_radians / 2) ** 2 + + numpy.cos(latitude_a_radians) + * numpy.cos(latitude_b_radians) + * numpy.sin(delta_longitude_radians / 2) ** 2 + ) + return 2 * _EARTH_RADIUS_KILOMETERS * numpy.arcsin(numpy.sqrt(haversine_term)) + class NdInterpolator: def __init__( @@ -17,6 +40,7 @@ def __init__( data_period=None, data_discont=None, nearest_neighbour=False, + nan_fallback_radius=0, ): self.get_data = get_data @@ -29,6 +53,7 @@ def __init__( self.data_period = data_period self.data_discont = data_discont self.nearest_neighbour = nearest_neighbour + self.nan_fallback_radius = nan_fallback_radius @property def passive_coordinate_names(self): @@ -150,9 +175,11 @@ def interpolate( number_points, indices_1d, weights_1d ) else: - return self._data_interpolator(number_points, indices_1d, weights_1d) + return self._data_interpolator( + number_points, indices_1d, weights_1d, points + ) - def _data_interpolator(self, number_points, indices_1d, weights_1d): + def _data_interpolator(self, number_points, indices_1d, weights_1d, points): # We keep a running sum of the weights, if a point is excluded because it # contains no data (NaN) the weights will no longer add up to 1 - and we @@ -184,7 +211,211 @@ def _data_interpolator(self, number_points, indices_1d, weights_1d): ) with numpy.errstate(invalid="ignore", divide="ignore"): - return numpy.where(weights_sum > 0.5, interp_val / weights_sum, numpy.nan) + primary_result = numpy.where( + weights_sum > 0.5, interp_val / weights_sum, numpy.nan + ) + + if self.nan_fallback_radius == 0 or not numpy.any(weights_sum <= 0.5): + return primary_result + + failed_point_position, resolved_value = self._radius_neighbor_fallback( + number_points, indices_1d, weights_1d, weights_sum, points + ) + if failed_point_position.size == 0: + return primary_result + + result = primary_result.copy() + result[self.output_indexing_full(failed_point_position)] = resolved_value + return result + + def _point_slice(self, output_shaped_array): + # weights_sum is constant across passive dimensions for a given point, + # so any single passive index recovers the per-point value. + indexer = [0] * self.output_ndims + indexer[self.output_index_coord_index] = slice(None) + return output_shaped_array[tuple(indexer)] + + def _radius_neighbor_fallback( + self, number_points, indices_1d, weights_1d, weights_sum, points + ): + # For a point whose primary lookup failed, inverse-distance-weight + # whichever of its radius-N neighbors (in source-grid index space) + # have valid data instead of returning NaN outright. Everything here + # is sized to the (typically tiny) failed-point subset, not to + # number_points -- for a global grid, failed points are a small + # fraction of the total, and full-size scratch buffers here would + # scale with the whole grid for no reason. + # + # The one point where this assumes the interpolation point axis is + # output axis 0 (i.e. interp_index_coord_name is the first + # data_coordinates entry) is the final scatter into result in + # _data_interpolator, which mirrors the same assumption the primary + # bilinear loop above already makes. + if ( + "latitude" not in self.interp_coord_names + or "longitude" not in self.interp_coord_names + ): + raise NotImplementedError( + "nan_fallback_radius requires 'latitude' and 'longitude' among " + "the interpolation coordinates" + ) + + coordinate_value_by_name = dict(self.data_coordinates) + latitude_values = coordinate_value_by_name["latitude"] + longitude_values = coordinate_value_by_name["longitude"] + latitude_axis_index = self.interp_coord_names.index("latitude") + longitude_axis_index = self.interp_coord_names.index("longitude") + + axis_length_by_axis = [ + len(coordinate_value_by_name[coordinate_name]) + for coordinate_name in self.interp_coord_names + ] + axis_period_by_axis = [ + self.coordinate_period(coordinate_name) + for coordinate_name in self.interp_coord_names + ] + + # Whichever of the two bilinear bracket points carries the larger + # weight is the nearest source-grid index on that axis. + bracket_argmax_per_axis = numpy.argmax(weights_1d, axis=1) + coincident_source_index_per_axis = numpy.take_along_axis( + indices_1d, bracket_argmax_per_axis[:, None, :], axis=1 + )[:, 0, :] + + # A point outside the source grid's domain gets NaN (not merely low) + # bilinear weights, and its bracket indices are meaningless clipped + # edge values -- exclude it here so out-of-domain queries still + # return NaN instead of a fabricated nearest-edge value. + point_is_in_domain = numpy.all(numpy.isfinite(weights_1d), axis=(0, 1)) + failed_point_position = numpy.flatnonzero( + (self._point_slice(weights_sum) <= 0.5) & point_is_in_domain + ) + if failed_point_position.size == 0: + return failed_point_position, None + + number_of_failed_points = failed_point_position.size + passive_shape = tuple( + int(size) + for axis, size in enumerate(self.output_shape(number_points)) + if axis != self.output_index_coord_index + ) + fallback_weight_sum = numpy.zeros(number_of_failed_points) + fallback_value_sum = numpy.zeros( + (number_of_failed_points,) + passive_shape, dtype=numpy.float64 + ) + + coincident_source_index_of_failed_points = coincident_source_index_per_axis[ + :, failed_point_position + ] + # Distances are measured from the actually-requested point, not the + # coincident source node -- for a general (non-grid-aligned) bilinear + # miss these are not the same location. + target_latitude = points["latitude"][failed_point_position] + target_longitude = points["longitude"][failed_point_position] + + # Includes the zero offset (the coincident node itself): for a + # general, non-grid-aligned bilinear miss, the coincident node can be + # valid and still fall below the weights_sum > 0.5 threshold (e.g. if + # it is the single largest of four roughly-even corner weights), and + # is then the best available candidate, not a redundant recheck. For + # today's grid-aligned use case this is a no-op: the coincident node + # always has weight exactly 0 or 1, so it is either already handled + # by the primary lookup or invalid here too. + radius = self.nan_fallback_radius + neighbor_offsets = list( + itertools.product(range(-radius, radius + 1), repeat=self.interp_ndims) + ) + + for neighbor_offset in neighbor_offsets: + neighbor_source_index_per_axis = ( + coincident_source_index_of_failed_points.copy() + ) + neighbor_within_bounds = numpy.ones(number_of_failed_points, dtype=bool) + for axis_index in range(self.interp_ndims): + neighbor_source_index_per_axis[axis_index] += neighbor_offset[ + axis_index + ] + if axis_period_by_axis[axis_index] is not None: + neighbor_source_index_per_axis[axis_index] %= axis_length_by_axis[ + axis_index + ] + else: + neighbor_within_bounds &= ( + neighbor_source_index_per_axis[axis_index] >= 0 + ) & ( + neighbor_source_index_per_axis[axis_index] + < axis_length_by_axis[axis_index] + ) + + if not numpy.any(neighbor_within_bounds): + continue + + in_bounds_local_position = numpy.flatnonzero(neighbor_within_bounds) + # Subset to the in-bounds candidates now, once, so every array + # below (neighbor_value, neighbor_value_is_valid, and the + # per-axis index lookups) is sized consistently. + neighbor_source_index_per_axis_in_bounds = neighbor_source_index_per_axis[ + :, neighbor_within_bounds + ] + neighbor_query_indices = [ + neighbor_source_index_per_axis_in_bounds[axis_index] + for axis_index in range(self.interp_ndims) + ] + neighbor_value = self.get_data( + neighbor_query_indices, self.interp_coord_dim_indices + ) + neighbor_value_is_valid = numpy.all( + ~numpy.isnan(neighbor_value), + axis=self.output_passive_coord_dim_indices, + ) + if not numpy.any(neighbor_value_is_valid): + continue + + usable_local_position = in_bounds_local_position[neighbor_value_is_valid] + + neighbor_latitude_value = latitude_values[ + neighbor_source_index_per_axis_in_bounds[latitude_axis_index][ + neighbor_value_is_valid + ] + ] + neighbor_longitude_value = longitude_values[ + neighbor_source_index_per_axis_in_bounds[longitude_axis_index][ + neighbor_value_is_valid + ] + ] + neighbor_distance_kilometers = _haversine_distance_kilometers( + target_latitude[usable_local_position], + target_longitude[usable_local_position], + neighbor_latitude_value, + neighbor_longitude_value, + ) + neighbor_inverse_distance_weight = numpy.where( + neighbor_distance_kilometers > 0, + 1.0 / neighbor_distance_kilometers, + 0.0, + ) + + usable_value = neighbor_value[neighbor_value_is_valid] + weight_broadcast_shape = (-1,) + (1,) * (usable_value.ndim - 1) + + fallback_weight_sum[ + usable_local_position + ] += neighbor_inverse_distance_weight + fallback_value_sum[usable_local_position] += ( + neighbor_inverse_distance_weight.reshape(weight_broadcast_shape) + * usable_value + ) + + weight_broadcast_shape = (-1,) + (1,) * (fallback_value_sum.ndim - 1) + with numpy.errstate(invalid="ignore", divide="ignore"): + resolved_value = numpy.where( + fallback_weight_sum.reshape(weight_broadcast_shape) > 0, + fallback_value_sum + / fallback_weight_sum.reshape(weight_broadcast_shape), + numpy.nan, + ) + + return failed_point_position, resolved_value def _periodic_data_interpolator(self, number_points, indices_1d, weights_1d): # We keep a running sum of the weights, if a point is excluded because it diff --git a/src/roguewave/wavewatch3/io.py b/src/roguewave/wavewatch3/io.py index 1d9e1c6..1f5490e 100644 --- a/src/roguewave/wavewatch3/io.py +++ b/src/roguewave/wavewatch3/io.py @@ -154,7 +154,7 @@ def write_restart_file( :return: None """ if isinstance(spectra, (Dataset, Spectrum)): - spectra = spectra.variance_density.values + spectra = spectra.directional_variance_density.values elif isinstance(spectra, DataArray): spectra = spectra.values diff --git a/src/roguewave/wavewatch3/restart_file.py b/src/roguewave/wavewatch3/restart_file.py index ba35ea9..4c1e7bb 100644 --- a/src/roguewave/wavewatch3/restart_file.py +++ b/src/roguewave/wavewatch3/restart_file.py @@ -36,9 +36,77 @@ from datetime import datetime from functools import cache from roguewave.tools.time import to_datetime64 +from roguewave.tools.grid import midpoint_rule_step from xarray import Dataset, DataArray from roguewavespectrum import Spectrum +_GRAVITATIONAL_ACCELERATION = 9.81 + + +def _cos_power_directional_density( + direction_degrees: numpy.ndarray, + mean_direction_degrees: float, + spreading_power: float, +) -> numpy.ndarray: + """ + cos^spreading_power(direction - mean_direction) for |direction - mean_direction| + <= 90 degrees, zero beyond -- normalized so the density sums to 1 over the + (assumed uniform) direction bins, matching ww3_strt ITYPE 1's directional shape. + """ + signed_angle_difference = ( + (direction_degrees - mean_direction_degrees + 180.0) % 360.0 + ) - 180.0 + raw_density = numpy.where( + numpy.abs(signed_angle_difference) <= 90.0, + numpy.cos(numpy.radians(signed_angle_difference)) ** spreading_power, + 0.0, + ) + direction_bin_width_degrees = 360.0 / len(direction_degrees) + return raw_density / (numpy.sum(raw_density) * direction_bin_width_degrees) + + +def _gaussian_frequency_density( + frequency: numpy.ndarray, peak_frequency: float, frequency_spread: float +) -> numpy.ndarray: + """ + Gaussian bump in frequency around peak_frequency, normalized so the density + integrates to 1 over the (possibly non-uniform) frequency bins, matching + ww3_strt ITYPE 1's frequency shape. + """ + raw_density = numpy.exp( + -((frequency - peak_frequency) ** 2) / (2 * frequency_spread**2) + ) + frequency_bin_width = midpoint_rule_step(frequency) + return raw_density / numpy.sum(raw_density * frequency_bin_width) + + +def _jonswap_frequency_density( + frequency: numpy.ndarray, + peak_frequency: float, + alpha: float, + gamma: float, + sigma_a: float, + sigma_b: float, +) -> numpy.ndarray: + """ + Standard five-parameter JONSWAP spectrum, matching ww3_strt ITYPE 2's + ALFA/FP/GAMMA/SIGA/SIGB parameterization. Unlike the Gaussian shape above, + alpha sets the absolute energy level directly -- this is not renormalized. + """ + sigma = numpy.where(frequency <= peak_frequency, sigma_a, sigma_b) + peak_enhancement = gamma ** numpy.exp( + -((frequency - peak_frequency) ** 2) / (2 * sigma**2 * peak_frequency**2) + ) + pierson_moskowitz = ( + alpha + * _GRAVITATIONAL_ACCELERATION**2 + * (2 * numpy.pi) ** -4 + * frequency**-5 + * numpy.exp(-1.25 * (peak_frequency / frequency) ** 4) + ) + return pierson_moskowitz * peak_enhancement + + MAXIMUM_NUMBER_OF_WORKERS = 10 @@ -278,6 +346,15 @@ def _fancy_index(self, indices: Union[Sequence, numpy.ndarray]) -> numpy.ndarray if isinstance(indices, Sequence): indices = numpy.array(indices, dtype="int32") + if len(indices) == 0: + # numpy.array([]) on an empty list of spectra collapses to shape + # (0,) instead of (0, number_of_frequencies, number_of_directions), + # which xarray then rejects as a dimension-size conflict. + return numpy.empty( + (0, self.number_of_frequencies, self.number_of_directions), + dtype=self._dtype, + ) + indices = indices + self._start_record slices = [ slice(self._byte_index(index), self._byte_index(index + 1), 1) @@ -326,6 +403,14 @@ def interpolate_in_space( Input can be either a single latitude and longitude pair, or a numpy array of latitudes and longitudes. + When a target point's source grid point is masked/land, this falls + back to an inverse-distance-weighted blend of its radius-1 neighbors + (see NdInterpolator's nan_fallback_radius) rather than returning NaN + outright. A point can still come back NaN if none of its radius-1 + neighbors have data either -- a genuine domain gap, not a coastal + artifact. See fill_missing_spectra to replace those with a WW3 + cold-start-style spectrum instead of leaving them as NaN. + :param latitude: latitudes to get interpolated spectra :param longitude: longitudes to get interpolated spectra :return: Interpolated spectra. Returned data is of type float32 and @@ -378,6 +463,7 @@ def _get_data(indices, _dummy): data_periodic_coordinates=periodic_coordinates, data_period=None, data_discont=None, + nan_fallback_radius=1, ) def _get_depth(indices, _dummy): @@ -422,6 +508,99 @@ def _get_depth(indices, _dummy): ) ) + def fill_missing_spectra( + self, spectrum: Spectrum, fill_type: str = "calm", **fill_parameters + ) -> Spectrum: + """ + Fill points in `spectrum` (as returned by interpolate_in_space) that + are still entirely NaN -- e.g. a domain gap with no nearby source data + at all -- with a WW3 cold-start-style parametric spectrum, rather than + leaving them as NaN. Points that already have data are left untouched. + + fill_type options, matching ww3_strt's cold-start initial-condition + types (these are roguewave's own implementations of the same named + spectral shapes, not a port of WW3's Fortran): + + - "calm" (default, matches ITYPE 5): zero energy everywhere. + - "user_defined" (matches ITYPE 4): broadcast a caller-supplied + spectrum to every missing point. Parameter: spectral_values, an + array shaped (number_of_frequencies, number_of_directions). + - "gaussian" (matches ITYPE 1): Gaussian in frequency, cos-power in + direction, normalized to a target significant wave height. + Parameters: peak_frequency, frequency_spread, mean_direction, + directional_spreading_power, significant_wave_height. + - "jonswap" (matches ITYPE 2): standard five-parameter JONSWAP, cos- + power in direction. Parameters: alpha, peak_frequency, gamma, + sigma_a, sigma_b, mean_direction, directional_spreading_power. + + :param spectrum: a Spectrum shaped (points, frequency, direction), + as returned by interpolate_in_space. + :param fill_type: one of "calm", "user_defined", "gaussian", "jonswap". + :param fill_parameters: parameters for the chosen fill_type, see above. + :return: a copy of spectrum with missing points filled. + """ + filled_spectrum = spectrum.copy() + variable = filled_spectrum.directional_variance_density + values = variable.values.copy() + + point_is_missing = numpy.all(numpy.isnan(values), axis=(1, 2)) + if numpy.any(point_is_missing): + values[point_is_missing, :, :] = self._fill_spectrum( + fill_type, **fill_parameters + ) + filled_spectrum.dataset[variable.name] = (variable.dims, values) + + return filled_spectrum + + def _fill_spectrum(self, fill_type: str, **fill_parameters) -> numpy.ndarray: + if fill_type == "calm": + return numpy.zeros((self.number_of_frequencies, self.number_of_directions)) + + if fill_type == "user_defined": + spectral_values = numpy.asarray(fill_parameters["spectral_values"]) + expected_shape = (self.number_of_frequencies, self.number_of_directions) + if spectral_values.shape != expected_shape: + raise ValueError( + f"spectral_values must have shape {expected_shape}, got " + f"{spectral_values.shape}" + ) + return spectral_values + + if fill_type in ("gaussian", "jonswap"): + directional_density = _cos_power_directional_density( + self.direction, + fill_parameters["mean_direction"], + fill_parameters["directional_spreading_power"], + ) + + if fill_type == "gaussian": + frequency_density = _gaussian_frequency_density( + self.frequency, + fill_parameters["peak_frequency"], + fill_parameters["frequency_spread"], + ) + target_variance = ( + fill_parameters["significant_wave_height"] / 4.0 + ) ** 2 + return target_variance * numpy.outer( + frequency_density, directional_density + ) + else: + frequency_density = _jonswap_frequency_density( + self.frequency, + fill_parameters["peak_frequency"], + fill_parameters["alpha"], + fill_parameters["gamma"], + fill_parameters["sigma_a"], + fill_parameters["sigma_b"], + ) + return numpy.outer(frequency_density, directional_density) + + raise ValueError( + f"Unknown fill_type '{fill_type}'; expected one of " + "'calm', 'user_defined', 'gaussian', 'jonswap'" + ) + def to_wavenumber_action_density( self, s: Union[slice, NDArray, Sequence] ) -> numpy.array: diff --git a/tests/interpolate/__init__.py b/tests/interpolate/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/interpolate/test_nd_interp.py b/tests/interpolate/test_nd_interp.py new file mode 100644 index 0000000..1423236 --- /dev/null +++ b/tests/interpolate/test_nd_interp.py @@ -0,0 +1,132 @@ +import numpy +from roguewave.interpolate.nd_interp import NdInterpolator + + +def _make_get_data(grid_value, grid_mask): + def get_data(indices, _dimension_indices): + latitude_index, longitude_index = indices[0], indices[1] + value = grid_value[latitude_index, longitude_index].astype(float).copy() + value[~grid_mask[latitude_index, longitude_index]] = numpy.nan + return value + + return get_data + + +def _make_interpolator( + grid_value, grid_mask, nan_fallback_radius, periodic_longitude=None +): + latitude_values = numpy.array([-1.0, 0.0, 1.0]) + longitude_values = numpy.array([-1.0, 0.0, 1.0]) + periodic_coordinates = ( + {"longitude": periodic_longitude} if periodic_longitude is not None else {} + ) + + return NdInterpolator( + get_data=_make_get_data(grid_value, grid_mask), + data_coordinates=( + ("latitude", latitude_values), + ("longitude", longitude_values), + ), + data_shape=[3, 3], + interp_coord_names=["latitude", "longitude"], + interp_index_coord_name="latitude", + data_periodic_coordinates=periodic_coordinates, + nan_fallback_radius=nan_fallback_radius, + ) + + +# Grid layout (row=latitude index, column=longitude index), matching +# latitude_values=[-1, 0, 1] and longitude_values=[-1, 0, 1]: +# SW S SE 10 20 30 +# W C E = 40 100 50 +# NW N NE 60 70 80 +GRID_VALUE = numpy.array([[10.0, 20.0, 30.0], [40.0, 100.0, 50.0], [60.0, 70.0, 80.0]]) + + +def test_coincident_point_valid_ignores_fallback(): + grid_mask = numpy.ones((3, 3), dtype=bool) + interpolator = _make_interpolator(GRID_VALUE, grid_mask, nan_fallback_radius=1) + + result = interpolator.interpolate( + {"latitude": numpy.array([0.0]), "longitude": numpy.array([0.0])} + ) + + assert result[0] == 100.0 + + +def test_masked_coincident_point_blends_valid_radius_one_neighbors(): + grid_mask = numpy.zeros((3, 3), dtype=bool) + grid_mask[0, 1] = True # S, value 20 + grid_mask[2, 1] = True # N, value 70, equidistant from the target as S + + interpolator = _make_interpolator(GRID_VALUE, grid_mask, nan_fallback_radius=1) + + result = interpolator.interpolate( + {"latitude": numpy.array([0.0]), "longitude": numpy.array([0.0])} + ) + + assert numpy.abs(result[0] - 45.0) < 1e-6 + + +def test_masked_coincident_point_with_no_valid_neighbors_stays_nan(): + grid_mask = numpy.zeros( + (3, 3), dtype=bool + ) # every point masked, including the target + + interpolator = _make_interpolator(GRID_VALUE, grid_mask, nan_fallback_radius=1) + + result = interpolator.interpolate( + {"latitude": numpy.array([0.0]), "longitude": numpy.array([0.0])} + ) + + assert numpy.isnan(result[0]) + + +def test_default_radius_zero_preserves_current_nan_behavior(): + grid_mask = numpy.zeros((3, 3), dtype=bool) + grid_mask[0, 1] = True + grid_mask[2, 1] = True + + interpolator = _make_interpolator(GRID_VALUE, grid_mask, nan_fallback_radius=0) + + result = interpolator.interpolate( + {"latitude": numpy.array([0.0]), "longitude": numpy.array([0.0])} + ) + + assert numpy.isnan(result[0]) + + +def test_fallback_wraps_around_periodic_longitude(): + # Target is at index 0, an array endpoint: its "west" neighbor (offset + # -1) only exists by wrapping around to index 2 via the modulo in + # _radius_neighbor_fallback. A masked target in the middle of the array + # would never exercise that wrap, since both its immediate neighbors are + # already in-bounds without it. + latitude_values = numpy.array([0.0]) + longitude_values = numpy.array([0.0, 1.0, 2.0]) + grid_value = numpy.array([[90.0, 100.0, 10.0]]) + # index 0 (target) masked; index 1 (offset +1, no wrap needed) also + # masked so the only valid neighbor is index 2, reached solely via the + # offset -1 -> -1 % 3 == 2 wraparound. + grid_mask = numpy.array([[False, False, True]]) + + interpolator = NdInterpolator( + get_data=_make_get_data(grid_value, grid_mask), + data_coordinates=( + ("latitude", latitude_values), + ("longitude", longitude_values), + ), + data_shape=[1, 3], + interp_coord_names=["latitude", "longitude"], + interp_index_coord_name="latitude", + data_periodic_coordinates={"longitude": 3.0}, + nan_fallback_radius=1, + ) + + result = interpolator.interpolate( + {"latitude": numpy.array([0.0]), "longitude": numpy.array([0.0])} + ) + + # The only valid neighbor is index 2, reachable only by wrapping; with a + # single valid neighbor the IDW blend degenerates to that exact value. + assert numpy.abs(result[0] - 10.0) < 1e-6 diff --git a/tests/restart_files/test_empty_index.py b/tests/restart_files/test_empty_index.py new file mode 100644 index 0000000..f3284eb --- /dev/null +++ b/tests/restart_files/test_empty_index.py @@ -0,0 +1,66 @@ +from datetime import datetime, timezone +import numpy +from roguewave.wavewatch3.grid_tools import Grid +from roguewave.wavewatch3.restart_file import RestartFile +from roguewave.wavewatch3.restart_file_metadata import MetaData + + +def _make_minimal_restart_file(): + """ + A fully synthetic, in-memory RestartFile: Grid and MetaData are plain + dataclasses with no file I/O, and the empty-index path under test never + touches `resource`, so this needs no real restart file on disk or S3. + """ + number_of_frequencies = 2 + number_of_directions = 4 + frequencies = numpy.array([0.1, 0.2]) + directions = numpy.array([0.0, 90.0, 180.0, 270.0]) + latitude = numpy.array([0.0, 1.0]) + longitude = numpy.array([0.0, 1.0]) + + # A 2x2 lat/lon grid with a single sea point at (latitude[0], longitude[0]). + to_linear_index = numpy.array([[0, -1], [-1, -1]]) + to_point_index = numpy.array([[0], [0]]) # [ilon, ilat] for linear index 0 + + grid = Grid( + number_of_spatial_points=1, + frequencies=frequencies, + directions=directions, + latitude=latitude, + longitude=longitude, + depth=numpy.array([100.0]), + mask=numpy.array([[1, 0], [0, 0]]), + _to_linear_index=to_linear_index, + _to_point_index=to_point_index, + ) + meta_data = MetaData( + name="test", + version="1", + grid_name="test", + restart_type="test", + nsea=1, + nspec=number_of_frequencies * number_of_directions, + record_size_bytes=number_of_frequencies * number_of_directions * 4, + time=datetime(2020, 1, 1, tzinfo=timezone.utc), + byte_order="<", + float_size=4, + ) + return RestartFile(grid=grid, meta_data=meta_data, resource=None) + + +def test_getitem_with_empty_fancy_index_does_not_crash(): + restart_file = _make_minimal_restart_file() + + spectra = restart_file[numpy.array([], dtype="int32")] + + assert spectra.number_of_spectra == 0 + values = spectra.directional_variance_density.values + assert values.shape == (0, 2, 4) + + +def test_fancy_index_with_empty_indices_has_correct_shape(): + restart_file = _make_minimal_restart_file() + + result = restart_file._fancy_index(numpy.array([], dtype="int32")) + + assert result.shape == (0, 2, 4) diff --git a/tests/restart_files/test_fill_missing_spectra.py b/tests/restart_files/test_fill_missing_spectra.py new file mode 100644 index 0000000..53ef369 --- /dev/null +++ b/tests/restart_files/test_fill_missing_spectra.py @@ -0,0 +1,181 @@ +import numpy +from xarray import Dataset +from roguewavespectrum import Spectrum +from roguewave.tools.grid import midpoint_rule_step +from roguewave.wavewatch3.restart_file import RestartFile + + +class _FakeRestartFile: + """ + A minimal duck-typed stand-in for RestartFile, exposing only what + fill_missing_spectra/_fill_spectrum actually use, so these can be + tested without constructing a real, binary-file-backed RestartFile. + """ + + fill_missing_spectra = RestartFile.fill_missing_spectra + _fill_spectrum = RestartFile._fill_spectrum + + def __init__(self, frequency, direction): + self.frequency = frequency + self.direction = direction + + @property + def number_of_frequencies(self): + return len(self.frequency) + + @property + def number_of_directions(self): + return len(self.direction) + + +def _make_spectrum(frequency, direction, values): + number_of_points = values.shape[0] + return Spectrum( + Dataset( + data_vars={ + "directional_variance_density": ( + ("points", "frequency", "direction"), + values, + ), + "longitude": (("points",), numpy.zeros(number_of_points)), + "latitude": (("points",), numpy.zeros(number_of_points)), + "depth": (("points",), numpy.full(number_of_points, 100.0)), + }, + coords={"frequency": frequency, "direction": direction}, + ) + ) + + +def test_fill_spectrum_calm_is_zero(): + fake_restart_file = _FakeRestartFile( + numpy.array([0.1, 0.2]), numpy.array([0.0, 90.0, 180.0, 270.0]) + ) + fill = fake_restart_file._fill_spectrum("calm") + + assert fill.shape == (2, 4) + assert numpy.all(fill == 0.0) + + +def test_fill_spectrum_user_defined_passes_through(): + fake_restart_file = _FakeRestartFile( + numpy.array([0.1, 0.2]), numpy.array([0.0, 90.0, 180.0, 270.0]) + ) + supplied_spectral_values = numpy.arange(8, dtype=float).reshape(2, 4) + + fill = fake_restart_file._fill_spectrum( + "user_defined", spectral_values=supplied_spectral_values + ) + + assert numpy.array_equal(fill, supplied_spectral_values) + + +def test_fill_spectrum_user_defined_wrong_shape_raises(): + fake_restart_file = _FakeRestartFile( + numpy.array([0.1, 0.2]), numpy.array([0.0, 90.0, 180.0, 270.0]) + ) + + try: + fake_restart_file._fill_spectrum( + "user_defined", spectral_values=numpy.zeros((3, 3)) + ) + assert False, "expected a ValueError for a mismatched spectral_values shape" + except ValueError: + pass + + +def test_fill_spectrum_unknown_fill_type_raises(): + fake_restart_file = _FakeRestartFile( + numpy.array([0.1, 0.2]), numpy.array([0.0, 90.0, 180.0, 270.0]) + ) + + try: + fake_restart_file._fill_spectrum("not_a_real_fill_type") + assert False, "expected a ValueError for an unknown fill_type" + except ValueError: + pass + + +def test_fill_spectrum_gaussian_matches_target_significant_wave_height(): + frequency = numpy.geomspace(0.035, 0.5, 30) + direction = numpy.linspace(0, 360, 36, endpoint=False) + fake_restart_file = _FakeRestartFile(frequency, direction) + + target_significant_wave_height = 2.0 + fill = fake_restart_file._fill_spectrum( + "gaussian", + peak_frequency=0.1, + frequency_spread=0.01, + mean_direction=90.0, + directional_spreading_power=4, + significant_wave_height=target_significant_wave_height, + ) + + frequency_bin_width = midpoint_rule_step(frequency) + direction_bin_width = 360.0 / len(direction) + total_variance = ( + numpy.sum(fill * frequency_bin_width[:, None]) * direction_bin_width + ) + resulting_significant_wave_height = 4.0 * numpy.sqrt(total_variance) + + assert ( + numpy.abs(resulting_significant_wave_height - target_significant_wave_height) + < 1e-3 + ) + assert numpy.argmax(numpy.sum(fill, axis=1)) == numpy.argmin( + numpy.abs(frequency - 0.1) + ) + + +def test_fill_spectrum_jonswap_peaks_near_peak_frequency(): + frequency = numpy.geomspace(0.035, 0.5, 60) + direction = numpy.linspace(0, 360, 36, endpoint=False) + fake_restart_file = _FakeRestartFile(frequency, direction) + + fill = fake_restart_file._fill_spectrum( + "jonswap", + peak_frequency=0.1, + alpha=0.01, + gamma=3.3, + sigma_a=0.07, + sigma_b=0.09, + mean_direction=180.0, + directional_spreading_power=2, + ) + + assert numpy.all(fill >= 0.0) + assert numpy.argmax(numpy.sum(fill, axis=1)) == numpy.argmin( + numpy.abs(frequency - 0.1) + ) + + +def test_fill_missing_spectra_leaves_valid_points_untouched_and_fills_missing(): + frequency = numpy.array([0.1, 0.2]) + direction = numpy.array([0.0, 90.0, 180.0, 270.0]) + + valid_values = numpy.ones((2, 4)) * 5.0 + values = numpy.stack([valid_values, numpy.full((2, 4), numpy.nan)], axis=0) + spectrum = _make_spectrum(frequency, direction, values) + fake_restart_file = _FakeRestartFile(frequency, direction) + + filled_spectrum = fake_restart_file.fill_missing_spectra(spectrum, fill_type="calm") + filled_values = filled_spectrum.directional_variance_density.values + + assert numpy.array_equal(filled_values[0], valid_values) + assert numpy.all(filled_values[1] == 0.0) + # the input spectrum passed in is not mutated + assert numpy.all(numpy.isnan(spectrum.directional_variance_density.values[1])) + + +def test_fill_missing_spectra_is_a_no_op_when_nothing_is_missing(): + frequency = numpy.array([0.1, 0.2]) + direction = numpy.array([0.0, 90.0, 180.0, 270.0]) + + valid_values = numpy.ones((1, 2, 4)) * 5.0 + spectrum = _make_spectrum(frequency, direction, valid_values) + fake_restart_file = _FakeRestartFile(frequency, direction) + + filled_spectrum = fake_restart_file.fill_missing_spectra(spectrum, fill_type="calm") + + assert numpy.array_equal( + filled_spectrum.directional_variance_density.values, valid_values + ) diff --git a/tests/restart_files/test_io.py b/tests/restart_files/test_io.py index 15e2848..ece1bed 100644 --- a/tests/restart_files/test_io.py +++ b/tests/restart_files/test_io.py @@ -1,6 +1,8 @@ from roguewave.wavewatch3.io import ( write_partial_restart_file, reassemble_restart_file_from_parts, + clone_restart_file, + open_restart_file, ) import os @@ -41,6 +43,25 @@ def test_local_partial_write(): ) +def test_clone_restart_file(): + # Regression test for #15: write_restart_file used to raise IndexError + # when given a Spectrum/Dataset (e.g. via clone_restart_file's use of + # RestartFile.__getitem__), because it read .variance_density instead of + # .directional_variance_density. + restart_file = clone_remote() + local_file = os.path.join(TEST_DIR, LOCAL_FILE_NAME) + model_definition_file = os.path.join(TEST_DIR, "mod_def.ww3") + output = "cloned_restart_test.file" + + clone_restart_file(local_file, model_definition_file, output) + + try: + cloned = open_restart_file(output, model_definition_file) + assert cloned.number_of_spatial_points == restart_file.number_of_spatial_points + finally: + os.remove(output) + + def test_local_reassemble(): restart_file = clone_remote() names = []