Source code for satpy.tests.writer_tests.test_geotiff

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright (c) 2017 Satpy developers
#
# This file is part of satpy.
#
# satpy 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.
#
# satpy 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
# satpy.  If not, see <http://www.gnu.org/licenses/>.
"""Tests for the geotiff writer."""

from datetime import datetime
from unittest import mock

import dask.array as da
import numpy as np
import pytest
import xarray as xr

# NOTE:
# The following fixtures are not defined in this file, but are used and injected by Pytest:
# - tmp_path


[docs] def _get_test_datasets_2d(): """Create a single 2D test dataset.""" from pyresample.geometry import AreaDefinition adef = AreaDefinition( "test", "test", "test", "EPSG:4326", 100, 200, (-180., -90., 180., 90.), ) ds1 = xr.DataArray( da.zeros((100, 200), chunks=50), dims=("y", "x"), attrs={"name": "test", "start_time": datetime.utcnow(), "units": "K", "area": adef} ) return [ds1]
[docs] def _get_test_datasets_2d_nonlinear_enhancement(): data_arrays = _get_test_datasets_2d() enh_history = [ {"gamma": 2.0}, ] for data_arr in data_arrays: data_arr.attrs["enhancement_history"] = enh_history return data_arrays
[docs] def _get_test_datasets_3d(): """Create a single 3D test dataset.""" from pyresample.geometry import AreaDefinition adef = AreaDefinition( "test", "test", "test", "EPSG:4326", 100, 200, (-180., -90., 180., 90.), ) ds1 = xr.DataArray( da.zeros((3, 100, 200), chunks=50), dims=("bands", "y", "x"), coords={"bands": ["R", "G", "B"]}, attrs={"name": "test", "start_time": datetime.utcnow(), "area": adef} ) return [ds1]
[docs] class TestGeoTIFFWriter: """Test the GeoTIFF Writer class."""
[docs] def test_init(self): """Test creating the writer with no arguments.""" from satpy.writers.geotiff import GeoTIFFWriter GeoTIFFWriter()
[docs] @pytest.mark.parametrize( "input_func", [ _get_test_datasets_2d, _get_test_datasets_3d ] ) def test_simple_write(self, input_func, tmp_path): """Test basic writer operation.""" from satpy.writers.geotiff import GeoTIFFWriter datasets = input_func() w = GeoTIFFWriter(base_dir=tmp_path) w.save_datasets(datasets)
[docs] def test_simple_delayed_write(self, tmp_path): """Test writing can be delayed.""" from satpy.writers.geotiff import GeoTIFFWriter datasets = _get_test_datasets_2d() w = GeoTIFFWriter(base_dir=tmp_path) # when we switch to rio_save on XRImage then this will be sources # and targets res = w.save_datasets(datasets, compute=False) # this will fail if rasterio isn't installed assert isinstance(res, tuple) # two lists, sources and destinations assert len(res) == 2 assert isinstance(res[0], list) assert isinstance(res[1], list) assert isinstance(res[0][0], da.Array) da.store(res[0], res[1]) for target in res[1]: if hasattr(target, "close"): target.close()
[docs] def test_colormap_write(self, tmp_path): """Test writing an image with a colormap.""" from trollimage.colormap import spectral from trollimage.xrimage import XRImage from satpy.writers.geotiff import GeoTIFFWriter datasets = _get_test_datasets_2d() w = GeoTIFFWriter(base_dir=tmp_path) # we'd have to customize enhancements to test this through # save_datasets. We'll use `save_image` as a workaround. img = XRImage(datasets[0]) img.palettize(spectral) w.save_image(img, keep_palette=True)
[docs] def test_float_write(self, tmp_path): """Test that geotiffs can be written as floats. NOTE: Does not actually check that the output is floats. """ from satpy.writers.geotiff import GeoTIFFWriter datasets = _get_test_datasets_2d() w = GeoTIFFWriter(base_dir=tmp_path, enhance=False, dtype=np.float32) w.save_datasets(datasets)
[docs] def test_dtype_for_enhance_false(self, tmp_path): """Test that dtype of dataset is used if parameters enhance=False and dtype=None.""" from satpy.writers.geotiff import GeoTIFFWriter datasets = _get_test_datasets_2d() w = GeoTIFFWriter(base_dir=tmp_path, enhance=False) with mock.patch("satpy.writers.XRImage.save") as save_method: save_method.return_value = None w.save_datasets(datasets, compute=False) assert save_method.call_args[1]["dtype"] == np.float64
[docs] def test_dtype_for_enhance_false_and_given_dtype(self, tmp_path): """Test that dtype of dataset is used if enhance=False and dtype=uint8.""" from satpy.writers.geotiff import GeoTIFFWriter datasets = _get_test_datasets_2d() w = GeoTIFFWriter(base_dir=tmp_path, enhance=False, dtype=np.uint8) with mock.patch("satpy.writers.XRImage.save") as save_method: save_method.return_value = None w.save_datasets(datasets, compute=False) assert save_method.call_args[1]["dtype"] == np.uint8
[docs] def test_fill_value_from_config(self, tmp_path): """Test fill_value coming from the writer config.""" from satpy.writers.geotiff import GeoTIFFWriter datasets = _get_test_datasets_2d() w = GeoTIFFWriter(base_dir=tmp_path) w.info["fill_value"] = 128 with mock.patch("satpy.writers.XRImage.save") as save_method: save_method.return_value = None w.save_datasets(datasets, compute=False) assert save_method.call_args[1]["fill_value"] == 128
[docs] def test_tags(self, tmp_path): """Test tags being added.""" from satpy.writers.geotiff import GeoTIFFWriter datasets = _get_test_datasets_2d() w = GeoTIFFWriter(tags={"test1": 1}, base_dir=tmp_path) w.info["fill_value"] = 128 with mock.patch("satpy.writers.XRImage.save") as save_method: save_method.return_value = None w.save_datasets(datasets, tags={"test2": 2}, compute=False) called_tags = save_method.call_args[1]["tags"] assert called_tags == {"test1": 1, "test2": 2}
[docs] @pytest.mark.parametrize( "input_func", [ _get_test_datasets_2d, _get_test_datasets_3d, _get_test_datasets_2d_nonlinear_enhancement, ] ) @pytest.mark.parametrize( "save_kwargs", [ {"include_scale_offset": True}, {"scale_offset_tags": ("scale", "offset")}, ] ) def test_scale_offset(self, input_func, save_kwargs, tmp_path): """Test tags being added.""" from satpy.writers.geotiff import GeoTIFFWriter datasets = input_func() w = GeoTIFFWriter(tags={"test1": 1}, base_dir=tmp_path) w.info["fill_value"] = 128 with mock.patch("satpy.writers.XRImage.save") as save_method: save_method.return_value = None w.save_datasets(datasets, tags={"test2": 2}, compute=False, **save_kwargs) kwarg_name = "include_scale_offset_tags" if "include_scale_offset" in save_kwargs else "scale_offset_tags" kwarg_value = save_method.call_args[1].get(kwarg_name) assert kwarg_value is not None
[docs] def test_tiled_value_from_config(self, tmp_path): """Test tiled value coming from the writer config.""" from satpy.writers.geotiff import GeoTIFFWriter datasets = _get_test_datasets_2d() w = GeoTIFFWriter(base_dir=tmp_path) with mock.patch("satpy.writers.XRImage.save") as save_method: save_method.return_value = None w.save_datasets(datasets, compute=False) assert save_method.call_args[1]["tiled"]
[docs] def test_float_write_with_unit_conversion(self, tmp_path): """Test that geotiffs can be written as floats and convert units.""" from satpy.writers.geotiff import GeoTIFFWriter dataset = _get_test_datasets_2d()[0] dtype = np.float32 w = GeoTIFFWriter(base_dir=tmp_path, enhance=False, dtype=dtype) filename = tmp_path / "data_in_C.tif" w.save_dataset(dataset, filename=filename, units="degC") ds = xr.open_dataset(filename, engine="rasterio") assert ds["band_data"].dtype == dtype np.testing.assert_allclose(ds["band_data"], -273.15)