From 363ef69d88dbf859bfa5ee2a057f2b3930012ad0 Mon Sep 17 00:00:00 2001 From: Ken Kehoe Date: Tue, 12 Dec 2023 12:48:51 -0700 Subject: [PATCH] Correctly inserting two nan values for 2D plots to ensure minimal streaking in plot. (#770) * Correctly inserting two nan values for 2D plots to ensure minimal streaking in plot. * ENH: Adding in simple test * ENH: PEP8 --------- Co-authored-by: AdamTheisen --- act/tests/test_utils.py | 11 +++++++++++ act/utils/data_utils.py | 23 +++++++++++++++++++---- 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/act/tests/test_utils.py b/act/tests/test_utils.py index 5a5faf5dd0..f62a6a9360 100644 --- a/act/tests/test_utils.py +++ b/act/tests/test_utils.py @@ -82,6 +82,17 @@ def test_add_in_nan(): index = np.where(time_filled == np.datetime64('2019-01-01T01:40'))[0] assert index.size == 0 + # Test for 2D data + time = np.arange('2019-01-01T01:00', '2019-01-01T02:00', dtype='datetime64[m]') + data = np.random.random((len(time), 25)) + + time = np.delete(time, range(3, 8)) + data = np.delete(data, range(3, 8), axis=0) + time_filled, data_filled = act.utils.add_in_nan(time, data) + + assert np.count_nonzero(np.isnan(data_filled[3, :])) == 25 + assert len(time_filled) == len(time) + 2 + def test_get_missing_value(): ds = act.io.arm.read_arm_netcdf(act.tests.sample_files.EXAMPLE_EBBR1) diff --git a/act/utils/data_utils.py b/act/utils/data_utils.py index 0e7e24fa99..5d499338dc 100644 --- a/act/utils/data_utils.py +++ b/act/utils/data_utils.py @@ -427,6 +427,7 @@ def add_in_nan(time, data): # Leaving code in here in case we need to update. # diff = np.diff(time.astype('datetime64[s]'), 1) diff = np.diff(time, 1) + # Wrapping in a try to catch error while switching between numpy 1.10 to 1.11 try: mode = stats.mode(diff, keepdims=True).mode[0] @@ -437,10 +438,24 @@ def add_in_nan(time, data): offset = 0 for i in index[0]: corr_i = i + offset - time_added = time[corr_i] + (time[corr_i + 1] - time[corr_i]) / 2.0 - time = np.insert(time, corr_i + 1, time_added) - data = np.insert(data, corr_i + 1, np.nan, axis=0) - offset += 1 + + if len(data.shape) == 1: + # For line plotting adding a NaN will stop the connection of the line + # between points. So we just need to add a NaN anywhere between the points. + corr_i = i + offset + time_added = time[corr_i] + (time[corr_i + 1] - time[corr_i]) / 2.0 + time = np.insert(time, corr_i + 1, time_added) + data = np.insert(data, corr_i + 1, np.nan, axis=0) + offset += 1 + else: + # For 2D plots need to add a NaN right after and right before the data + # to correctly mitigate streaking with pcolormesh. + time_added_1 = time[corr_i] + 1 # One time step after + time_added_2 = time[corr_i + 1] - 1 # One time step before + time = np.insert(time, corr_i + 1, [time_added_1, time_added_2]) + data = np.insert(data, corr_i + 1, np.nan, axis=0) + data = np.insert(data, corr_i + 2, np.nan, axis=0) + offset += 2 if time_is_DataArray: time = xr.DataArray(time, attrs=time_attributes, dims=time_dims)