From 95018ca9a21252be11906e8b3f1c23651455fc93 Mon Sep 17 00:00:00 2001
From: Nick Papior <nickpapior@gmail.com>
Date: Tue, 24 Sep 2024 11:25:21 +0200
Subject: [PATCH] enabled times file which returns xarray data for routines

Signed-off-by: Nick Papior <nickpapior@gmail.com>
---
 CHANGELOG.md                   |   3 +
 docs/api/io/siesta.rst         |   1 +
 src/sisl/io/siesta/__init__.py |   2 +
 src/sisl/io/siesta/times.py    | 120 +++++++++++++++++++++++++++++++++
 4 files changed, 126 insertions(+)
 create mode 100644 src/sisl/io/siesta/times.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 98df783c68..47de451175 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,9 @@ we hit release version 1.0.0.
 
 ## [0.15.2] - YYYY-MM-DD
 
+### Added
+- added `timesSileSiesta` which can read siesta TIMES output
+
 ### Fixed
 - fixed warning raised when reading electronic structure without Ef, #826
 - fixed precision writing of energies in GF files (siesta), #827
diff --git a/docs/api/io/siesta.rst b/docs/api/io/siesta.rst
index b789cb8707..9eb16ef2e0 100644
--- a/docs/api/io/siesta.rst
+++ b/docs/api/io/siesta.rst
@@ -30,6 +30,7 @@ Siesta
    kpSileSiesta - k-points from simulation
    rkpSileSiesta - k-points to simulation
    structSileSiesta - geometry in STRUCT_* files
+   timesSileSiesta - timings from siesta
 
 
 .. autosummary::
diff --git a/src/sisl/io/siesta/__init__.py b/src/sisl/io/siesta/__init__.py
index 36940a2d4b..3877d685ee 100644
--- a/src/sisl/io/siesta/__init__.py
+++ b/src/sisl/io/siesta/__init__.py
@@ -37,6 +37,7 @@
    kpSileSiesta - k-points from simulation
    rkpSileSiesta - k-points to simulation
    structSileSiesta - geometry in STRUCT_* files
+   timesSileSiesta - timing output from siesta
 
 
 The TranSiesta specific output files are:
@@ -63,6 +64,7 @@
 from .siesta_nc import *
 from .stdout import *
 from .struct import *
+from .times import *
 from .transiesta_grid import *
 from .vibra import *
 from .xv import *
diff --git a/src/sisl/io/siesta/times.py b/src/sisl/io/siesta/times.py
new file mode 100644
index 0000000000..a22feb61ae
--- /dev/null
+++ b/src/sisl/io/siesta/times.py
@@ -0,0 +1,120 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at https://mozilla.org/MPL/2.0/.
+from __future__ import annotations
+
+import numpy as np
+import xarray as xr
+
+from sisl._array import arrayd, arrayi
+from sisl._internal import set_module
+
+from ..sile import add_sile, sile_fh_open
+from .sile import SileSiesta
+
+__all__ = ["timesSileSiesta"]
+
+_A = SileSiesta.InfoAttr
+
+
+@set_module("sisl.io.siesta")
+class timesSileSiesta(SileSiesta):
+    """Time information from the Siesta run"""
+
+    _info_attributes_ = [
+        _A(
+            "processors",
+            r"^timer: Number of nodes",
+            lambda attr, instance, match: int(match.string.split()[-1]),
+            not_found="error",
+        ),
+        _A(
+            "threads",
+            r"^timer: Number of threads per node",
+            lambda attr, instance, match: int(match.string.split()[-1]),
+            default=1,
+            not_found="info",
+        ),
+        _A(
+            "processor_reference",
+            r"^timer: Times refer to node",
+            lambda attr, instance, match: int(match.string.split()[-1]),
+            not_found="error",
+        ),
+        _A(
+            "wall_clock",
+            r"^timer: Total elapsed wall-clock",
+            lambda attr, instance, match: float(match.string.split()[-1]),
+            not_found="error",
+        ),
+    ]
+
+    @sile_fh_open(True)
+    def read_data(self):
+        r"""Returns data associated with the times file
+
+        Returns
+        -------
+        times : xarray.DataArray
+            containing timings.
+        """
+        rl = self.readline
+        # first line is empty
+        rl()
+
+        line = rl()  # Number of nodes
+        line = rl()  # Busiest calculating
+        line = rl()  # Times refer to node
+        line = rl()  # Total elapsed
+
+        line = rl()  # CPU times
+        for _ in range(4):
+            rl()
+
+        self.step_to("Program")
+        # header
+        rl()
+
+        # now we read content
+        routines = []
+        calls = []
+        comm = []
+        times = []
+        imbalance = []
+
+        while len(cols := rl().split()) > 1:
+            if cols[0] == "MPI":
+                break
+            routines.append(cols[0])
+            calls.append(cols[1])
+            comm.append(cols[2])
+            times.append(cols[4])
+            imbalance.append(cols[6])
+
+        def ar(ar):
+            return ar.reshape(1, -1)
+
+        dims = ["parallel", "routine"]
+        data = arrayd([calls, comm, times, imbalance]).T
+        data_vars = dict(
+            calls=(dims, ar(arrayi(calls))),
+            comm=(dims, ar(arrayd(comm))),
+            time=(dims, ar(arrayd(times))),
+            imbalance=(dims, ar(arrayd(imbalance))),
+        )
+
+        coords = dict(
+            processors=("parallel", [self.info.processors]),
+            threads=("parallel", [self.info.threads]),
+            routine=routines,
+        )
+
+        data = xr.Dataset(
+            data_vars=data_vars, coords=coords, attrs=dict(file=f"{self!s}", unit="s")
+        )
+        data = data.set_index(parallel=["processors", "threads"])
+
+        return data
+
+
+add_sile("TIMES", timesSileSiesta, gzip=True)