make color_by_time type agnostic and respect tlim

Author: wx4stg

pyxlma/plot/xlma_plot_feature.py

@@ -25,19 +25,24 @@ def subset(lon_data, lat_data, alt_data, time_data, chi_data,station_data,
     return lon_data, lat_data, alt_data, time_data, selection
 
 
-def color_by_time(time_array, tlim):
+def color_by_time(time_array, tlim=None):
     """
-    Generates colormap values for plotting VHF sources by time in a
+    Generates colormap values for plotting scatter points by time in a
     given time window
 
     Returns: min, max values, array by time
     """
-    vmax = (tlim[1] - time_array.min()).total_seconds()
-    ldiff = time_array - time_array.min()
-    ldf = []
-    for df in ldiff:
-        ldf.append(df.total_seconds())
-    c = np.array(ldf)
+    if tlim is None:
+        tlim = np.array([np.atleast_1d(time_array.min())[0],
+                         np.atleast_1d(time_array.max())[0]])
+    nsToS = 1e9
+    time_array = np.array(time_array).astype('datetime64[ns]').astype(float)/nsToS
+    tlim = np.atleast_1d(tlim).astype('datetime64[ns]').astype(float)/nsToS
+    
+    ldiff = time_array - tlim[0]
+
+    vmax = tlim[1] - tlim[0]
+    c = ldiff.astype(float)
     vmin = 0
 
     return vmin, vmax, c

tests/test_plot_feature.py

@@ -17,14 +17,29 @@ def test_subset():
                                               '2023-12-24T00:57:07.814960674', '2023-12-24T00:57:07.826344209']).astype(np.datetime64).astype(float))
     assert np.sum(selection) == 10
 
-def test_color_by_time_datetime():
+def test_color_by_time_datetime_nolimit():
+    some_datetimes = np.array([dt(2021, 4, 9, 1, 51, 0), dt(2021, 4, 9, 1, 52, 0), dt(2021, 4, 9, 1, 53, 0), dt(2021, 4, 9, 1, 54, 0), dt(2021, 4, 9, 1, 59, 0)])
+    vmin, vmax, colors = color_by_time(some_datetimes)
+    assert vmin == 0
+    assert vmax == 480
+    assert np.allclose(colors, np.array([0, 60, 120, 180, 480]))
+
+
+def test_color_by_time_datetime_limit():
     some_datetimes = np.array([dt(2021, 4, 9, 1, 51, 0), dt(2021, 4, 9, 1, 52, 0), dt(2021, 4, 9, 1, 53, 0), dt(2021, 4, 9, 1, 54, 0), dt(2021, 4, 9, 1, 59, 0)])
     limits = [dt(2021, 4, 9, 1, 50, 0), dt(2021, 4, 9, 2, 0, 0)]
     vmin, vmax, colors = color_by_time(some_datetimes, limits)
     assert vmin == 0
-    assert vmax == 540
-    assert np.allclose(colors, np.array([0, 60, 120, 180, 480]))
+    assert vmax == 600
+    assert np.allclose(colors, np.array([60, 120, 180, 240, 540]))
 
+def test_color_by_time_xarray():
+    dataset = xr.open_dataset('tests/truth/lma_netcdf/lma.nc')
+    vmin, vmax, colors = color_by_time(dataset.event_time)
+    assert vmin == 0
+    assert np.isclose(vmax, 57.943683385849)
+    assert np.isclose(np.mean(colors), 30.483982899376258)
+    assert np.isclose(np.std(colors), 17.25687093241869)
 
 def test_setup_hist():
     lma = xr.open_dataset('tests/truth/lma_netcdf/lma.nc')