#!/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/
"""The Binary writer writes band data to a flat binary file.
By default it enhances the data based on the enhancement configuration file
and then saves the data to a flat binary file. The output data type will match
the input data for integer types (ex. uint8 -> uint8), but floating point types
will always be forced to 32-bit floats for consistency between readers and
changes in the low-level Satpy library. A different output type can be
specified using the ``--dtype`` flag. To turn off scaling of the data (a.k.a.
enhancements) the ``--no-enhance`` command line flag can be specified to write
the "raw" band data.
"""
from __future__ import annotations
import logging
import dask.array as da
import numpy as np
import xarray as xr
from satpy.writers import ImageWriter, get_enhanced_image
from polar2grid.core.dtype import NUMPY_DTYPE_STRS, clip_to_data_type, dtype_to_str, int_or_float, str_to_dtype
from polar2grid.core.script_utils import NumpyDtypeList
from polar2grid.utils.legacy_compat import convert_p2g_pattern_to_satpy
logger = logging.getLogger(__name__)
DEFAULT_OUTPUT_FILENAMES = {
"polar2grid": {
None: "{platform_name!l}_{sensor!l}_{p2g_name}_{start_time:%Y%m%d_%H%M%S}_{area.area_id}.dat",
},
"geo2grid": {
None: "{platform_name!u}_{sensor!u}_{p2g_name}_{start_time:%Y%m%d_%H%M%S}_{area.area_id}.dat",
"abi_l1b": "{platform_name!u}_{sensor!u}_{observation_type}{scene_abbr}_"
"{p2g_name}_{start_time:%Y%m%d_%H%M%S}_{area.area_id}.dat",
"clavrx": "{platform_name!l}_{sensor!l}_{p2g_name}_{start_time:%Y%m%d_%H%M%S}_{area.area_id}.dat",
},
}
[docs]
class FlatBinaryWriter(ImageWriter):
"""Write data to disk as flat binary files."""
[docs]
def save_dataset(
self, dataset, filename=None, fill_value=None, overlay=None, decorate=None, compute=True, **kwargs
):
"""Save the ``dataset`` to a given ``filename``.
This method creates an enhanced image using :func:`get_enhanced_image`.
The image is then passed to :meth:`save_image`. See both of these
functions for more details on the arguments passed to this method.
"""
img = get_enhanced_image(
dataset.squeeze(),
enhance=self.enhancer,
overlay=overlay,
decorate=decorate,
)
kwargs["dtype"] = kwargs.get("dtype") or self._get_default_dtype(dataset)
return self.save_image(img, filename=filename, compute=compute, fill_value=fill_value, **kwargs)
[docs]
@staticmethod
def _get_default_dtype(data_arr: xr.DataArray) -> np.dtype:
if data_arr.dtype == np.float64:
return np.float32
return data_arr.dtype
[docs]
def save_image(self, img, filename=None, compute=True, dtype=None, fill_value=None, **kwargs):
filename = filename or self.get_filename(
data_type=dtype_to_str(dtype), rows=img.data.shape[0], columns=img.data.shape[1], **img.data.attrs
)
data = self._prep_data(img.data, dtype, fill_value)
logger.info("Saving product %s to binary file %s", img.data.attrs["p2g_name"], filename)
dst = np.memmap(filename, shape=img.data.shape, dtype=dtype, mode="w+")
if compute:
da.store(data, dst)
return
return [[data], [dst]]
[docs]
def _prep_data(self, data: xr.DataArray, dtype: np.dtype, fill_value) -> da.Array:
fill = data.attrs.get("_FillValue", np.nan)
if fill_value is None:
fill_value = fill
final_data = data.data
if self.enhancer and np.issubdtype(data.dtype, np.floating) and not np.issubdtype(dtype, np.floating):
# going from float -> int and the data was enhanced
# scale the data to fit the integer dtype
rmin, rmax = np.iinfo(dtype).min, np.iinfo(dtype).max
final_data = final_data * (rmax - rmin) + rmin
final_data = clip_to_data_type(final_data, dtype)
same_fill = np.isnan(fill) and np.isnan(fill_value) or fill == fill_value
if data.dtype == dtype and same_fill:
return final_data
final_data = final_data.astype(dtype)
if same_fill:
return final_data
fill_mask = np.isnan(final_data) if np.isnan(fill) else final_data == fill
final_data = da.where(fill_mask, fill_value, final_data)
return final_data
[docs]
def add_writer_argument_groups(parser, group=None):
if group is None:
group = parser.add_argument_group(title="Binary Writer")
group.add_argument(
"--output-filename",
dest="filename",
type=convert_p2g_pattern_to_satpy,
help="Custom file pattern to save dataset to",
)
group.add_argument(
"--dtype",
choices=NumpyDtypeList(NUMPY_DTYPE_STRS),
type=str_to_dtype,
help="Data type of the output file (8-bit unsigned integer by default - uint8)",
)
group.add_argument(
"--no-enhance",
dest="enhance",
action="store_false",
help="Don't enhance the data before saving it",
)
group.add_argument(
"--fill-value",
dest="fill_value",
type=int_or_float,
help="Replace invalid values with the specified value. Floating-point "
"products typically use NaN while integer fields will use 0 or "
"the max value for that data type.",
)
return group, None
