Correctly inserting two nan values for 2D plots to ensure minimal str…

…eaking 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 <[email protected]>
ARM-DOE · Dec 12, 2023 · 363ef69 · 363ef69
1 parent e9cc16e
commit 363ef69
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 4 deletions.
diff --git 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
@@ -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)