Sort datasets by shot/time

disruption_py/settings/__init__.py

@@ -3,10 +3,7 @@
 """Package initialization for the settings module."""
 
 from .log_settings import LogSettings
-from .output_setting import (
-    OutputSetting,
-    OutputSettingParams,
-)
+from .output_setting import OutputSetting, OutputSettingParams
 from .retrieval_settings import RetrievalSettings
 from .shotlist_setting import DatabaseShotlistSetting, FileShotlistSetting
 from .time_setting import TimeSetting, TimeSettingParams

disruption_py/settings/output_setting.py

@@ -266,8 +266,13 @@ def get_results(self) -> OutputSingleType:
         xr.Dataset | xr.DataTree | pd.DataFrame
             The resulting object.
         """
-        logger.debug("Concatenating {tot} shots.", tot=len(self.results))
+        logger.debug("Concatenating {tot:,} shots...", tot=len(self.results))
+        took = -time.time()
         self.result = self.concat()
+        took += time.time()
+        logger.info(
+            "Concatenated {tot:,} shots in {sec:.3f}s.", tot=len(self.results), sec=took
+        )
         self.results = {}
         return self.result
 
@@ -319,7 +324,16 @@ def concat(self) -> xr.Dataset:
         if not self.results:
             logger.critical("Nothing to concatenate!")
             return xr.Dataset()
-        return xr.concat(self.results.values(), dim="idx", combine_attrs="no_conflicts")
+
+        ds = xr.concat(self.results.values(), dim="idx", combine_attrs="no_conflicts")
+        if "shot" not in ds.coords or "time" not in ds.coords:
+            return ds
+
+        took = -time.time()
+        ds = ds.sortby(["shot", "time"])
+        took += time.time()
+        logger.debug("Sorted {tot:,} rows in {sec:.3f}s.", tot=len(ds.idx), sec=took)
+        return ds
 
 
 class DataTreeOutputSetting(SingleOutputSetting):

disruption_py/workflow.py

@@ -16,10 +16,7 @@
 
 from disruption_py.config import config
 from disruption_py.core.retrieval_manager import RetrievalManager
-from disruption_py.core.utils.misc import (
-    get_elapsed_time,
-    without_duplicates,
-)
+from disruption_py.core.utils.misc import get_elapsed_time, without_duplicates
 from disruption_py.inout.mds import ProcessMDSConnection
 from disruption_py.inout.sql import ShotDatabase
 from disruption_py.inout.xr import XarrayConnection

tests/test_output_setting.py

@@ -94,8 +94,8 @@ def test_output_exists(fresh_data, test_folder_m):
     pd.testing.assert_frame_equal(df_out, df_dsk)
 
     # format equivalence
-    xr.testing.assert_identical(ds_out, xr.concat(dict_out.values(), dim="idx"))
-    xr.testing.assert_identical(
-        ds_out, xr.concat([dt.to_dataset() for dt in dt_out.values()], dim="idx")
-    )
+    for ds_list in [dict_out.values(), [dt.to_dataset() for dt in dt_out.values()]]:
+        xr.testing.assert_identical(
+            ds_out, xr.concat(ds_list, dim="idx").sortby(["shot", "time"])
+        )
     pd.testing.assert_frame_equal(df_out, ds_out.to_dataframe()[df_out.columns])