-
Notifications
You must be signed in to change notification settings - Fork 303
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #481 from djhoese/feature-scmi-rgbs
Refactor SCMI writer to be dask friendly
- Loading branch information
Showing
3 changed files
with
355 additions
and
298 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,11 +1,7 @@ | ||
#!/usr/bin/env python | ||
# -*- coding: utf-8 -*- | ||
# | ||
# Copyright (c) 2017 David Hoese | ||
# | ||
# Author(s): | ||
# | ||
# David Hoese <[email protected]> | ||
# Copyright (c) 2017-2018 SatPy Developers | ||
# | ||
# 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 | ||
|
@@ -23,9 +19,11 @@ | |
""" | ||
import os | ||
import sys | ||
from glob import glob | ||
from datetime import datetime, timedelta | ||
|
||
import numpy as np | ||
import dask.array as da | ||
|
||
try: | ||
from unittest import mock | ||
|
@@ -70,14 +68,15 @@ def test_basic_numbered_1_tile(self): | |
'test', | ||
'test', | ||
'test', | ||
proj_dict=proj4_str_to_dict('+proj=lcc +datum=WGS84 +ellps=WGS84 +lon_0=-95. +lat_0=25 +lat_1=25 +units=m +no_defs'), | ||
proj_dict=proj4_str_to_dict('+proj=lcc +datum=WGS84 +ellps=WGS84 +lon_0=-95. ' | ||
'+lat_0=25 +lat_1=25 +units=m +no_defs'), | ||
x_size=100, | ||
y_size=200, | ||
area_extent=(-1000., -1500., 1000., 1500.), | ||
) | ||
now = datetime.utcnow() | ||
now = datetime(2018, 1, 1, 12, 0, 0) | ||
ds = DataArray( | ||
np.linspace(0., 1., 20000, dtype=np.float32).reshape((200, 100)), | ||
da.from_array(np.linspace(0., 1., 20000, dtype=np.float32).reshape((200, 100)), chunks=50), | ||
attrs=dict( | ||
name='test_ds', | ||
platform_name='PLAT', | ||
|
@@ -87,8 +86,10 @@ def test_basic_numbered_1_tile(self): | |
start_time=now, | ||
end_time=now + timedelta(minutes=20)) | ||
) | ||
fn = w.save_datasets([ds], sector_id='TEST', source_name="TESTS") | ||
self.assertTrue(os.path.isfile(fn)) | ||
w.save_datasets([ds], sector_id='TEST', source_name='TESTS') | ||
all_files = glob(os.path.join(self.base_dir, 'TESTS_AII*.nc')) | ||
self.assertEqual(len(all_files), 1) | ||
self.assertEqual(os.path.basename(all_files[0]), 'TESTS_AII_PLAT_SENSOR_test_ds_TEST_T001_20180101_1200.nc') | ||
|
||
def test_basic_numbered_tiles(self): | ||
"""Test creating a multiple numbered tiles""" | ||
|
@@ -101,14 +102,15 @@ def test_basic_numbered_tiles(self): | |
'test', | ||
'test', | ||
'test', | ||
proj_dict=proj4_str_to_dict('+proj=lcc +datum=WGS84 +ellps=WGS84 +lon_0=-95. +lat_0=25 +lat_1=25 +units=m +no_defs'), | ||
proj_dict=proj4_str_to_dict('+proj=lcc +datum=WGS84 +ellps=WGS84 +lon_0=-95. ' | ||
'+lat_0=25 +lat_1=25 +units=m +no_defs'), | ||
x_size=100, | ||
y_size=200, | ||
area_extent=(-1000., -1500., 1000., 1500.), | ||
) | ||
now = datetime.utcnow() | ||
now = datetime(2018, 1, 1, 12, 0, 0) | ||
ds = DataArray( | ||
np.linspace(0., 1., 20000, dtype=np.float32).reshape((200, 100)), | ||
da.from_array(np.linspace(0., 1., 20000, dtype=np.float32).reshape((200, 100)), chunks=50), | ||
attrs=dict( | ||
name='test_ds', | ||
platform_name='PLAT', | ||
|
@@ -118,13 +120,9 @@ def test_basic_numbered_tiles(self): | |
start_time=now, | ||
end_time=now + timedelta(minutes=20)) | ||
) | ||
fn = w.save_datasets([ds], | ||
sector_id='TEST', | ||
source_name="TESTS", | ||
tile_count=(3, 3)) | ||
# `fn` is currently the last file created | ||
self.assertTrue(os.path.isfile(fn)) | ||
self.assertIn('T009', fn) | ||
w.save_datasets([ds], sector_id='TEST', source_name="TESTS", tile_count=(3, 3)) | ||
all_files = glob(os.path.join(self.base_dir, 'TESTS_AII*.nc')) | ||
self.assertEqual(len(all_files), 9) | ||
|
||
def test_basic_lettered_tiles(self): | ||
"""Test creating a lettered grid""" | ||
|
@@ -137,14 +135,15 @@ def test_basic_lettered_tiles(self): | |
'test', | ||
'test', | ||
'test', | ||
proj_dict=proj4_str_to_dict('+proj=lcc +datum=WGS84 +ellps=WGS84 +lon_0=-95. +lat_0=25 +lat_1=25 +units=m +no_defs'), | ||
proj_dict=proj4_str_to_dict('+proj=lcc +datum=WGS84 +ellps=WGS84 +lon_0=-95. ' | ||
'+lat_0=25 +lat_1=25 +units=m +no_defs'), | ||
x_size=1000, | ||
y_size=2000, | ||
area_extent=(-1000000., -1500000., 1000000., 1500000.), | ||
) | ||
now = datetime.utcnow() | ||
now = datetime(2018, 1, 1, 12, 0, 0) | ||
ds = DataArray( | ||
np.linspace(0., 1., 2000000, dtype=np.float32).reshape((2000, 1000)), | ||
da.from_array(np.linspace(0., 1., 2000000, dtype=np.float32).reshape((2000, 1000)), chunks=500), | ||
attrs=dict( | ||
name='test_ds', | ||
platform_name='PLAT', | ||
|
@@ -154,13 +153,9 @@ def test_basic_lettered_tiles(self): | |
start_time=now, | ||
end_time=now + timedelta(minutes=20)) | ||
) | ||
fn = w.save_datasets([ds], | ||
sector_id='LCC', | ||
source_name="TESTS", | ||
tile_count=(3, 3), | ||
lettered_grid=True) | ||
# `fn` is currently the last file created | ||
self.assertTrue(os.path.isfile(fn)) | ||
w.save_datasets([ds], sector_id='LCC', source_name="TESTS", tile_count=(3, 3), lettered_grid=True) | ||
all_files = glob(os.path.join(self.base_dir, 'TESTS_AII*.nc')) | ||
self.assertEqual(len(all_files), 16) | ||
|
||
def test_lettered_tiles_no_fit(self): | ||
"""Test creating a lettered grid with no data""" | ||
|
@@ -173,14 +168,15 @@ def test_lettered_tiles_no_fit(self): | |
'test', | ||
'test', | ||
'test', | ||
proj_dict=proj4_str_to_dict('+proj=lcc +datum=WGS84 +ellps=WGS84 +lon_0=-95. +lat_0=25 +lat_1=25 +units=m +no_defs'), | ||
proj_dict=proj4_str_to_dict('+proj=lcc +datum=WGS84 +ellps=WGS84 +lon_0=-95. ' | ||
'+lat_0=25 +lat_1=25 +units=m +no_defs'), | ||
x_size=1000, | ||
y_size=2000, | ||
area_extent=(4000000., 5000000., 5000000., 6000000.), | ||
) | ||
now = datetime.utcnow() | ||
now = datetime(2018, 1, 1, 12, 0, 0) | ||
ds = DataArray( | ||
np.linspace(0., 1., 2000000, dtype=np.float32).reshape((2000, 1000)), | ||
da.from_array(np.linspace(0., 1., 2000000, dtype=np.float32).reshape((2000, 1000)), chunks=500), | ||
attrs=dict( | ||
name='test_ds', | ||
platform_name='PLAT', | ||
|
@@ -190,14 +186,10 @@ def test_lettered_tiles_no_fit(self): | |
start_time=now, | ||
end_time=now + timedelta(minutes=20)) | ||
) | ||
fn = w.save_datasets([ds], | ||
sector_id='LCC', | ||
source_name="TESTS", | ||
tile_count=(3, 3), | ||
lettered_grid=True) | ||
# `fn` is currently the last file created | ||
w.save_datasets([ds], sector_id='LCC', source_name="TESTS", tile_count=(3, 3), lettered_grid=True) | ||
# No files created | ||
self.assertIsNone(fn) | ||
all_files = glob(os.path.join(self.base_dir, 'TESTS_AII*.nc')) | ||
self.assertEqual(len(all_files), 0) | ||
|
||
def test_lettered_tiles_bad_filename(self): | ||
"""Test creating a lettered grid with a bad filename""" | ||
|
@@ -210,14 +202,15 @@ def test_lettered_tiles_bad_filename(self): | |
'test', | ||
'test', | ||
'test', | ||
proj_dict=proj4_str_to_dict('+proj=lcc +datum=WGS84 +ellps=WGS84 +lon_0=-95. +lat_0=25 +lat_1=25 +units=m +no_defs'), | ||
proj_dict=proj4_str_to_dict('+proj=lcc +datum=WGS84 +ellps=WGS84 +lon_0=-95. ' | ||
'+lat_0=25 +lat_1=25 +units=m +no_defs'), | ||
x_size=1000, | ||
y_size=2000, | ||
area_extent=(-1000000., -1500000., 1000000., 1500000.), | ||
) | ||
now = datetime.utcnow() | ||
now = datetime(2018, 1, 1, 12, 0, 0) | ||
ds = DataArray( | ||
np.linspace(0., 1., 2000000, dtype=np.float32).reshape((2000, 1000)), | ||
da.from_array(np.linspace(0., 1., 2000000, dtype=np.float32).reshape((2000, 1000)), chunks=500), | ||
attrs=dict( | ||
name='test_ds', | ||
platform_name='PLAT', | ||
|
@@ -230,14 +223,52 @@ def test_lettered_tiles_bad_filename(self): | |
self.assertRaises(KeyError, w.save_datasets, | ||
[ds], | ||
sector_id='LCC', | ||
source_name="TESTS", | ||
source_name='TESTS', | ||
tile_count=(3, 3), | ||
lettered_grid=True) | ||
|
||
def test_basic_numbered_tiles_rgb(self): | ||
"""Test creating a multiple numbered tiles with RGB""" | ||
from satpy.writers.scmi import SCMIWriter | ||
from xarray import DataArray | ||
from pyresample.geometry import AreaDefinition | ||
from pyresample.utils import proj4_str_to_dict | ||
w = SCMIWriter(base_dir=self.base_dir, compress=True) | ||
area_def = AreaDefinition( | ||
'test', | ||
'test', | ||
'test', | ||
proj_dict=proj4_str_to_dict('+proj=lcc +datum=WGS84 +ellps=WGS84 +lon_0=-95. ' | ||
'+lat_0=25 +lat_1=25 +units=m +no_defs'), | ||
x_size=100, | ||
y_size=200, | ||
area_extent=(-1000., -1500., 1000., 1500.), | ||
) | ||
now = datetime(2018, 1, 1, 12, 0, 0) | ||
ds = DataArray( | ||
da.from_array(np.linspace(0., 1., 60000, dtype=np.float32).reshape((3, 200, 100)), chunks=50), | ||
dims=('bands', 'y', 'x'), | ||
coords={'bands': ['R', 'G', 'B']}, | ||
attrs=dict( | ||
name='test_ds', | ||
platform_name='PLAT', | ||
sensor='SENSOR', | ||
units='1', | ||
area=area_def, | ||
start_time=now, | ||
end_time=now + timedelta(minutes=20)) | ||
) | ||
w.save_datasets([ds], sector_id='TEST', source_name="TESTS", tile_count=(3, 3)) | ||
all_files = glob(os.path.join(self.base_dir, 'TESTS_AII*test_ds_R*.nc')) | ||
self.assertEqual(len(all_files), 9) | ||
all_files = glob(os.path.join(self.base_dir, 'TESTS_AII*test_ds_G*.nc')) | ||
self.assertEqual(len(all_files), 9) | ||
all_files = glob(os.path.join(self.base_dir, 'TESTS_AII*test_ds_B*.nc')) | ||
self.assertEqual(len(all_files), 9) | ||
|
||
|
||
def suite(): | ||
"""The test suite for this writer's tests. | ||
""" | ||
"""The test suite for this writer's tests.""" | ||
loader = unittest.TestLoader() | ||
mysuite = unittest.TestSuite() | ||
mysuite.addTest(loader.loadTestsFromTestCase(TestSCMIWriter)) | ||
|
Oops, something went wrong.