From 931ec3addc499fc9321983045f6618f726f057e7 Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Thu, 9 May 2024 19:03:16 +0200
Subject: [PATCH 01/15] Output -> ElementOutput, ConnectionOutput

---
 tests/helpers.py                         |  17 +-
 tests/test_cli.py                        |   2 +-
 tests/test_output.py                     |   4 +-
 toughio/__init__.py                      |   6 +-
 toughio/_cli/_export.py                  |   5 +-
 toughio/_io/__init__.py                  |   5 +-
 toughio/_io/h5/_write.py                 |  12 +-
 toughio/_io/output/__init__.py           |   5 +-
 toughio/_io/output/_common.py            | 288 +++++++++++++++++++----
 toughio/_io/output/_helpers.py           |  46 ++--
 toughio/_io/output/csv/_csv.py           |  61 +++--
 toughio/_io/output/petrasim/_petrasim.py |  73 ++++--
 toughio/_io/output/save/_save.py         |  32 +--
 toughio/_io/output/tecplot/_tecplot.py   |  12 +-
 toughio/_io/output/tough/_tough.py       |  24 +-
 toughio/_mesh/_mesh.py                   |  16 +-
 16 files changed, 425 insertions(+), 183 deletions(-)

diff --git a/tests/helpers.py b/tests/helpers.py
index 194d7793..d9043b1b 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -72,10 +72,8 @@
 )
 
 output_eleme = [
-    toughio.Output(
-        "element",
+    toughio.ElementOutput(
         float(time),
-        np.array([f"AAA0{i}" for i in range(10)]),
         {
             "X": np.random.rand(10),
             "Y": np.random.rand(10),
@@ -83,15 +81,14 @@
             "PRES": np.random.rand(10),
             "TEMP": np.random.rand(10),
         },
+        np.array([f"AAA0{i}" for i in range(10)]),
     )
     for time in range(3)
 ]
 
 output_conne = [
-    toughio.Output(
-        "connection",
+    toughio.ConnectionOutput(
         float(time),
-        np.array([[f"AAA0{i}", f"AAA0{i}"] for i in range(10)]),
         {
             "X": np.random.rand(10),
             "Y": np.random.rand(10),
@@ -99,6 +96,7 @@
             "HEAT": np.random.rand(10),
             "FLOW": np.random.rand(10),
         },
+        np.array([[f"AAA0{i}", f"AAA0{i}"] for i in range(10)]),
     )
     for time in range(3)
 ]
@@ -170,10 +168,9 @@ def allclose(x, y, atol=1.0e-8, ignore_keys=None, ignore_none=False):
                 if x.cell_data:
                     assert allclose(x.cell_data, y.cell_data, atol=atol)
 
-            elif isinstance(x, toughio.Output):
-                assert isinstance(y, toughio.Output)
+            elif isinstance(x, (toughio.ElementOutput, toughio.ConnectionOutput)):
+                assert isinstance(y, (toughio.ElementOutput, toughio.ConnectionOutput))
 
-                assert allclose(x.type, y.type, atol=atol)
                 assert allclose(x.time, y.time, atol=atol)
                 assert allclose(x.data, y.data, atol=atol)
 
@@ -205,7 +202,7 @@ def convert_outputs_labels(outputs, connection=False):
                 output.labels[:] = toughio.convert_labels(output.labels)
 
             else:
-                labels = toughio.convert_labels(output.labels.ravel())
+                labels = toughio.convert_labels(np.ravel(output.labels))
                 output.labels[:] = labels.reshape((labels.size // 2, 2))
 
         except TypeError:
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 51179da1..1c08dfc3 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -275,5 +275,5 @@ def test_save2incon(reset):
 
     incon = toughio.read_output(output_filename)
 
-    assert save.labels.tolist() == incon.labels.tolist()
+    assert list(save.labels) == list(incon.labels)
     helpers.allclose(save.data, incon.data)
diff --git a/tests/test_output.py b/tests/test_output.py
index 29217853..7a0b0f9d 100644
--- a/tests/test_output.py
+++ b/tests/test_output.py
@@ -47,9 +47,9 @@ def test_output_eleme(filename, filename_ref, file_format):
     for output, time_ref in zip(outputs, times_ref):
         assert time_ref == output.time
         assert (
-            save.labels.tolist() == output.labels.tolist()
+            list(save.labels) == list(output.labels)
             if file_format in {"csv", "petrasim", "tough"}
-            else len(output.labels) == 0
+            else not output.labels
         )
         if file_format != "tough":
             assert keys_ref == sorted(list(output.data))
diff --git a/toughio/__init__.py b/toughio/__init__.py
index 27906915..f68c4895 100644
--- a/toughio/__init__.py
+++ b/toughio/__init__.py
@@ -2,7 +2,8 @@
 from .__about__ import __version__
 from ._helpers import convert_labels
 from ._io import (
-    Output,
+    ElementOutput,
+    ConnectionOutput,
     read_input,
     read_output,
     read_table,
@@ -24,7 +25,8 @@
 __all__ = [
     "Mesh",
     "CellBlock",
-    "Output",
+    "ElementOutput",
+    "ConnectionOutput",
     "meshmaker",
     "register_input",
     "register_output",
diff --git a/toughio/_cli/_export.py b/toughio/_cli/_export.py
index e2e845d7..fad8b53d 100644
--- a/toughio/_cli/_export.py
+++ b/toughio/_cli/_export.py
@@ -18,7 +18,6 @@ def export(argv=None):
     import sys
 
     from .. import read_mesh, read_output, write_time_series
-    from .._io.output._common import reorder_labels
     from ..meshmaker import triangulate, voxelize
 
     parser = _get_parser()
@@ -144,13 +143,13 @@ def export(argv=None):
 
         if args.file_format != "xdmf":
             mesh.point_data = {}
-            mesh.cell_dada = {}
+            mesh.cell_data = {}
             mesh.field_data = {}
             mesh.point_sets = {}
             mesh.cell_sets = {}
             mesh.read_output(output)
         else:
-            output = [reorder_labels(data, mesh.labels) for data in output]
+            output = [out[mesh.labels] for out in output]
     print(" Done!")
 
     # Output file name
diff --git a/toughio/_io/__init__.py b/toughio/_io/__init__.py
index cff65c2f..c59fd595 100644
--- a/toughio/_io/__init__.py
+++ b/toughio/_io/__init__.py
@@ -2,7 +2,7 @@
 from .input import read as read_input
 from .input import register as register_input
 from .input import write as write_input
-from .output import Output
+from .output import ElementOutput, ConnectionOutput
 from .output import read as read_output
 from .output import register as register_output
 from .output import write as write_output
@@ -10,7 +10,8 @@
 from .table import register as register_table
 
 __all__ = [
-    "Output",
+    "ElementOutput",
+    "ConnectionOutput",
     "register_input",
     "register_output",
     "read_input",
diff --git a/toughio/_io/h5/_write.py b/toughio/_io/h5/_write.py
index 6c89319b..16fbd833 100644
--- a/toughio/_io/h5/_write.py
+++ b/toughio/_io/h5/_write.py
@@ -2,7 +2,7 @@
 
 import h5py
 
-from ..output import Output
+from ..output import ElementOutput, ConnectionOutput
 from ..output import read as read_output
 from ..table import read as read_table
 
@@ -25,9 +25,9 @@ def write(
     ----------
     filename : str or pathlike
         Output file name.
-    elements : namedtuple, list of namedtuple, str, pathlike or None, optional, default None
+    elements : str, pathlike, :class:`toughio.ElementOutput`, sequence of :class:`toughio.ElementOutput` or None, optional, default None
         Element outputs to export.
-    connections : namedtuple, list of namedtuple, str, pathlike or None, optional, default None
+    connections : str, pathlike, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ConnectionOutput` or None, optional, default None
         Connection outputs to export.
     element_history : dict or None, optional, default None
         Element history to export.
@@ -79,12 +79,12 @@ def _write_output(f, outputs, labels_order, connection, **kwargs):
     if isinstance(outputs, (str, pathlib.Path)):
         outputs = read_output(outputs, labels_order=labels_order, connection=connection)
 
-    if isinstance(outputs, Output):
+    if isinstance(outputs, (ElementOutput, ConnectionOutput)):
         outputs = [outputs]
 
     elif isinstance(outputs, (list, tuple)):
         for output in outputs:
-            if not isinstance(output, Output):
+            if not isinstance(output, (ElementOutput, ConnectionOutput)):
                 raise ValueError()
 
     else:
@@ -92,7 +92,7 @@ def _write_output(f, outputs, labels_order, connection, **kwargs):
 
     for output in outputs:
         group = f.create_group(f"time={output.time}")
-        group.create_dataset("labels", data=output.labels.astype("S"), **kwargs)
+        group.create_dataset("labels", data=list(output.labels), **kwargs)
 
         for k, v in output.data.items():
             group.create_dataset(k, data=v, **kwargs)
diff --git a/toughio/_io/output/__init__.py b/toughio/_io/output/__init__.py
index 2568037c..40b9e3d4 100644
--- a/toughio/_io/output/__init__.py
+++ b/toughio/_io/output/__init__.py
@@ -1,9 +1,10 @@
 from . import csv, petrasim, save, tecplot, tough
-from ._common import Output
+from ._common import ElementOutput, ConnectionOutput
 from ._helpers import read, register, write
 
 __all__ = [
-    "Output",
+    "ElementOutput",
+    "ConnectionOutput",
     "register",
     "read",
     "write",
diff --git a/toughio/_io/output/_common.py b/toughio/_io/output/_common.py
index 5182cd3b..dc405960 100644
--- a/toughio/_io/output/_common.py
+++ b/toughio/_io/output/_common.py
@@ -1,4 +1,4 @@
-import collections
+from abc import ABC, abstractmethod
 import logging
 
 import numpy as np
@@ -8,20 +8,240 @@
 ]
 
 
-Output = collections.namedtuple("Output", ["type", "time", "labels", "data"])
+class Output(ABC):
+    def __init__(self, time, data, labels=None):
+        """
+        Base class for output data.
 
+        Do not use.
 
-def to_output(file_type, labels_order, headers, times, labels, variables):
-    """Create an Output namedtuple."""
-    outputs = [
-        Output(
-            file_type,
-            time,
-            np.array(label) if len(label) else np.array(label, dtype="<U1"),
-            {k: v for k, v in zip(headers, np.transpose(variable))},
+        """
+        self._time = time
+        self._data = data
+        self._labels = list(labels) if labels is not None else labels
+
+    @abstractmethod
+    def __getitem__(self):
+        """Slice output."""
+        pass
+    
+    @abstractmethod
+    def index(self):
+        """Get index of element or connection."""
+        if self.labels is None:
+            raise AttributeError()
+        
+    @property
+    def n_data(self):
+        """Return number of data points."""
+        return len(self.data[list(self.data)[0]])
+
+    @property
+    def time(self):
+        """Return time step (in seconds)."""
+        return self._time
+    
+    @time.setter
+    def time(self, value):
+        self._time = value
+    
+    @property
+    def data(self):
+        """Return data arrays."""
+        return self._data
+    
+    @data.setter
+    def data(self, value):
+        self._data = value
+    
+    @property
+    def labels(self):
+        """Return labels."""
+        return self._labels
+    
+    @labels.setter
+    def labels(self, value):
+        if value is not None:
+            if len(value) != self.n_data:
+                raise ValueError()
+            
+            self._labels = list(value)
+
+        else:
+            self._labels = None
+
+
+class ElementOutput(Output):
+    def __init__(self, time, data, labels=None):
+        """
+        Element output data.
+
+        Parameters
+        ----------
+        time : scalar
+            Time step (in seconds).
+        data : dict
+            Data arrays.
+        labels : sequence of str or None, default, None
+            Labels of elements.
+        
+        """
+        super().__init__(time, data, labels)
+
+    def __getitem__(self, islice):
+        """
+        Slice element output.
+        
+        Parameters
+        ----------
+        islice : int, str, slice, sequence of int or sequence of str
+            Indices or labels of elements to slice.
+
+        Returns
+        -------
+        dict or :class:`toughio.ElementOutput`
+            Sliced element outputs.
+        
+        """
+        if self.labels is None:
+            raise AttributeError()
+        
+        if np.ndim(islice) == 0:
+            if isinstance(islice, slice):
+                islice = np.arange(self.n_data)[islice]
+
+            else:
+                islice = self.index(islice) if isinstance(islice, str) else islice
+            
+                return {k: v[islice] for k, v in self.data.items()}
+        
+        elif np.ndim(islice) == 1:
+            islice = [self.index(i) if isinstance(i, str) else i for i in islice]
+
+        else:
+            raise ValueError()
+        
+        return ElementOutput(
+            self.time,
+            {k: v[islice] for k, v in self.data.items()},
+            [self._labels[i] for i in islice],
         )
-        for time, label, variable in zip(times, labels, variables)
-    ]
+
+    def index(self, label):
+        """
+        Get index of element.
+        
+        Parameters
+        ----------
+        label : str
+            Label of element.
+
+        Returns
+        -------
+        int
+            Index of element.
+        
+        """
+        super().index()
+        
+        return self.labels.index(label)
+
+
+class ConnectionOutput(Output):
+    def __init__(self, time, data, labels=None):
+        """
+        Connection output data.
+
+        Parameters
+        ----------
+        time : scalar
+            Time step (in seconds).
+        data : dict
+            Data arrays.
+        labels : sequence of str or None, default, None
+            Labels of connections.
+        
+        """
+        super().__init__(time, data, labels)
+
+    def __getitem__(self, islice):
+        """
+        Slice connection output.
+        
+        Parameters
+        ----------
+        islice : int, str, slice, sequence of int or sequence of str
+            Indices or labels of connections to slice.
+
+        Returns
+        -------
+        dict or :class:`toughio.ConnectionOutput`
+            Sliced connection outputs.
+        
+        """
+        if self.labels is None:
+            raise AttributeError()
+        
+        if np.ndim(islice) == 0:
+            if isinstance(islice, str):
+                islice = [
+                    i for i, (label1, label2) in enumerate(self.labels)
+                    if label1 == islice or label2 == islice
+                ]
+
+            elif isinstance(islice, slice):
+                islice = np.arange(self.n_data)[islice]
+
+            else:
+                return {k: v[islice] for k, v in self.data.items()}
+        
+        elif np.ndim(islice) <= 2:
+            islice = [self.index(*i) if np.ndim(i) == 1 else i for i in islice]
+
+        else:
+            raise ValueError()
+        
+        return ConnectionOutput(
+            self.time,
+            {k: v[islice] for k, v in self.data.items()},
+            [self._labels[i] for i in islice],
+        )
+
+    def index(self, label1, label2):
+        """
+        Get index of connection.
+        
+        Parameters
+        ----------
+        label1 : str
+            Label of first element of connection.
+        label2 : str
+            Label of second element of connection.
+
+        Returns
+        -------
+        int
+            Index of connection.
+        
+        """
+        super().index()
+        labels = ["".join(label) for label in self.labels]
+
+        return labels.index(f"{label1}{label2}")
+
+
+def to_output(file_type, labels_order, headers, times, labels, data):
+    """Helper function to create output data objects."""
+    outputs = []
+    for time, labels_, data_ in zip(times, labels, data):
+        kwargs = {
+            "time": time,
+            "labels": labels_ if labels_ is not None and len(labels_) else None,
+            "data": {k: v for k, v in zip(headers, np.transpose(data_))},
+        }
+
+        output = ElementOutput(**kwargs) if file_type == "element" else ConnectionOutput(**kwargs)
+        outputs.append(output)
 
     # Some older versions of TOUGH3 have duplicate connection outputs when running in parallel
     # Fix the outputs here by summing the duplicate connections
@@ -43,36 +263,14 @@ def to_output(file_type, labels_order, headers, times, labels, variables):
                 "Found duplicate connections. Fixing outputs by summing duplicate connections."
             )
 
-            outputs = [
-                Output(
-                    output.type,
-                    output.time,
-                    np.array(list(connections)),
-                    {
-                        k: np.array([v[idx].sum() for idx in connections.values()])
-                        for k, v in output.data.items()
-                    },
-                )
-                for output in outputs
-            ]
-
-    return (
-        [reorder_labels(out, labels_order) for out in outputs]
-        if labels_order is not None and file_type == "element"
-        else outputs
-    )
-
-
-def reorder_labels(data, labels):
-    """Reorder output or save cell data according to input labels."""
-    if len(data.labels) != len(labels):
-        raise ValueError()
-
-    mapper = {k: v for v, k in enumerate(data.labels)}
-    idx = [mapper[label] for label in labels]
-    data.labels[:] = data.labels[idx]
-
-    for k, v in data.data.items():
-        data.data[k] = v[idx]
-
-    return data
+            for output in outputs:
+                output.labels = list(connections)
+                output.data = {
+                    k: np.array([v[idx].sum() for idx in connections.values()])
+                    for k, v in output.data.items()
+                }
+
+    if file_type == "element" and labels_order is not None:
+        outputs = [output[labels_order] for output in outputs]
+
+    return outputs
diff --git a/toughio/_io/output/_helpers.py b/toughio/_io/output/_helpers.py
index 41b9e9ce..8ebc4dbc 100644
--- a/toughio/_io/output/_helpers.py
+++ b/toughio/_io/output/_helpers.py
@@ -55,17 +55,17 @@ def read(
     ----------
     filename : str, pathlike or buffer
         Input file name or buffer.
-    file_format : str ('csv', 'petrasim', 'save', 'tecplot', 'tough') or None, optional, default None
+    file_format : {'csv', 'petrasim', 'save', 'tecplot', 'tough'} or None, optional, default None
         Input file format.
-    labels_order : list of array_like or None, optional, default None
+    labels_order : sequence of array_like or None, optional, default None
         List of labels. If None, output will be assumed ordered.
     connection : bool, optional, default False
         Only for standard TOUGH output file. If `True`, return data related to connections.
 
     Returns
     -------
-    namedtuple or list of namedtuple
-        namedtuple (type, format, time, labels, data) or list of namedtuple for each time step.
+    :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+        Output data for each time step.
 
     """
     if not (
@@ -104,9 +104,9 @@ def write(filename, output, file_format=None, **kwargs):
     ----------
     filename : str, pathlike or buffer
         Output file name or buffer.
-    output : namedtuple or list of namedtuple
-        namedtuple (type, format, time, labels, data) or list of namedtuple for each time step to export.
-    file_format : str ('csv', 'petrasim', 'tecplot') or None, optional, default None
+    output : :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+        Output data to export for each time step.
+    file_format : {'csv', 'petrasim', 'tecplot'} or None, optional, default None
         Output file format.
 
     Other Parameters
@@ -139,31 +139,47 @@ def get_output_type(filename):
         if not line:
             line = f.readline().strip()
             if line.startswith("@@@@@"):
-                return "element", "tough"
+                file_format = "tough"
+                file_type = "element"
 
             else:
                 raise ValueError()
 
         elif line.startswith("1      @@@@@"):
-            return "element", "tough"
+            file_format = "tough"
+            file_type = "element"
 
         elif line.startswith("INCON"):
-            return "element", "save"
+            file_format = "save"
+            file_type = "element"
 
         elif "=" in line:
-            return "element", "tecplot"
+            file_format = "tecplot"
+            file_type = (
+                "connection"
+                if "HEAT" in line or "FLOW" in line or "VEL" in line
+                else "element"
+            )
 
         elif line.startswith("TIME"):
-            return "element", "petrasim"
+            file_format = "petrasim"
+            file_type = "connection" if "ELEM1" in line else "element"
 
         else:
             header = line.split(",")[0].replace('"', "").strip()
+            file_format = "csv"
+            file_type = "connection" if header == "ELEM1" else "element"
 
             if header == "ELEM":
-                return "element", "csv"
+                file_format = "csv"
+                file_type = "element"
 
             elif header == "ELEM1":
-                return "connection", "csv"
+                file_format = "csv"
+                file_type = "connection"
 
             else:
-                return "element", None
+                file_format = "element"
+                file_type = None
+            
+    return file_type, file_format
diff --git a/toughio/_io/output/csv/_csv.py b/toughio/_io/output/csv/_csv.py
index db4ef5c6..cdfc3538 100644
--- a/toughio/_io/output/csv/_csv.py
+++ b/toughio/_io/output/csv/_csv.py
@@ -1,7 +1,5 @@
-import numpy as np
-
 from ...._common import open_file
-from .._common import to_output
+from .._common import to_output, ElementOutput
 
 __all__ = [
     "read",
@@ -42,70 +40,67 @@ def read(filename, file_type, labels_order=None):
         Input file name or buffer.
     file_type : str
         Input file type.
-    labels_order : list of array_like
+    labels_order : sequence of array_like
         List of labels. If None, output will be assumed ordered.
 
     Returns
     -------
-    namedtuple or list of namedtuple
-        namedtuple (type, format, time, labels, data) or list of namedtuple for each time step.
+    :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+        Output data for each time step.
 
     """
     with open_file(filename, "r") as f:
-        headers, times, variables = _read_csv(f, file_type)
-
-        ilab = 1 if file_type == "element" else 2
-        headers = headers[ilab:]
-        labels = [[v[:ilab] for v in variable] for variable in variables]
-        labels = (
-            [[l[0] for l in label] for label in labels]
-            if file_type == "element"
-            else labels
-        )
-        variables = np.array([[v[ilab:] for v in variable] for variable in variables])
+        headers, times, labels, data = _read_csv(f, file_type)
 
-    return to_output(file_type, labels_order, headers, times, labels, variables)
+        return to_output(file_type, labels_order, headers, times, labels, data)
 
 
 def _read_csv(f, file_type):
     """Read CSV table."""
+    # Label index
+    ilab = 1 if file_type == "element" else 2
+
     # Read header
     line = f.readline().replace('"', "")
-    headers = [l.strip() for l in line.split(",")]
+    headers = [l.strip() for l in line.split(",")[ilab:]]
 
     # Skip second line (unit)
     line = f.readline()
 
     # Check third line (does it start with TIME?)
     line = f.readline()
-    single = not line.startswith('"TIME')
+    single = not line.startswith('"TIME [sec]')
 
     # Read data
     if single:
-        times, variables = [None], [[]]
+        times, labels, data = [None], [[]], [[]]
+
     else:
-        times, variables = [], []
+        times, labels, data = [], [], []
 
-    line = line.replace('"', "").strip()
-    ilab = 1 if file_type == "element" else 2
     while line:
         line = line.split(",")
 
         # Time step
-        if line[0].startswith("TIME"):
-            line = line[0].split()
+        if line[0].startswith('"TIME [sec]'):
+            line = line[0].replace('"', "").split()
             times.append(float(line[-1]))
-            variables.append([])
+            labels.append([])
+            data.append([])
 
         # Output
         else:
-            tmp = [l.strip() for l in line[:ilab]]
-            tmp += [float(l.strip()) for l in line[ilab:]]
-            variables[-1].append(tmp)
+            if ilab == 1:
+                labels[-1].append(line[0].replace('"', "").strip())
+
+            else:
+                labels[-1].append([l.replace('"', "").strip() for l in line[:ilab]])
+            
+            data[-1].append([float(l.strip()) for l in line[ilab:]])
 
-        line = f.readline().strip().replace('"', "")
+        line = f.readline()
 
-    return headers, times, variables
+    return headers, times, labels, data
 
 
 def write(filename, output, unit=None):
@@ -124,7 +119,7 @@ def write(filename, output, unit=None):
         raise TypeError()
 
     out = output[-1]
-    headers = ["ELEM"] if out.type == "element" else ["ELEM1", "ELEM2"]
+    headers = ["ELEM"] if isinstance(out, ElementOutput) else ["ELEM1", "ELEM2"]
     headers += ["X"] if "X" in out.data else []
     headers += ["Y"] if "Y" in out.data else []
     headers += ["Z"] if "Z" in out.data else []
diff --git a/toughio/_io/output/petrasim/_petrasim.py b/toughio/_io/output/petrasim/_petrasim.py
index 45802522..18fb39cf 100644
--- a/toughio/_io/output/petrasim/_petrasim.py
+++ b/toughio/_io/output/petrasim/_petrasim.py
@@ -1,7 +1,7 @@
 import numpy as np
 
 from ...._common import open_file
-from .._common import to_output
+from .._common import to_output, ElementOutput
 
 __all__ = [
     "read",
@@ -19,45 +19,61 @@ def read(filename, file_type, labels_order=None):
         Input file name or buffer.
     file_type : str
         Input file type.
-    labels_order : list of array_like
+    labels_order : sequence of array_like
         List of labels. If None, output will be assumed ordered.
 
     Returns
     -------
-    namedtuple or list of namedtuple
-        namedtuple (type, format, time, labels, data) or list of namedtuple for each time step.
+    :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+        Output data for each time step.
 
     """
     with open_file(filename, "r") as f:
+        # Label index
+        ilab = 3 if file_type == "element" else 4
+
         # Headers
         line = f.readline().strip()
-        headers = [header.strip() for header in line.split(",")[3:]]
+        headers = [header.strip() for header in line.split(",")[ilab:]]
 
         # Data
-        times, elements, data = [], [], []
+        count, tcur, offset = 0, None, []
+        times, labels, data = [], [], []
+
         while True:
             line = f.readline().strip()
 
             if line:
                 line = line.split(",")
-                times.append(float(line[0]))
-                elements.append(line[1].strip())
-                data.append([float(x) for x in line[3:]])
+
+                if line[0] != tcur:
+                    tcur = line[0]
+                    offset.append(count)
+                    times.append(float(tcur))
+
+                if file_type == "element":
+                    labels.append(line[1].strip())
+
+                else:
+                    labels.append([line[1].strip(), line[2].strip()])
+
+                data.append([float(x) for x in line[ilab:]])
+
             else:
                 break
 
-    times = np.array(times)
-    elements = np.array(elements)
-    data = np.array(data)
+            count += 1
 
-    labels, unique_times, variables = [], [], []
-    for time in np.unique(times):
-        idx = times == time
-        labels.append(elements[idx])
-        unique_times.append(time)
-        variables.append(data[idx])
+        offset.append(count)
 
-    return to_output(file_type, labels_order, headers, unique_times, labels, variables)
+    return to_output(
+        file_type,
+        labels_order,
+        headers,
+        times,
+        [labels[i1:i2] for i1, i2 in zip(offset[:-1], offset[1:])],
+        [data[i1:i2] for i1, i2 in zip(offset[:-1], offset[1:])],
+    )
 
 
 def write(filename, output):
@@ -81,20 +97,33 @@ def write(filename, output):
 
     with open_file(filename, "w") as f:
         # Headers
+        headers_ = (
+            ["TIME [sec]", "ELEM", "INDEX"]
+            if isinstance(out, ElementOutput)
+            else ["TIME [sec]", "ELEM1", "ELEM2", "INDEX"]
+        )
         record = ",".join(
-            f"{header:>18}" for header in ["TIME [sec]", "ELEM", "INDEX"] + headers
+            f"{header:>18}" for header in headers_ + headers
         )
         f.write(f"{record}\n")
 
         # Data
         for out in output:
             data = np.transpose([out.data[k] for k in headers])
-            formats = ["{:20.12e}", "{:>18}", "{:20d}"]
+            formats = (
+                ["{:20.12e}", "{:>18}", "{:20d}"]
+                if isinstance(out, ElementOutput)
+                else ["{:20.12e}", "{:>18}", "{:>18}", "{:20d}"]
+            )
             formats += ["{:20.12e}"] * len(out.data)
 
             i = 0
             for d in data:
-                tmp = [out.time, out.labels[i], i + 1]
+                tmp = (
+                    [out.time, out.labels[i], i + 1]
+                    if isinstance(out, ElementOutput)
+                    else [out.time, *out.labels[i], i + 1]
+                )
                 tmp += [x for x in d]
                 record = ",".join(fmt.format(x) for fmt, x in zip(formats, tmp))
                 f.write(f"{record}\n")
diff --git a/toughio/_io/output/save/_save.py b/toughio/_io/output/save/_save.py
index 8a98b53e..a7b68649 100644
--- a/toughio/_io/output/save/_save.py
+++ b/toughio/_io/output/save/_save.py
@@ -1,14 +1,14 @@
 import numpy as np
 
 from ...input import tough
-from .._common import Output, reorder_labels
+from .._common import ElementOutput
 
 __all__ = [
     "read",
 ]
 
 
-def read(filename, file_type, labels_order=None):
+def read(filename, file_type=None, labels_order=None):
     """
     Read SAVE file.
 
@@ -16,27 +16,25 @@ def read(filename, file_type, labels_order=None):
     ----------
     filename : str, pathlike or buffer
         Input file name or buffer.
-    file_type : str
+    file_type : str or None, default, None
         Input file type.
     labels_order : list of array_like
         List of labels. If None, output will be assumed ordered.
 
     Returns
     -------
-    namedtuple or list of namedtuple
-        namedtuple (type, format, time, labels, data) or list of namedtuple for each time step.
+    :class:`toughio.ElementOutput`
+        Output data.
 
     """
     parameters = tough.read(filename)
-
-    labels = list(parameters["initial_conditions"])
-    variables = [v["values"] for v in parameters["initial_conditions"].values()]
-
-    data = {f"X{i + 1}": x for i, x in enumerate(np.transpose(variables))}
-
+    
+    data = [v["values"] for v in parameters["initial_conditions"].values()]
+    data = {f"X{i + 1}": x for i, x in enumerate(np.transpose(data))}
     data["porosity"] = np.array(
         [v["porosity"] for v in parameters["initial_conditions"].values()]
     )
+    labels = list(parameters["initial_conditions"])
 
     userx = [
         v["userx"] for v in parameters["initial_conditions"].values() if "userx" in v
@@ -44,6 +42,12 @@ def read(filename, file_type, labels_order=None):
     if userx:
         data["userx"] = np.array(userx)
 
-    labels_order = labels_order if labels_order else parameters["initial_conditions"]
-    output = Output(file_type, None, np.array(labels), data)
-    return reorder_labels(output, labels_order)
+    try:
+        time = float(parameters["end_comments"][1].split()[-1])
+
+    except Exception:
+        time = None
+
+    output = ElementOutput(time, data, labels)
+
+    return output
diff --git a/toughio/_io/output/tecplot/_tecplot.py b/toughio/_io/output/tecplot/_tecplot.py
index 54aa2953..9d493c8e 100644
--- a/toughio/_io/output/tecplot/_tecplot.py
+++ b/toughio/_io/output/tecplot/_tecplot.py
@@ -38,27 +38,27 @@ def read(filename, file_type, labels_order=None):
         Input file name or buffer.
     file_type : str
         Input file type.
-    labels_order : list of array_like
+    labels_order : sequence of array_like
         List of labels. If None, output will be assumed ordered.
 
     Returns
     -------
-    namedtuple or list of namedtuple
-        namedtuple (type, format, time, labels, data) or list of namedtuple for each time step.
+    :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+        Output data for each time step.
 
     """
     with open_file(filename, "r") as f:
         headers, zones = read_buffer(f)
 
-    times, labels, variables = [], [], []
+    times, labels, data = [], [], []
     for zone in zones:
         time = float(zone["title"].split()[0]) if "title" in zone else None
 
         times.append(time)
         labels.append([])
-        variables.append(zone["data"])
+        data.append(zone["data"])
 
-    return to_output(file_type, labels_order, headers, times, labels, variables)
+    return to_output(file_type, labels_order, headers, times, labels, data)
 
 
 def read_buffer(f):
diff --git a/toughio/_io/output/tough/_tough.py b/toughio/_io/output/tough/_tough.py
index 939ece99..48773698 100644
--- a/toughio/_io/output/tough/_tough.py
+++ b/toughio/_io/output/tough/_tough.py
@@ -21,20 +21,20 @@ def read(filename, file_type, labels_order=None):
         Input file name or buffer.
     file_type : str
         Input file type.
-    labels_order : list of array_like
+    labels_order : sequence of array_like
         List of labels. If None, output will be assumed ordered.
 
     Returns
     -------
-    namedtuple or list of namedtuple
-        namedtuple (type, format, time, labels, data) or list of namedtuple for each time step.
+    :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+        Output data for each time step.
 
     """
     with open_file(filename, "r") as f:
-        headers, times, variables = _read_table(f, file_type)
+        headers, times, data = _read_table(f, file_type)
 
         # Postprocess labels
-        labels = [v[0].lstrip() for v in variables[0]]
+        labels = [v[0].lstrip() for v in data[0]]
 
         if file_type == "element":
             label_length = max(len(label) for label in labels)
@@ -74,10 +74,10 @@ def read(filename, file_type, labels_order=None):
 
         ilab = 1 if file_type == "element" else 2
         headers = headers[ilab + 1 :]
-        labels = [labels.copy() for _ in variables]
-        variables = np.array([[v[2:] for v in variable] for variable in variables])
+        labels = [labels.copy() for _ in data]
+        data = np.array([[v[2:] for v in data] for data in data])
 
-    return to_output(file_type, labels_order, headers, times, labels, variables)
+    return to_output(file_type, labels_order, headers, times, labels, data)
 
 
 def _read_table(f, file_type):
@@ -85,7 +85,7 @@ def _read_table(f, file_type):
     labels_key = "ELEM." if file_type == "element" else "ELEM1"
 
     first = True
-    times, variables = [], []
+    times, data = [], []
     for line in f:
         line = line.strip()
 
@@ -94,7 +94,7 @@ def _read_table(f, file_type):
             # Read time step in following line
             line = next(f).strip()
             times.append(float(line.split()[0]))
-            variables.append([])
+            data.append([])
 
             # Look for "ELEM." or "ELEM1"
             while True:
@@ -172,13 +172,13 @@ def _read_table(f, file_type):
                             first = False
 
                     tmp += reader(line)
-                    variables[-1].append([x for x in tmp if x is not None])
+                    data[-1].append([x for x in tmp if x is not None])
 
                 line = next(f)
                 if line[1:].startswith("@@@@@"):
                     break
 
-    return headers, times, variables
+    return headers, times, data
 
 
 def to_float(x):
diff --git a/toughio/_mesh/_mesh.py b/toughio/_mesh/_mesh.py
index a7e8c371..33472c4b 100644
--- a/toughio/_mesh/_mesh.py
+++ b/toughio/_mesh/_mesh.py
@@ -26,7 +26,7 @@
 CellBlock = collections.namedtuple("CellBlock", ["type", "data"])
 
 
-class Mesh(object):
+class Mesh:
     def __init__(
         self,
         points,
@@ -479,18 +479,18 @@ def read_output(
 
         Parameters
         ----------
-        file_or_output : str, pathlike, buffer, namedtuple or list of namedtuple
+        file_or_output : str, pathlike, buffer, :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
             Input file name or buffer, or output data.
         time_step : int, optional, default -1
             Data for given time step to import. Default is last time step.
-        labels_order : list of array_like or None, optional, default None
+        labels_order : sequence of array_like or None, optional, default None
             List of labels. If None, output will be assumed ordered.
         connection : bool, optional, default False
             Only for standard TOUGH output file. If `True`, read data related to connections.
 
         """
         from .. import read_output
-        from .._io.output._common import Output, reorder_labels
+        from .._io.output._common import ElementOutput, ConnectionOutput
 
         if not isinstance(time_step, int):
             raise TypeError()
@@ -501,19 +501,19 @@ def read_output(
         else:
             out = file_or_output
 
-        if not isinstance(out, Output):
+        if not isinstance(out, (ElementOutput, ConnectionOutput)):
             if not (-len(out) <= time_step < len(out)):
                 raise ValueError()
 
             out = out[time_step]
 
-        if out.type == "element":
+        if isinstance(out, ElementOutput):
             if labels_order is not None:
-                out = reorder_labels(out, labels_order)
+                out = out[labels_order]
 
             self.cell_data.update(out.data)
 
-        elif out.type == "connection":
+        elif isinstance(out, ConnectionOutput):
             centers = self.centers
             labels_map = {k: v for v, k in enumerate(self.labels)}
 

From c19a339e0f4c0bcba9251b5f1504ea4221b7e1c6 Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Fri, 10 May 2024 15:29:44 +0200
Subject: [PATCH 02/15] add option time_steps

---
 toughio/_io/output/_helpers.py           |  5 ++-
 toughio/_io/output/csv/_csv.py           | 46 ++++++++++++++++----
 toughio/_io/output/petrasim/_petrasim.py | 46 +++++++++++++++-----
 toughio/_io/output/save/_save.py         |  4 +-
 toughio/_io/output/tecplot/_tecplot.py   | 53 +++++++++++++++++++-----
 toughio/_io/output/tough/_tough.py       | 37 +++++++++++++++--
 toughio/_mesh/_mesh.py                   |  2 +-
 7 files changed, 156 insertions(+), 37 deletions(-)

diff --git a/toughio/_io/output/_helpers.py b/toughio/_io/output/_helpers.py
index 8ebc4dbc..0146881d 100644
--- a/toughio/_io/output/_helpers.py
+++ b/toughio/_io/output/_helpers.py
@@ -46,6 +46,7 @@ def read(
     filename,
     file_format=None,
     labels_order=None,
+    time_steps=None,
     connection=False,
 ):
     """
@@ -59,6 +60,8 @@ def read(
         Input file format.
     labels_order : sequence of array_like or None, optional, default None
         List of labels. If None, output will be assumed ordered.
+    time_steps : int or sequence of int
+        List of time steps to read. If None, all time steps will be read.
     connection : bool, optional, default False
         Only for standard TOUGH output file. If `True`, return data related to connections.
 
@@ -93,7 +96,7 @@ def read(
     if connection:
         file_type = "connection" if connection else "element"
 
-    return _reader_map[file_format](filename, file_type, labels_order)
+    return _reader_map[file_format](filename, file_type, labels_order, time_steps)
 
 
 def write(filename, output, file_format=None, **kwargs):
diff --git a/toughio/_io/output/csv/_csv.py b/toughio/_io/output/csv/_csv.py
index cdfc3538..bbb4ff99 100644
--- a/toughio/_io/output/csv/_csv.py
+++ b/toughio/_io/output/csv/_csv.py
@@ -30,7 +30,7 @@
 }
 
 
-def read(filename, file_type, labels_order=None):
+def read(filename, file_type, labels_order=None, time_steps=None):
     """
     Read OUTPUT_{ELEME, CONNE}.csv.
 
@@ -42,20 +42,32 @@ def read(filename, file_type, labels_order=None):
         Input file type.
     labels_order : sequence of array_like
         List of labels. If None, output will be assumed ordered.
+    time_steps : int or sequence of int
+        List of time steps to read. If None, all time steps will be read.
 
     Returns
     -------
-    :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+    sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
         Output data for each time step.
 
     """
+    if time_steps is not None:
+        if isinstance(time_steps, int):
+            time_steps = [time_steps]
+
+        if any(i < 0 for i in time_steps):
+            n_steps = _count_time_steps(filename)
+            time_steps = [i if i >= 0 else n_steps + i for i in time_steps]
+        
+        time_steps = set(time_steps)
+
     with open_file(filename, "r") as f:
-        headers, times, labels, data = _read_csv(f, file_type)
+        headers, times, labels, data = _read_csv(f, file_type, time_steps)
 
         return to_output(file_type, labels_order, headers, times, labels, data)
 
 
-def _read_csv(f, file_type):
+def _read_csv(f, file_type, time_steps=None):
     """Read CSV table."""
     # Label index
     ilab = 1 if file_type == "element" else 2
@@ -73,9 +85,11 @@ def _read_csv(f, file_type):
 
     # Read data
     if single:
+        t_step = 0
         times, labels, data = [None], [[]], [[]]
 
     else:
+        t_step = -1
         times, labels, data = [], [], []
 
     while line:
@@ -83,13 +97,16 @@ def _read_csv(f, file_type):
 
         # Time step
         if line[0].startswith('"TIME [sec]'):
-            line = line[0].replace('"', "").split()
-            times.append(float(line[-1]))
-            labels.append([])
-            data.append([])
+            t_step += 1
+
+            if time_steps is None or t_step in time_steps:
+                line = line[0].replace('"', "").split()
+                times.append(float(line[-1]))
+                labels.append([])
+                data.append([])
 
         # Output
-        else:
+        elif time_steps is None or t_step in time_steps:
             if ilab == 1:
                 labels[-1].append(line[0].replace('"', "").strip())
 
@@ -163,3 +180,14 @@ def _write_csv(f, output, headers, unit=None):
                 + "\n"
             )
             f.write(record)
+
+
+def _count_time_steps(filename):
+    """Count the number of time steps."""
+    with open_file(filename, "r") as f:
+        count = 0
+
+        for line in f:
+            count += int(line.startswith('"TIME [sec]'))
+
+    return count
diff --git a/toughio/_io/output/petrasim/_petrasim.py b/toughio/_io/output/petrasim/_petrasim.py
index 18fb39cf..71b2be6e 100644
--- a/toughio/_io/output/petrasim/_petrasim.py
+++ b/toughio/_io/output/petrasim/_petrasim.py
@@ -9,7 +9,7 @@
 ]
 
 
-def read(filename, file_type, labels_order=None):
+def read(filename, file_type, labels_order=None, time_steps=None):
     """
     Read Petrasim OUTPUT_ELEME.csv.
 
@@ -21,13 +21,25 @@ def read(filename, file_type, labels_order=None):
         Input file type.
     labels_order : sequence of array_like
         List of labels. If None, output will be assumed ordered.
+    time_steps : int or sequence of int
+        List of time steps to read. If None, all time steps will be read.
 
     Returns
     -------
-    :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+    sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
         Output data for each time step.
 
     """
+    if time_steps is not None:
+        if isinstance(time_steps, int):
+            time_steps = [time_steps]
+
+        if any(i < 0 for i in time_steps):
+            n_steps = _count_time_steps(filename)
+            time_steps = [i if i >= 0 else n_steps + i for i in time_steps]
+        
+        time_steps = set(time_steps)
+
     with open_file(filename, "r") as f:
         # Label index
         ilab = 3 if file_type == "element" else 4
@@ -37,6 +49,7 @@ def read(filename, file_type, labels_order=None):
         headers = [header.strip() for header in line.split(",")[ilab:]]
 
         # Data
+        t_step = -1
         count, tcur, offset = 0, None, []
         times, labels, data = [], [], []
 
@@ -47,23 +60,26 @@ def read(filename, file_type, labels_order=None):
                 line = line.split(",")
 
                 if line[0] != tcur:
+                    t_step += 1
                     tcur = line[0]
-                    offset.append(count)
-                    times.append(float(tcur))
 
-                if file_type == "element":
-                    labels.append(line[1].strip())
+                    if time_steps is None or t_step in time_steps:
+                        offset.append(count)
+                        times.append(float(tcur))
 
-                else:
-                    labels.append([line[1].strip(), line[2].strip()])
+                if time_steps is None or t_step in time_steps:
+                    if file_type == "element":
+                        labels.append(line[1].strip())
 
-                data.append([float(x) for x in line[ilab:]])
+                    else:
+                        labels.append([line[1].strip(), line[2].strip()])
+
+                    data.append([float(x) for x in line[ilab:]])
+                    count += 1
 
             else:
                 break
 
-            count += 1
-
         offset.append(count)
 
     return to_output(
@@ -128,3 +144,11 @@ def write(filename, output):
                 record = ",".join(fmt.format(x) for fmt, x in zip(formats, tmp))
                 f.write(f"{record}\n")
                 i += 1
+
+
+def _count_time_steps(filename):
+    """Count the number of time steps."""
+    with open_file(filename, "r") as f:
+        x = np.genfromtxt(f, delimiter=",", skip_header=1, usecols=0)
+
+    return np.unique(x).size
diff --git a/toughio/_io/output/save/_save.py b/toughio/_io/output/save/_save.py
index a7b68649..e1191fe8 100644
--- a/toughio/_io/output/save/_save.py
+++ b/toughio/_io/output/save/_save.py
@@ -8,7 +8,7 @@
 ]
 
 
-def read(filename, file_type=None, labels_order=None):
+def read(filename, file_type=None, labels_order=None, time_steps=None):
     """
     Read SAVE file.
 
@@ -20,6 +20,8 @@ def read(filename, file_type=None, labels_order=None):
         Input file type.
     labels_order : list of array_like
         List of labels. If None, output will be assumed ordered.
+    time_steps : int or sequence of int
+        List of time steps to read. If None, all time steps will be read.
 
     Returns
     -------
diff --git a/toughio/_io/output/tecplot/_tecplot.py b/toughio/_io/output/tecplot/_tecplot.py
index 9d493c8e..a9599a15 100644
--- a/toughio/_io/output/tecplot/_tecplot.py
+++ b/toughio/_io/output/tecplot/_tecplot.py
@@ -28,7 +28,7 @@
 }
 
 
-def read(filename, file_type, labels_order=None):
+def read(filename, file_type, labels_order=None, time_steps=None):
     """
     Read OUTPUT_ELEME.tec.
 
@@ -40,15 +40,27 @@ def read(filename, file_type, labels_order=None):
         Input file type.
     labels_order : sequence of array_like
         List of labels. If None, output will be assumed ordered.
+    time_steps : int or sequence of int
+        List of time steps to read. If None, all time steps will be read.
 
     Returns
     -------
-    :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+    sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
         Output data for each time step.
 
     """
+    if time_steps is not None:
+        if isinstance(time_steps, int):
+            time_steps = [time_steps]
+
+        if any(i < 0 for i in time_steps):
+            n_steps = _count_time_steps(filename)
+            time_steps = [i if i >= 0 else n_steps + i for i in time_steps]
+        
+        time_steps = set(time_steps)
+
     with open_file(filename, "r") as f:
-        headers, zones = read_buffer(f)
+        headers, zones = read_buffer(f, time_steps)
 
     times, labels, data = [], [], []
     for zone in zones:
@@ -61,11 +73,13 @@ def read(filename, file_type, labels_order=None):
     return to_output(file_type, labels_order, headers, times, labels, data)
 
 
-def read_buffer(f):
+def read_buffer(f, time_steps=None):
     """Read OUTPUT_ELEME.tec."""
     zones = []
 
     # Loop until end of file
+    t_step = -1
+
     while True:
         line = f.readline().strip()
 
@@ -79,14 +93,22 @@ def read_buffer(f):
 
             if "I" not in zone:
                 raise ValueError()
+            
+            else:
+                t_step += 1
+
+            if time_steps is None or t_step in time_steps:
+                # Read data
+                data = np.genfromtxt(f, max_rows=zone["I"])
 
-            # Read data
-            data = np.genfromtxt(f, max_rows=zone["I"])
+                # Output
+                tmp = {"data": data}
+                tmp["title"] = zone["T"]
+                zones.append(tmp)
 
-            # Output
-            tmp = {"data": data}
-            tmp["title"] = zone["T"]
-            zones.append(tmp)
+            else:
+                for _ in range(zone["I"]):
+                    _ = f.readline()
 
         elif not line:
             break
@@ -215,3 +237,14 @@ def _read_zone(line):
         zone["T"] = zone_title.strip()
 
     return zone
+
+
+def _count_time_steps(filename):
+    """Count the number of time steps."""
+    with open_file(filename, "r") as f:
+        count = 0
+
+        for line in f:
+            count += int(line.strip().upper().startswith("ZONE"))
+
+    return count
diff --git a/toughio/_io/output/tough/_tough.py b/toughio/_io/output/tough/_tough.py
index 48773698..390ba61b 100644
--- a/toughio/_io/output/tough/_tough.py
+++ b/toughio/_io/output/tough/_tough.py
@@ -11,7 +11,7 @@
 ]
 
 
-def read(filename, file_type, labels_order=None):
+def read(filename, file_type, labels_order=None, time_steps=None):
     """
     Read standard TOUGH output file.
 
@@ -23,15 +23,27 @@ def read(filename, file_type, labels_order=None):
         Input file type.
     labels_order : sequence of array_like
         List of labels. If None, output will be assumed ordered.
+    time_steps : int or sequence of int
+        List of time steps to read. If None, all time steps will be read.
 
     Returns
     -------
-    :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+    sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
         Output data for each time step.
 
     """
+    if time_steps is not None:
+        if isinstance(time_steps, int):
+            time_steps = [time_steps]
+
+        if any(i < 0 for i in time_steps):
+            n_steps = _count_time_steps(filename)
+            time_steps = [i if i >= 0 else n_steps + i for i in time_steps]
+        
+        time_steps = set(time_steps)
+
     with open_file(filename, "r") as f:
-        headers, times, data = _read_table(f, file_type)
+        headers, times, data = _read_table(f, file_type, time_steps)
 
         # Postprocess labels
         labels = [v[0].lstrip() for v in data[0]]
@@ -80,10 +92,11 @@ def read(filename, file_type, labels_order=None):
     return to_output(file_type, labels_order, headers, times, labels, data)
 
 
-def _read_table(f, file_type):
+def _read_table(f, file_type, time_steps=None):
     """Read data table for current time step."""
     labels_key = "ELEM." if file_type == "element" else "ELEM1"
 
+    t_step = -1
     first = True
     times, data = [], []
     for line in f:
@@ -91,6 +104,11 @@ def _read_table(f, file_type):
 
         # Look for "TOTAL TIME"
         if line.startswith("TOTAL TIME"):
+            t_step += 1
+
+            if not (time_steps is None or t_step in time_steps):
+                continue
+            
             # Read time step in following line
             line = next(f).strip()
             times.append(float(line.split()[0]))
@@ -190,3 +208,14 @@ def to_float(x):
 
     except ValueError:
         return np.nan
+
+
+def _count_time_steps(filename):
+    """Count the number of time steps."""
+    with open_file(filename, "r") as f:
+        count = 0
+
+        for line in f:
+            count += int(line.strip().startswith("TOTAL TIME"))
+
+    return count
diff --git a/toughio/_mesh/_mesh.py b/toughio/_mesh/_mesh.py
index 33472c4b..36a090da 100644
--- a/toughio/_mesh/_mesh.py
+++ b/toughio/_mesh/_mesh.py
@@ -496,7 +496,7 @@ def read_output(
             raise TypeError()
 
         if isinstance(file_or_output, str):
-            out = read_output(file_or_output, connection=connection)
+            out = read_output(file_or_output, time_steps=time_step, connection=connection)[0]
 
         else:
             out = file_or_output

From 2ef8bf3b885a7c70a627aa92204a7e3757e822ee Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Fri, 10 May 2024 16:05:20 +0200
Subject: [PATCH 03/15] return output instead of list if only one output

---
 tests/test_cli.py                        |  5 +++++
 tests/test_output.py                     |  1 +
 toughio/_cli/_export.py                  | 13 ++++++-------
 toughio/_io/output/_common.py            |  2 +-
 toughio/_io/output/csv/_csv.py           |  2 +-
 toughio/_io/output/petrasim/_petrasim.py |  2 +-
 toughio/_io/output/tecplot/_tecplot.py   |  2 +-
 toughio/_io/output/tough/_tough.py       |  2 +-
 8 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/tests/test_cli.py b/tests/test_cli.py
index 1c08dfc3..db7f45c9 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -166,18 +166,23 @@ def test_extract(file_format, split, connection):
         this_dir, "support_files", "outputs", f"{base_filename}.csv"
     )
     outputs_ref = toughio.read_output(filename_ref)
+    outputs_ref = outputs_ref if isinstance(outputs_ref, list) else [outputs_ref]
 
     if not split:
         outputs = toughio.read_output(output_filename, connection=connection)
+        outputs = outputs if isinstance(outputs, list) else [outputs]
 
         for output_ref, output in zip(outputs_ref, outputs):
             assert output_ref.time == output.time
             for k, v in output_ref.data.items():
                 assert helpers.allclose(v.mean(), output.data[k].mean(), atol=1.0e-2)
+
     else:
         filenames = glob.glob(os.path.join(tempdir, f"{base_filename}_*.csv"))
+
         for i, output_filename in enumerate(sorted(filenames)):
             outputs = toughio.read_output(output_filename)
+            outputs = outputs if isinstance(outputs, list) else [outputs]
 
             assert len(outputs) == 1
 
diff --git a/tests/test_output.py b/tests/test_output.py
index 7a0b0f9d..063f4965 100644
--- a/tests/test_output.py
+++ b/tests/test_output.py
@@ -110,6 +110,7 @@ def test_output(output_ref, file_format):
         reader_kws={},
     )
 
+    output = output if isinstance(output, list) else [output]
     output_ref = output_ref if isinstance(output_ref, list) else [output_ref]
     for out_ref, out in zip(output_ref, output):
         # Careful here, tecplot format has no label
diff --git a/toughio/_cli/_export.py b/toughio/_cli/_export.py
index fad8b53d..e1f6e9f8 100644
--- a/toughio/_cli/_export.py
+++ b/toughio/_cli/_export.py
@@ -44,18 +44,17 @@ def export(argv=None):
 
     # Read output file
     print(f"Reading file '{args.infile}' ...", end="")
+
     sys.stdout.flush()
-    output = read_output(args.infile)
     if args.file_format != "xdmf":
-        if args.time_step is not None:
-            if not (-len(output) <= args.time_step < len(output)):
-                raise ValueError("Inconsistent time step value.")
-            output = output[args.time_step]
-        else:
-            output = output[-1]
+        time_step = args.time_step if args.time_step is not None else -1
+        output = read_output(args.infile, time_steps=time_step)
         labels = output.labels
+
     else:
+        output = read_output(args.infile)
         labels = output[-1].labels
+
     print(" Done!")
 
     with_mesh = bool(args.mesh)
diff --git a/toughio/_io/output/_common.py b/toughio/_io/output/_common.py
index dc405960..f84c694c 100644
--- a/toughio/_io/output/_common.py
+++ b/toughio/_io/output/_common.py
@@ -273,4 +273,4 @@ def to_output(file_type, labels_order, headers, times, labels, data):
     if file_type == "element" and labels_order is not None:
         outputs = [output[labels_order] for output in outputs]
 
-    return outputs
+    return outputs[0] if len(outputs) == 1 else outputs
diff --git a/toughio/_io/output/csv/_csv.py b/toughio/_io/output/csv/_csv.py
index bbb4ff99..fc8581d3 100644
--- a/toughio/_io/output/csv/_csv.py
+++ b/toughio/_io/output/csv/_csv.py
@@ -47,7 +47,7 @@ def read(filename, file_type, labels_order=None, time_steps=None):
 
     Returns
     -------
-    sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+    :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
         Output data for each time step.
 
     """
diff --git a/toughio/_io/output/petrasim/_petrasim.py b/toughio/_io/output/petrasim/_petrasim.py
index 71b2be6e..c1fd1187 100644
--- a/toughio/_io/output/petrasim/_petrasim.py
+++ b/toughio/_io/output/petrasim/_petrasim.py
@@ -26,7 +26,7 @@ def read(filename, file_type, labels_order=None, time_steps=None):
 
     Returns
     -------
-    sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+    :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
         Output data for each time step.
 
     """
diff --git a/toughio/_io/output/tecplot/_tecplot.py b/toughio/_io/output/tecplot/_tecplot.py
index a9599a15..a204cb2b 100644
--- a/toughio/_io/output/tecplot/_tecplot.py
+++ b/toughio/_io/output/tecplot/_tecplot.py
@@ -45,7 +45,7 @@ def read(filename, file_type, labels_order=None, time_steps=None):
 
     Returns
     -------
-    sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+    :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
         Output data for each time step.
 
     """
diff --git a/toughio/_io/output/tough/_tough.py b/toughio/_io/output/tough/_tough.py
index 390ba61b..92089358 100644
--- a/toughio/_io/output/tough/_tough.py
+++ b/toughio/_io/output/tough/_tough.py
@@ -28,7 +28,7 @@ def read(filename, file_type, labels_order=None, time_steps=None):
 
     Returns
     -------
-    sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
+    :class:`toughio.ElementOutput`, :class:`toughio.ConnectionOutput`, sequence of :class:`toughio.ElementOutput` or sequence of :class:`toughio.ConnectionOutput`
         Output data for each time step.
 
     """

From 3c95061f90e1adef6987ba577e8e402b55102847 Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Fri, 10 May 2024 16:17:34 +0200
Subject: [PATCH 04/15] test option time_steps

---
 tests/test_output.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/tests/test_output.py b/tests/test_output.py
index 063f4965..f58eb449 100644
--- a/tests/test_output.py
+++ b/tests/test_output.py
@@ -117,6 +117,28 @@ def test_output(output_ref, file_format):
         helpers.allclose(out, out_ref)
 
 
+@pytest.mark.parametrize(
+    "filename",
+    [
+        "OUTPUT_ELEME.csv",
+        "OUTPUT_ELEME.tec",
+        "OUTPUT_ELEME_PETRASIM.csv",
+        "OUTPUT_CONNE.csv",
+    ],
+)
+def test_output_time_steps(filename):
+    this_dir = os.path.dirname(os.path.abspath(__file__))
+    filename = os.path.join(this_dir, "support_files", "outputs", filename)
+    outputs_ref = toughio.read_output(filename)
+
+    time_steps = [0, 2, -1]
+    outputs = toughio.read_output(filename, time_steps=time_steps)
+    outputs_ref = [outputs_ref[time_step] for time_step in time_steps]
+
+    for out_ref, out in zip(outputs_ref, outputs):
+        helpers.allclose(out, out_ref)
+
+
 def test_save():
     this_dir = os.path.dirname(os.path.abspath(__file__))
     filename = os.path.join(this_dir, "support_files", "outputs", "SAVE.out")

From 7ad5a6d4552c686117525766e990911533ad72c3 Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Fri, 10 May 2024 16:31:47 +0200
Subject: [PATCH 05/15] some optimization

---
 tests/test_output.py                     | 1 +
 toughio/_io/output/csv/_csv.py           | 3 +++
 toughio/_io/output/petrasim/_petrasim.py | 4 ++++
 toughio/_io/output/tecplot/_tecplot.py   | 3 +++
 toughio/_io/output/tough/_tough.py       | 3 +++
 5 files changed, 14 insertions(+)

diff --git a/tests/test_output.py b/tests/test_output.py
index f58eb449..8a0c7c0e 100644
--- a/tests/test_output.py
+++ b/tests/test_output.py
@@ -123,6 +123,7 @@ def test_output(output_ref, file_format):
         "OUTPUT_ELEME.csv",
         "OUTPUT_ELEME.tec",
         "OUTPUT_ELEME_PETRASIM.csv",
+        "OUTPUT.out",
         "OUTPUT_CONNE.csv",
     ],
 )
diff --git a/toughio/_io/output/csv/_csv.py b/toughio/_io/output/csv/_csv.py
index fc8581d3..6c1a9448 100644
--- a/toughio/_io/output/csv/_csv.py
+++ b/toughio/_io/output/csv/_csv.py
@@ -99,6 +99,9 @@ def _read_csv(f, file_type, time_steps=None):
         if line[0].startswith('"TIME [sec]'):
             t_step += 1
 
+            if time_steps is not None and t_step > max(time_steps):
+                break
+
             if time_steps is None or t_step in time_steps:
                 line = line[0].replace('"', "").split()
                 times.append(float(line[-1]))
diff --git a/toughio/_io/output/petrasim/_petrasim.py b/toughio/_io/output/petrasim/_petrasim.py
index c1fd1187..b8af3e75 100644
--- a/toughio/_io/output/petrasim/_petrasim.py
+++ b/toughio/_io/output/petrasim/_petrasim.py
@@ -61,6 +61,10 @@ def read(filename, file_type, labels_order=None, time_steps=None):
 
                 if line[0] != tcur:
                     t_step += 1
+
+                    if time_steps is not None and t_step > max(time_steps):
+                        break
+
                     tcur = line[0]
 
                     if time_steps is None or t_step in time_steps:
diff --git a/toughio/_io/output/tecplot/_tecplot.py b/toughio/_io/output/tecplot/_tecplot.py
index a204cb2b..ccb4af8a 100644
--- a/toughio/_io/output/tecplot/_tecplot.py
+++ b/toughio/_io/output/tecplot/_tecplot.py
@@ -97,6 +97,9 @@ def read_buffer(f, time_steps=None):
             else:
                 t_step += 1
 
+                if time_steps is not None and t_step > max(time_steps):
+                    break
+
             if time_steps is None or t_step in time_steps:
                 # Read data
                 data = np.genfromtxt(f, max_rows=zone["I"])
diff --git a/toughio/_io/output/tough/_tough.py b/toughio/_io/output/tough/_tough.py
index 92089358..100b2ac7 100644
--- a/toughio/_io/output/tough/_tough.py
+++ b/toughio/_io/output/tough/_tough.py
@@ -106,6 +106,9 @@ def _read_table(f, file_type, time_steps=None):
         if line.startswith("TOTAL TIME"):
             t_step += 1
 
+            if time_steps is not None and t_step > max(time_steps):
+                break
+
             if not (time_steps is None or t_step in time_steps):
                 continue
             

From 847212dbbc19693b32a71aa9c5e4db6c3b0095bf Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Sat, 18 May 2024 17:36:17 +0200
Subject: [PATCH 06/15] use second line parser when failing with default line
 parser

---
 toughio/_io/output/tough/_tough.py | 69 ++++++++++++------------------
 1 file changed, 28 insertions(+), 41 deletions(-)

diff --git a/toughio/_io/output/tough/_tough.py b/toughio/_io/output/tough/_tough.py
index 100b2ac7..f3d89b7c 100644
--- a/toughio/_io/output/tough/_tough.py
+++ b/toughio/_io/output/tough/_tough.py
@@ -3,7 +3,7 @@
 import numpy as np
 
 from ...._common import open_file
-from ..._common import read_record
+from ..._common import read_record, to_float
 from .._common import to_output
 
 __all__ = [
@@ -97,7 +97,6 @@ def _read_table(f, file_type, time_steps=None):
     labels_key = "ELEM." if file_type == "element" else "ELEM1"
 
     t_step = -1
-    first = True
     times, data = [], []
     for line in f:
         line = line.strip()
@@ -143,9 +142,12 @@ def _read_table(f, file_type, time_steps=None):
                     break
 
             # Loop until end of output block
+            reader = lambda line: [to_float(x) for x in line.split()]
+            reader2 = None
+
             while True:
                 if line[:nwsp].strip() and not line.strip().startswith("ELEM"):
-                    if first:
+                    if reader2 is None:
                         # Find first floating point
                         x = line.split()[-headers[::-1].index("INDEX")]
 
@@ -158,41 +160,37 @@ def _read_table(f, file_type, time_steps=None):
                     tmp = [line[:iend]]
                     line = line[iend:]
 
-                    if first:
-                        try:
-                            # Set line parser and try parsing first line
-                            reader = lambda line: [to_float(x) for x in line.split()]
-                            _ = reader(line)
+                    if reader2 is None:
+                        # Determine number of characters for index
+                        idx = line.replace("-", " ").split()[0]
+                        nidx = line.index(idx) + len(idx)
+                        ifmt = f"{nidx}s"
 
-                        except ValueError:
-                            # Determine number of characters for index
-                            idx = line.replace("-", " ").split()[0]
-                            nidx = line.index(idx) + len(idx)
-                            ifmt = f"{nidx}s"
+                        # Determine number of characters between two Es
+                        i1 = line.find("E")
+                        i2 = line.find("E", i1 + 1)
 
-                            # Determine number of characters between two Es
-                            i1 = line.find("E")
-                            i2 = line.find("E", i1 + 1)
+                        # Initialize data format
+                        fmt = [ifmt]
+                        if i2 >= 0:
+                            di = i2 - i1
+                            dfmt = f"{di}.{di - 7}e"
+                            fmt += 20 * [dfmt]  # Read 20 data columns at most
 
-                            # Initialize data format
-                            fmt = [ifmt]
-                            if i2 >= 0:
-                                di = i2 - i1
-                                dfmt = f"{di}.{di - 7}e"
-                                fmt += 20 * [dfmt]  # Read 20 data columns at most
+                        else:
+                            fmt += ["12.5e"]
 
-                            else:
-                                fmt += ["12.5e"]
+                        fmt = ",".join(fmt)
 
-                            fmt = ",".join(fmt)
+                        # Set second line parser
+                        reader2 = partial(read_record, fmt=fmt)
 
-                            # Set line parser
-                            reader = partial(read_record, fmt=fmt)
+                    try:
+                        tmp += reader(line)
 
-                        finally:
-                            first = False
+                    except ValueError:
+                        tmp += reader2(line)
 
-                    tmp += reader(line)
                     data[-1].append([x for x in tmp if x is not None])
 
                 line = next(f)
@@ -202,17 +200,6 @@ def _read_table(f, file_type, time_steps=None):
     return headers, times, data
 
 
-def to_float(x):
-    """Return np.nan if x cannot be converted."""
-    from ..._common import to_float as _to_float
-
-    try:
-        return _to_float(x)
-
-    except ValueError:
-        return np.nan
-
-
 def _count_time_steps(filename):
     """Count the number of time steps."""
     with open_file(filename, "r") as f:

From fad2464bc16a50c5a82ad5508e3a2e2eba746333 Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Sun, 19 May 2024 09:52:49 +0200
Subject: [PATCH 07/15] version bump

---
 toughio/VERSION | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/toughio/VERSION b/toughio/VERSION
index a32d5a6a..d19d0890 100644
--- a/toughio/VERSION
+++ b/toughio/VERSION
@@ -1 +1 @@
-1.14.2
\ No newline at end of file
+1.15.0
\ No newline at end of file

From caefb3bb624d4701e046d8b318d7b40f20bc08b9 Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Sun, 19 May 2024 09:55:45 +0200
Subject: [PATCH 08/15] invoke format

---
 toughio/__init__.py                      |  2 +-
 toughio/_io/__init__.py                  |  2 +-
 toughio/_io/h5/_write.py                 |  2 +-
 toughio/_io/output/__init__.py           |  2 +-
 toughio/_io/output/_common.py            | 63 +++++++++++++-----------
 toughio/_io/output/_helpers.py           |  2 +-
 toughio/_io/output/csv/_csv.py           |  6 +--
 toughio/_io/output/petrasim/_petrasim.py |  8 ++-
 toughio/_io/output/save/_save.py         |  2 +-
 toughio/_io/output/tecplot/_tecplot.py   |  4 +-
 toughio/_io/output/tough/_tough.py       |  4 +-
 toughio/_mesh/_mesh.py                   |  6 ++-
 12 files changed, 54 insertions(+), 49 deletions(-)

diff --git a/toughio/__init__.py b/toughio/__init__.py
index f68c4895..91893487 100644
--- a/toughio/__init__.py
+++ b/toughio/__init__.py
@@ -2,8 +2,8 @@
 from .__about__ import __version__
 from ._helpers import convert_labels
 from ._io import (
-    ElementOutput,
     ConnectionOutput,
+    ElementOutput,
     read_input,
     read_output,
     read_table,
diff --git a/toughio/_io/__init__.py b/toughio/_io/__init__.py
index c59fd595..9c50ea33 100644
--- a/toughio/_io/__init__.py
+++ b/toughio/_io/__init__.py
@@ -2,7 +2,7 @@
 from .input import read as read_input
 from .input import register as register_input
 from .input import write as write_input
-from .output import ElementOutput, ConnectionOutput
+from .output import ConnectionOutput, ElementOutput
 from .output import read as read_output
 from .output import register as register_output
 from .output import write as write_output
diff --git a/toughio/_io/h5/_write.py b/toughio/_io/h5/_write.py
index 16fbd833..4f9ffede 100644
--- a/toughio/_io/h5/_write.py
+++ b/toughio/_io/h5/_write.py
@@ -2,7 +2,7 @@
 
 import h5py
 
-from ..output import ElementOutput, ConnectionOutput
+from ..output import ConnectionOutput, ElementOutput
 from ..output import read as read_output
 from ..table import read as read_table
 
diff --git a/toughio/_io/output/__init__.py b/toughio/_io/output/__init__.py
index 40b9e3d4..b02cea87 100644
--- a/toughio/_io/output/__init__.py
+++ b/toughio/_io/output/__init__.py
@@ -1,5 +1,5 @@
 from . import csv, petrasim, save, tecplot, tough
-from ._common import ElementOutput, ConnectionOutput
+from ._common import ConnectionOutput, ElementOutput
 from ._helpers import read, register, write
 
 __all__ = [
diff --git a/toughio/_io/output/_common.py b/toughio/_io/output/_common.py
index f84c694c..b559e805 100644
--- a/toughio/_io/output/_common.py
+++ b/toughio/_io/output/_common.py
@@ -1,5 +1,5 @@
-from abc import ABC, abstractmethod
 import logging
+from abc import ABC, abstractmethod
 
 import numpy as np
 
@@ -24,13 +24,13 @@ def __init__(self, time, data, labels=None):
     def __getitem__(self):
         """Slice output."""
         pass
-    
+
     @abstractmethod
     def index(self):
         """Get index of element or connection."""
         if self.labels is None:
             raise AttributeError()
-        
+
     @property
     def n_data(self):
         """Return number of data points."""
@@ -40,31 +40,31 @@ def n_data(self):
     def time(self):
         """Return time step (in seconds)."""
         return self._time
-    
+
     @time.setter
     def time(self, value):
         self._time = value
-    
+
     @property
     def data(self):
         """Return data arrays."""
         return self._data
-    
+
     @data.setter
     def data(self, value):
         self._data = value
-    
+
     @property
     def labels(self):
         """Return labels."""
         return self._labels
-    
+
     @labels.setter
     def labels(self, value):
         if value is not None:
             if len(value) != self.n_data:
                 raise ValueError()
-            
+
             self._labels = list(value)
 
         else:
@@ -84,14 +84,14 @@ def __init__(self, time, data, labels=None):
             Data arrays.
         labels : sequence of str or None, default, None
             Labels of elements.
-        
+
         """
         super().__init__(time, data, labels)
 
     def __getitem__(self, islice):
         """
         Slice element output.
-        
+
         Parameters
         ----------
         islice : int, str, slice, sequence of int or sequence of str
@@ -101,26 +101,26 @@ def __getitem__(self, islice):
         -------
         dict or :class:`toughio.ElementOutput`
             Sliced element outputs.
-        
+
         """
         if self.labels is None:
             raise AttributeError()
-        
+
         if np.ndim(islice) == 0:
             if isinstance(islice, slice):
                 islice = np.arange(self.n_data)[islice]
 
             else:
                 islice = self.index(islice) if isinstance(islice, str) else islice
-            
+
                 return {k: v[islice] for k, v in self.data.items()}
-        
+
         elif np.ndim(islice) == 1:
             islice = [self.index(i) if isinstance(i, str) else i for i in islice]
 
         else:
             raise ValueError()
-        
+
         return ElementOutput(
             self.time,
             {k: v[islice] for k, v in self.data.items()},
@@ -130,7 +130,7 @@ def __getitem__(self, islice):
     def index(self, label):
         """
         Get index of element.
-        
+
         Parameters
         ----------
         label : str
@@ -140,10 +140,10 @@ def index(self, label):
         -------
         int
             Index of element.
-        
+
         """
         super().index()
-        
+
         return self.labels.index(label)
 
 
@@ -160,14 +160,14 @@ def __init__(self, time, data, labels=None):
             Data arrays.
         labels : sequence of str or None, default, None
             Labels of connections.
-        
+
         """
         super().__init__(time, data, labels)
 
     def __getitem__(self, islice):
         """
         Slice connection output.
-        
+
         Parameters
         ----------
         islice : int, str, slice, sequence of int or sequence of str
@@ -177,15 +177,16 @@ def __getitem__(self, islice):
         -------
         dict or :class:`toughio.ConnectionOutput`
             Sliced connection outputs.
-        
+
         """
         if self.labels is None:
             raise AttributeError()
-        
+
         if np.ndim(islice) == 0:
             if isinstance(islice, str):
                 islice = [
-                    i for i, (label1, label2) in enumerate(self.labels)
+                    i
+                    for i, (label1, label2) in enumerate(self.labels)
                     if label1 == islice or label2 == islice
                 ]
 
@@ -194,13 +195,13 @@ def __getitem__(self, islice):
 
             else:
                 return {k: v[islice] for k, v in self.data.items()}
-        
+
         elif np.ndim(islice) <= 2:
             islice = [self.index(*i) if np.ndim(i) == 1 else i for i in islice]
 
         else:
             raise ValueError()
-        
+
         return ConnectionOutput(
             self.time,
             {k: v[islice] for k, v in self.data.items()},
@@ -210,7 +211,7 @@ def __getitem__(self, islice):
     def index(self, label1, label2):
         """
         Get index of connection.
-        
+
         Parameters
         ----------
         label1 : str
@@ -222,7 +223,7 @@ def index(self, label1, label2):
         -------
         int
             Index of connection.
-        
+
         """
         super().index()
         labels = ["".join(label) for label in self.labels]
@@ -240,7 +241,11 @@ def to_output(file_type, labels_order, headers, times, labels, data):
             "data": {k: v for k, v in zip(headers, np.transpose(data_))},
         }
 
-        output = ElementOutput(**kwargs) if file_type == "element" else ConnectionOutput(**kwargs)
+        output = (
+            ElementOutput(**kwargs)
+            if file_type == "element"
+            else ConnectionOutput(**kwargs)
+        )
         outputs.append(output)
 
     # Some older versions of TOUGH3 have duplicate connection outputs when running in parallel
diff --git a/toughio/_io/output/_helpers.py b/toughio/_io/output/_helpers.py
index 0146881d..0d569ee1 100644
--- a/toughio/_io/output/_helpers.py
+++ b/toughio/_io/output/_helpers.py
@@ -184,5 +184,5 @@ def get_output_type(filename):
             else:
                 file_format = "element"
                 file_type = None
-            
+
     return file_type, file_format
diff --git a/toughio/_io/output/csv/_csv.py b/toughio/_io/output/csv/_csv.py
index 6c1a9448..b7ba3142 100644
--- a/toughio/_io/output/csv/_csv.py
+++ b/toughio/_io/output/csv/_csv.py
@@ -1,5 +1,5 @@
 from ...._common import open_file
-from .._common import to_output, ElementOutput
+from .._common import ElementOutput, to_output
 
 __all__ = [
     "read",
@@ -58,7 +58,7 @@ def read(filename, file_type, labels_order=None, time_steps=None):
         if any(i < 0 for i in time_steps):
             n_steps = _count_time_steps(filename)
             time_steps = [i if i >= 0 else n_steps + i for i in time_steps]
-        
+
         time_steps = set(time_steps)
 
     with open_file(filename, "r") as f:
@@ -115,7 +115,7 @@ def _read_csv(f, file_type, time_steps=None):
 
             else:
                 labels[-1].append([l.replace('"', "").strip() for l in line[:ilab]])
-            
+
             data[-1].append([float(l.strip()) for l in line[ilab:]])
 
         line = f.readline()
diff --git a/toughio/_io/output/petrasim/_petrasim.py b/toughio/_io/output/petrasim/_petrasim.py
index b8af3e75..b9d5b916 100644
--- a/toughio/_io/output/petrasim/_petrasim.py
+++ b/toughio/_io/output/petrasim/_petrasim.py
@@ -1,7 +1,7 @@
 import numpy as np
 
 from ...._common import open_file
-from .._common import to_output, ElementOutput
+from .._common import ElementOutput, to_output
 
 __all__ = [
     "read",
@@ -37,7 +37,7 @@ def read(filename, file_type, labels_order=None, time_steps=None):
         if any(i < 0 for i in time_steps):
             n_steps = _count_time_steps(filename)
             time_steps = [i if i >= 0 else n_steps + i for i in time_steps]
-        
+
         time_steps = set(time_steps)
 
     with open_file(filename, "r") as f:
@@ -122,9 +122,7 @@ def write(filename, output):
             if isinstance(out, ElementOutput)
             else ["TIME [sec]", "ELEM1", "ELEM2", "INDEX"]
         )
-        record = ",".join(
-            f"{header:>18}" for header in headers_ + headers
-        )
+        record = ",".join(f"{header:>18}" for header in headers_ + headers)
         f.write(f"{record}\n")
 
         # Data
diff --git a/toughio/_io/output/save/_save.py b/toughio/_io/output/save/_save.py
index e1191fe8..0400c970 100644
--- a/toughio/_io/output/save/_save.py
+++ b/toughio/_io/output/save/_save.py
@@ -30,7 +30,7 @@ def read(filename, file_type=None, labels_order=None, time_steps=None):
 
     """
     parameters = tough.read(filename)
-    
+
     data = [v["values"] for v in parameters["initial_conditions"].values()]
     data = {f"X{i + 1}": x for i, x in enumerate(np.transpose(data))}
     data["porosity"] = np.array(
diff --git a/toughio/_io/output/tecplot/_tecplot.py b/toughio/_io/output/tecplot/_tecplot.py
index ccb4af8a..768bd172 100644
--- a/toughio/_io/output/tecplot/_tecplot.py
+++ b/toughio/_io/output/tecplot/_tecplot.py
@@ -56,7 +56,7 @@ def read(filename, file_type, labels_order=None, time_steps=None):
         if any(i < 0 for i in time_steps):
             n_steps = _count_time_steps(filename)
             time_steps = [i if i >= 0 else n_steps + i for i in time_steps]
-        
+
         time_steps = set(time_steps)
 
     with open_file(filename, "r") as f:
@@ -93,7 +93,7 @@ def read_buffer(f, time_steps=None):
 
             if "I" not in zone:
                 raise ValueError()
-            
+
             else:
                 t_step += 1
 
diff --git a/toughio/_io/output/tough/_tough.py b/toughio/_io/output/tough/_tough.py
index f3d89b7c..5b10da92 100644
--- a/toughio/_io/output/tough/_tough.py
+++ b/toughio/_io/output/tough/_tough.py
@@ -39,7 +39,7 @@ def read(filename, file_type, labels_order=None, time_steps=None):
         if any(i < 0 for i in time_steps):
             n_steps = _count_time_steps(filename)
             time_steps = [i if i >= 0 else n_steps + i for i in time_steps]
-        
+
         time_steps = set(time_steps)
 
     with open_file(filename, "r") as f:
@@ -110,7 +110,7 @@ def _read_table(f, file_type, time_steps=None):
 
             if not (time_steps is None or t_step in time_steps):
                 continue
-            
+
             # Read time step in following line
             line = next(f).strip()
             times.append(float(line.split()[0]))
diff --git a/toughio/_mesh/_mesh.py b/toughio/_mesh/_mesh.py
index 36a090da..a4ea03b6 100644
--- a/toughio/_mesh/_mesh.py
+++ b/toughio/_mesh/_mesh.py
@@ -490,13 +490,15 @@ def read_output(
 
         """
         from .. import read_output
-        from .._io.output._common import ElementOutput, ConnectionOutput
+        from .._io.output._common import ConnectionOutput, ElementOutput
 
         if not isinstance(time_step, int):
             raise TypeError()
 
         if isinstance(file_or_output, str):
-            out = read_output(file_or_output, time_steps=time_step, connection=connection)[0]
+            out = read_output(
+                file_or_output, time_steps=time_step, connection=connection
+            )[0]
 
         else:
             out = file_or_output

From 5e1231b65eaafe92bead143471d783de84b0ff82 Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Sun, 19 May 2024 10:08:00 +0200
Subject: [PATCH 09/15] make codacy happy

---
 toughio/_io/output/_common.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/toughio/_io/output/_common.py b/toughio/_io/output/_common.py
index b559e805..07cae11f 100644
--- a/toughio/_io/output/_common.py
+++ b/toughio/_io/output/_common.py
@@ -21,12 +21,12 @@ def __init__(self, time, data, labels=None):
         self._labels = list(labels) if labels is not None else labels
 
     @abstractmethod
-    def __getitem__(self):
+    def __getitem__(self, slice):
         """Slice output."""
-        pass
+        raise NotImplementedError()
 
     @abstractmethod
-    def index(self):
+    def index(self, *args):
         """Get index of element or connection."""
         if self.labels is None:
             raise AttributeError()

From a31ad2a04014e8d031af37432dcd7ca37bc52827 Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Sun, 19 May 2024 10:14:48 +0200
Subject: [PATCH 10/15] fix method read_output

---
 toughio/_mesh/_mesh.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/toughio/_mesh/_mesh.py b/toughio/_mesh/_mesh.py
index a4ea03b6..bbaa7f57 100644
--- a/toughio/_mesh/_mesh.py
+++ b/toughio/_mesh/_mesh.py
@@ -498,7 +498,7 @@ def read_output(
         if isinstance(file_or_output, str):
             out = read_output(
                 file_or_output, time_steps=time_step, connection=connection
-            )[0]
+            )
 
         else:
             out = file_or_output

From 1d97827d00bcd6cf5be1eea31ed6189b76eb389f Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Sun, 19 May 2024 10:27:18 +0200
Subject: [PATCH 11/15] more codacy

---
 toughio/_io/output/_common.py | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/toughio/_io/output/_common.py b/toughio/_io/output/_common.py
index 07cae11f..0a7400a9 100644
--- a/toughio/_io/output/_common.py
+++ b/toughio/_io/output/_common.py
@@ -21,12 +21,12 @@ def __init__(self, time, data, labels=None):
         self._labels = list(labels) if labels is not None else labels
 
     @abstractmethod
-    def __getitem__(self, slice):
+    def __getitem__(self, islice):
         """Slice output."""
         raise NotImplementedError()
 
     @abstractmethod
-    def index(self, *args):
+    def index(self, label, *args):
         """Get index of element or connection."""
         if self.labels is None:
             raise AttributeError()
@@ -208,16 +208,16 @@ def __getitem__(self, islice):
             [self._labels[i] for i in islice],
         )
 
-    def index(self, label1, label2):
+    def index(self, label, label2=None):
         """
         Get index of connection.
 
         Parameters
         ----------
-        label1 : str
-            Label of first element of connection.
-        label2 : str
-            Label of second element of connection.
+        label : str
+            Label of connection or label of first element of connection.
+        label2 : str or None, optional, default None
+            Label of second element of connection (if `label` is the label of the first element).
 
         Returns
         -------
@@ -228,7 +228,10 @@ def index(self, label1, label2):
         super().index()
         labels = ["".join(label) for label in self.labels]
 
-        return labels.index(f"{label1}{label2}")
+        if label2 is not None:
+            label = f"{label}{label2}"
+
+        return labels.index(label)
 
 
 def to_output(file_type, labels_order, headers, times, labels, data):

From fdd498f6ce2928038d991b7340f75192a4d48e07 Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Sun, 19 May 2024 10:33:45 +0200
Subject: [PATCH 12/15] more codacy

---
 toughio/_io/output/_common.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/toughio/_io/output/_common.py b/toughio/_io/output/_common.py
index 0a7400a9..defdc5dd 100644
--- a/toughio/_io/output/_common.py
+++ b/toughio/_io/output/_common.py
@@ -26,7 +26,7 @@ def __getitem__(self, islice):
         raise NotImplementedError()
 
     @abstractmethod
-    def index(self, label, *args):
+    def index(self, label, *args, **kwargs):
         """Get index of element or connection."""
         if self.labels is None:
             raise AttributeError()
@@ -127,7 +127,7 @@ def __getitem__(self, islice):
             [self._labels[i] for i in islice],
         )
 
-    def index(self, label):
+    def index(self, label, *args, **kwargs):
         """
         Get index of element.
 
@@ -142,7 +142,7 @@ def index(self, label):
             Index of element.
 
         """
-        super().index()
+        super().index(label, *args, **kwargs)
 
         return self.labels.index(label)
 
@@ -208,7 +208,7 @@ def __getitem__(self, islice):
             [self._labels[i] for i in islice],
         )
 
-    def index(self, label, label2=None):
+    def index(self, label, label2=None, *args, **kwargs):
         """
         Get index of connection.
 
@@ -225,7 +225,7 @@ def index(self, label, label2=None):
             Index of connection.
 
         """
-        super().index()
+        super().index(label, *args, **kwargs)
         labels = ["".join(label) for label in self.labels]
 
         if label2 is not None:

From e8ab7986bd8fd9095e98340fe7537324542946bf Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Sun, 19 May 2024 11:09:00 +0200
Subject: [PATCH 13/15] test getitem method for outputs

---
 tests/test_output.py | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/tests/test_output.py b/tests/test_output.py
index 8a0c7c0e..6aa629fd 100644
--- a/tests/test_output.py
+++ b/tests/test_output.py
@@ -153,3 +153,41 @@ def test_save():
     assert helpers.allclose(0.01, save.data["porosity"].mean())
 
     assert "userx" not in save.data
+
+
+@pytest.mark.parametrize(
+    "output_ref, islice",
+    [
+        (helpers.output_eleme[0], 0),
+        (helpers.output_eleme[0], [0, 2]),
+        (helpers.output_eleme[0], "AAA00"),
+        (helpers.output_eleme[0], ["AAA00", "AAA02"]),
+        (helpers.output_conne[0], 0),
+        (helpers.output_conne[0], [0, 2]),
+        (helpers.output_conne[0], "AAA00"),
+    ],
+)
+def test_getitem(output_ref, islice):
+    output = output_ref[islice]
+
+    idx = [islice] if isinstance(islice, (int, str)) else islice
+    idx = [i if isinstance(i, int) else int(i[-1]) for i in idx]
+
+    if not isinstance(output, dict):
+        assert np.allclose(output.time, output_ref.time)
+        assert len(idx) == output.n_data
+
+        for i, iref in enumerate(idx):
+            if isinstance(output.labels[i], str):
+                assert output.labels[i] == output_ref.labels[iref]
+
+            else:
+                for label, label_ref in zip(output.labels[i], output_ref.labels[iref]):
+                    assert label == label_ref
+
+            for k, v in output.data.items():
+                assert np.allclose(v[i], output_ref.data[k][iref])
+
+    else:
+        for k, v in output.items():
+            assert np.allclose(v, output_ref.data[k][idx[0]])

From 88c2bc2285e1b82c5f7d394dcca693b54e9c2dd7 Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Sun, 19 May 2024 11:13:48 +0200
Subject: [PATCH 14/15] support python 3.12

---
 .github/workflows/ci.yml | 2 +-
 setup.cfg                | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 52e5e073..b9bfb9b8 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"]
 
     steps:
     - name: Checkout code
diff --git a/setup.cfg b/setup.cfg
index 4e87b6d3..c9c49b46 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -21,6 +21,7 @@ classifiers =
     Programming Language :: Python :: 3.9
     Programming Language :: Python :: 3.10
     Programming Language :: Python :: 3.11
+    Programming Language :: Python :: 3.12
     Natural Language :: English
     Intended Audience :: Science/Research
     Topic :: Scientific/Engineering

From 636014e08a30ed2857645e816eae838e2b80156e Mon Sep 17 00:00:00 2001
From: Keurfon Luu <keurfonluu@outlook.com>
Date: Sun, 19 May 2024 11:27:53 +0200
Subject: [PATCH 15/15] temporarily disable test with python 3.12

---
 .github/workflows/ci.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index b9bfb9b8..52e5e073 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
 
     steps:
     - name: Checkout code