xlma-python/tests/test_xarray_util.py at 7541af5a433ec99312e013f20f4a20e5a8aa9e8b · deeplycloudy/xlma-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
from glob import glob
import xarray as xr
from pyxlma.xarray_util import *

def test_get_1d_dims():
    lma = xr.open_dataset('tests/truth/lma_netcdf/lma_stats.nc')
    dims_1d = get_1d_dims(lma)
    print(lma)
    assert dims_1d == ['number_of_flashes']


def test_gen_1d_datasets():
    lma = xr.open_dataset('tests/truth/lma_netcdf/lma_stats.nc')
    datasets_1d = list(gen_1d_datasets(lma))
    assert len(datasets_1d) == 1
    assert list(datasets_1d[0].dims.keys()) == ['number_of_flashes']


def test_get_1d_datasets():
    lma = xr.open_dataset('tests/truth/lma_netcdf/lma_stats.nc')
    datasets_1d = get_1d_datasets(lma)
    assert len(datasets_1d) == 1
    assert list(datasets_1d[0].dims.keys()) == ['number_of_flashes']


def test_get_scalar_vars():
    lma = xr.open_dataset('tests/truth/lma_netcdf/lma_stats.nc')
    scalar_vars = get_scalar_vars(lma)
    assert scalar_vars == ['network_center_latitude', 'network_center_longitude', 'network_center_altitude',
                           'flash_distance_separation_threshold', 'flash_time_separation_threshold']


def test_concat_1d_dims_no_scalars():
    glm_datasets = [xr.open_dataset(path) for path in glob('examples/network_samples/OR_GLM*.nc')]
    concatenated = concat_1d_dims(glm_datasets)
    assert dict(concatenated.dims) == {'number_of_events': 29664, 'number_of_groups': 10420, 'number_of_flashes': 606,
                                       'number_of_time_bounds': 6, 'number_of_wavelength_bounds': 6, 'number_of_field_of_view_bounds': 6}


def test_concat_1d_dims_scalars():
    glm_datasets = [xr.open_dataset(path) for path in glob('examples/network_samples/OR_GLM*.nc')]
    concatenated = concat_1d_dims(glm_datasets, stack_scalars='scalars')
    assert dict(concatenated.dims) == {'number_of_events': 29664, 'number_of_groups': 10420, 'number_of_flashes': 606, 'number_of_time_bounds': 6,
                                       'number_of_wavelength_bounds': 6, 'number_of_field_of_view_bounds': 6, 'scalars': 3}