[codecs/HAPI/CVS] fix time index loading

Signed-off-by: Alexis Jeandet <[email protected]>
brenard-irap · Jan 7, 2025 · eabe900 · eabe900
1 parent f1b2197
commit eabe900
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 5 deletions.
diff --git a/speasy/core/codecs/bundled_codecs/hapi_csv/reader.py b/speasy/core/codecs/bundled_codecs/hapi_csv/reader.py
@@ -29,11 +29,12 @@ def _extract_headers(file: io.IOBase) -> Dict[str, Any]:
 
 def _extract_data(file: io.IOBase) -> pds.DataFrame:
     file.seek(0)
-    return pds.read_csv(file, comment='#', header=None, skiprows=0)
+    return pds.read_csv(file, comment='#', sep=',', header=None, skiprows=0, parse_dates=[0], index_col=0)
 
 
 def _parse_hapi_csv(file: io.IOBase) -> Tuple[pds.DataFrame, Dict[str, Any]]:
     headers = _extract_headers(file)
+    assert headers["parameters"][0]["type"] == "isotime"
     data = _extract_data(file)
     return data, headers
 
@@ -51,9 +52,11 @@ def load_hapi_csv(file: Union[Buffer, str, io.IOBase]) -> Optional[HapiCsvFile]:
     data, headers = _load_csv(file)
     hapi_csv_file = HapiCsvFile()
     if data is not None and headers is not None:
-        assert headers["parameters"][0]["type"] == "isotime"
-        hapi_csv_file.create_parameter(data.index.to_numpy(dtype="datetime64[ns]"), meta=headers["parameters"][0])
-        column_offset = 1
+        time_header = headers["parameters"][0]
+        assert time_header["type"] == "isotime"
+        hapi_csv_file.create_parameter(data.index.to_numpy(dtype="datetime64[ns]"),
+                                       meta=time_header)
+        column_offset = 0
         for i, param_meta in enumerate(headers["parameters"][1:]):
             shape = param_meta.get("size", [1])
             flatten_shape = np.prod(shape)

diff --git a/tests/test_hapi_codecs.py b/tests/test_hapi_codecs.py
@@ -1,4 +1,6 @@
 import unittest
+
+import numpy as np
 from ddt import ddt, data, unpack
 
 import os
@@ -27,6 +29,19 @@ def test_loads_headers(self, fname, var_name, unit, description):
             self.assertEqual(v.unit, unit)
             self.assertEqual(v.meta['description'], description)
 
+    @data(
+        ('HAPI_sample_csv.csv', 'Magnitude', np.datetime64("1997-09-02T00:00:00.000Z"), np.datetime64("1997-09-03T23:00:00.000Z")),
+        ('HAPI_sample_csv_multiple_vars.csv', 'Magnitude', np.datetime64("1997-09-02T00:00:00.000Z"), np.datetime64("1997-09-03T23:00:00.000Z"))
+    )
+    @unpack
+    def test_load_time_index(self, fname, var_name, first_value, last_value):
+        hapi_csv_codec: CodecInterface = get_codec('hapi/csv')
+        v: SpeasyVariable = hapi_csv_codec.load_variable(
+            file=str(os.path.join(__HERE__, 'resources', 'HAPI_sample_csv_multiple_vars.csv')), variable=var_name,
+            disable_cache=True)
+        self.assertEqual(v.time[0], first_value)
+        self.assertEqual(v.time[-1], last_value)
+
     @data(
         ('HAPI_sample_csv.csv', 'Magnitude', (48, 1), 2.658, 14.743),
         ('HAPI_sample_csv_multiple_vars.csv', 'Magnitude', (48, 1), 2.658, 14.743),
@@ -39,7 +54,9 @@ def test_loads_headers(self, fname, var_name, unit, description):
     @unpack
     def test_load_values(self, fname, var_name, shape, first_value, last_value):
         hapi_csv_codec: CodecInterface = get_codec('hapi/csv')
-        v: SpeasyVariable = hapi_csv_codec.load_variable(file=str(os.path.join(__HERE__, 'resources', 'HAPI_sample_csv_multiple_vars.csv')) , variable=var_name, disable_cache=True)
+        v: SpeasyVariable = hapi_csv_codec.load_variable(
+            file=str(os.path.join(__HERE__, 'resources', 'HAPI_sample_csv_multiple_vars.csv')), variable=var_name,
+            disable_cache=True)
         self.assertEqual(v.values.shape, shape)
         if type(first_value) is list:
             self.assertListEqual(v.values[0].tolist(), first_value)