wehub-resource-sync
169 lines
6.5 KiB
Python
169 lines
6.5 KiB
Python
# Copyright 2025 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Tests for NaN-handling and interpolation utilities in the base module.
|
|
|
|
``strip_leading_nans`` and ``linear_interpolation`` sit on the critical
|
|
inference path: every user input passes through them before being
|
|
patched and fed to the transformer. Incorrect behavior here — silently
|
|
keeping NaN values or interpolating the wrong indices — causes NaN
|
|
propagation through the entire model and produces garbage forecasts.
|
|
"""
|
|
|
|
import numpy as np
|
|
|
|
from timesfm.timesfm_2p5.timesfm_2p5_base import (
|
|
linear_interpolation,
|
|
strip_leading_nans,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# strip_leading_nans
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestStripLeadingNans:
|
|
"""Tests for strip_leading_nans — removes leading NaN prefix."""
|
|
|
|
def test_no_nans_returns_unchanged(self):
|
|
"""An array without NaN values must pass through unmodified."""
|
|
arr = np.array([1.0, 2.0, 3.0])
|
|
result = strip_leading_nans(arr)
|
|
np.testing.assert_array_equal(result, arr)
|
|
|
|
def test_strips_leading_nans_only(self):
|
|
"""Leading NaNs are removed; NaNs embedded in the middle are kept."""
|
|
arr = np.array([np.nan, np.nan, 1.0, np.nan, 3.0])
|
|
result = strip_leading_nans(arr)
|
|
expected = np.array([1.0, np.nan, 3.0])
|
|
np.testing.assert_array_equal(result, expected)
|
|
|
|
def test_single_leading_nan(self):
|
|
"""Edge case: exactly one leading NaN."""
|
|
arr = np.array([np.nan, 5.0, 6.0])
|
|
result = strip_leading_nans(arr)
|
|
np.testing.assert_array_equal(result, np.array([5.0, 6.0]))
|
|
|
|
def test_no_leading_nan_with_internal_nans(self):
|
|
"""If the first element is valid, nothing is stripped regardless of
|
|
internal NaNs."""
|
|
arr = np.array([1.0, np.nan, np.nan, 4.0])
|
|
result = strip_leading_nans(arr)
|
|
np.testing.assert_array_equal(result, arr)
|
|
|
|
def test_single_valid_element(self):
|
|
"""A single non-NaN element must be returned as-is."""
|
|
arr = np.array([42.0])
|
|
result = strip_leading_nans(arr)
|
|
np.testing.assert_array_equal(result, np.array([42.0]))
|
|
|
|
def test_all_nans_returns_full_array(self):
|
|
"""When every element is NaN, ``np.argmax`` on an all-False mask
|
|
returns 0 — so the implementation returns the original array, not an
|
|
empty one.
|
|
|
|
This documents the *actual* behavior (which differs from the
|
|
docstring claim of returning an empty array). Downstream code
|
|
(``linear_interpolation``) is designed to handle this case.
|
|
"""
|
|
arr = np.array([np.nan, np.nan, np.nan])
|
|
result = strip_leading_nans(arr)
|
|
# Actual behavior: argmax(~isnan) = 0 when all NaN → returns full array.
|
|
assert len(result) == 3
|
|
assert np.all(np.isnan(result))
|
|
|
|
def test_preserves_dtype(self):
|
|
"""Output dtype must match input dtype (float32 stays float32)."""
|
|
arr = np.array([np.nan, 1.0, 2.0], dtype=np.float32)
|
|
result = strip_leading_nans(arr)
|
|
assert result.dtype == np.float32
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# linear_interpolation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestLinearInterpolation:
|
|
"""Tests for linear_interpolation — fills NaN gaps via ``np.interp``."""
|
|
|
|
def test_no_nans_returns_identical(self):
|
|
"""Without NaN values the array is returned as-is (fast path)."""
|
|
arr = np.array([1.0, 2.0, 3.0])
|
|
result = linear_interpolation(arr.copy())
|
|
np.testing.assert_array_equal(result, arr)
|
|
|
|
def test_interpolates_single_interior_nan(self):
|
|
"""A single interior NaN is linearly interpolated from neighbors."""
|
|
arr = np.array([0.0, np.nan, 2.0])
|
|
result = linear_interpolation(arr)
|
|
np.testing.assert_allclose(result, [0.0, 1.0, 2.0])
|
|
|
|
def test_interpolates_multiple_interior_nans(self):
|
|
"""Multiple consecutive interior NaN values are interpolated."""
|
|
arr = np.array([0.0, np.nan, np.nan, 3.0])
|
|
result = linear_interpolation(arr)
|
|
np.testing.assert_allclose(result, [0.0, 1.0, 2.0, 3.0])
|
|
|
|
def test_extrapolates_trailing_nans(self):
|
|
"""Trailing NaN values are filled via ``np.interp`` which holds the
|
|
last known value (nearest-neighbor extrapolation)."""
|
|
arr = np.array([1.0, 2.0, np.nan, np.nan])
|
|
result = linear_interpolation(arr)
|
|
# np.interp extrapolates by clamping to boundary values.
|
|
np.testing.assert_allclose(result, [1.0, 2.0, 2.0, 2.0])
|
|
|
|
def test_extrapolates_leading_nans(self):
|
|
"""Leading NaN values are filled with the first valid value.
|
|
|
|
In practice ``strip_leading_nans`` runs first, but the function must
|
|
still be robust on its own.
|
|
"""
|
|
arr = np.array([np.nan, np.nan, 3.0, 4.0])
|
|
result = linear_interpolation(arr)
|
|
np.testing.assert_allclose(result, [3.0, 3.0, 3.0, 4.0])
|
|
|
|
def test_output_has_no_nans(self):
|
|
"""After interpolation, no NaN values should remain."""
|
|
arr = np.array([np.nan, 1.0, np.nan, np.nan, 4.0, np.nan])
|
|
result = linear_interpolation(arr)
|
|
assert not np.any(np.isnan(result))
|
|
|
|
def test_preserves_non_nan_values(self):
|
|
"""Non-NaN values in the original array must never be modified."""
|
|
arr = np.array([10.0, np.nan, 30.0, np.nan, 50.0])
|
|
original_valid = arr[~np.isnan(arr)].copy()
|
|
result = linear_interpolation(arr)
|
|
np.testing.assert_array_equal(
|
|
result[~np.isnan(np.array([10.0, np.nan, 30.0, np.nan, 50.0]))],
|
|
original_valid,
|
|
)
|
|
|
|
def test_interpolation_is_monotone_for_monotone_input(self):
|
|
"""If the known values are strictly increasing, the interpolated
|
|
result must also be non-decreasing — a basic sanity check on the
|
|
interpolation direction."""
|
|
arr = np.array([1.0, np.nan, np.nan, 4.0, np.nan, 6.0])
|
|
result = linear_interpolation(arr)
|
|
diffs = np.diff(result)
|
|
assert np.all(diffs >= 0)
|
|
|
|
def test_single_non_nan_fills_all_gaps(self):
|
|
"""With only one valid value, every NaN is replaced by that value
|
|
(np.interp clamps to the single known point)."""
|
|
arr = np.array([np.nan, 5.0, np.nan])
|
|
result = linear_interpolation(arr)
|
|
np.testing.assert_allclose(result, [5.0, 5.0, 5.0])
|