#!/usr/bin/env python3
# encoding: utf-8
# Copyright (C) 2021 Space Science and Engineering Center (SSEC),
# University of Wisconsin-Madison.
#
# This program 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
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
# This file is part of the polar2grid software package. Polar2grid takes
# satellite observation data, remaps it, and writes it to a file format for
# input into another program.
# Documentation: http://www.ssec.wisc.edu/software/polar2grid/
"""Tests for the binary writer."""
import os
from datetime import datetime
from unittest import mock
import dask.array as da
import numpy as np
import pytest
import satpy
import xarray as xr
from satpy import Scene
TEST_ETC_DIR = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "etc"))
[docs]def _int_min_max(dtype):
info = np.iinfo(dtype)
return info.min, info.max
[docs]def _min_max_for_two_dtypes(src_dtype, dst_dtype):
src_is_int = not np.issubdtype(src_dtype, np.floating)
dst_is_int = not np.issubdtype(dst_dtype, np.floating)
if src_is_int and dst_is_int:
smin, smax = _int_min_max(src_dtype)
dmin, dmax = _int_min_max(dst_dtype)
return max(smin, dmin), min(smax, dmax)
elif src_is_int:
return _int_min_max(src_dtype)
return _int_min_max(dst_dtype)
[docs]def _create_fake_data_arr(shape=(100, 50), dims=("y", "x"), dtype=np.float64):
area_mock = mock.MagicMock()
area_mock.area_id = "fake_area"
attrs = {
"name": "fake_name",
"p2g_name": "fake_p2g_name",
"platform_name": "NOAA-20",
"sensor": "viirs",
"start_time": datetime(2021, 1, 1, 12, 0, 0),
"end_time": datetime(2021, 1, 1, 12, 10, 0),
"area": area_mock,
"standard_name": "test_arange",
}
data = np.arange(np.prod(shape), dtype=np.float64)
if not np.issubdtype(dtype, np.floating):
data = np.clip(data, *_int_min_max(dtype))
attrs["_FillValue"] = 0
data = data.astype(dtype)
data = da.from_array(data, chunks=10).reshape(shape)
data_arr = xr.DataArray(data, attrs=attrs, dims=dims)
return data_arr
[docs]class TestBinaryWriter:
"""Test the binary writer."""
[docs] def setup_method(self):
"""Add P2G configs to the Satpy path."""
from polar2grid.utils.config import add_polar2grid_config_paths
self._old_path = satpy.config.get("config_path")
add_polar2grid_config_paths()
# add test specific configs
curr_path = satpy.config.get("config_path")
satpy.config.set(config_path=[TEST_ETC_DIR] + curr_path)
[docs] def teardown_method(self):
"""Reset Satpy config path back to the original value."""
satpy.config.set(config_path=self._old_path)
[docs] @pytest.mark.parametrize("enhance", [True, False])
@pytest.mark.parametrize("dst_dtype", [None, np.float32, np.float64, np.uint8])
@pytest.mark.parametrize("src_dtype", [np.float32, np.float64, np.uint8])
def test_basic_write(self, tmpdir, src_dtype, dst_dtype, enhance):
"""Test writing data to disk."""
src_data_arr = _create_fake_data_arr(dtype=src_dtype)
scn = Scene()
scn[src_data_arr.attrs["name"]] = src_data_arr
scn.save_datasets(writer="binary", base_dir=str(tmpdir), dtype=dst_dtype, enhance=enhance)
exp_fn = tmpdir.join("noaa-20_viirs_fake_p2g_name_20210101_120000_fake_area.dat")
if dst_dtype is None:
dst_dtype = src_dtype if src_dtype != np.float64 else np.float32
assert os.path.isfile(exp_fn)
data = np.memmap(str(exp_fn), mode="r", dtype=dst_dtype)
exp_data = self._generate_expected_output(src_data_arr, dst_dtype, enhance)
np.testing.assert_allclose(data, exp_data, atol=2e-7)
[docs] @staticmethod
def _generate_expected_output(src_data_arr, dst_dtype, enhance):
src_dtype = src_data_arr.dtype
exp_data = src_data_arr.values.astype(np.float64).ravel()
if enhance:
exp_data = exp_data / 500.0 # see polar2grid/tests/etc/enhancements/generic.yaml
if not np.issubdtype(dst_dtype, np.floating):
rmin, rmax = _int_min_max(dst_dtype)
exp_data = exp_data * (rmax - rmin) + rmin
exp_data = np.clip(exp_data, rmin, rmax)
exp_data = exp_data.astype(dst_dtype)
elif not np.issubdtype(dst_dtype, np.floating) or not np.issubdtype(src_dtype, np.floating):
rmin, rmax = _min_max_for_two_dtypes(src_dtype, dst_dtype)
exp_data = np.clip(exp_data, rmin, rmax)
return exp_data
