From b1657fd371663428f13aa2e8596135b9f62fdb3c Mon Sep 17 00:00:00 2001
From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com>
Date: Mon, 6 Jan 2025 15:53:49 +0000
Subject: [PATCH 01/11] Limited SciPy version in DAE project (#510)

---
 pySDC/projects/DAE/environment.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pySDC/projects/DAE/environment.yml b/pySDC/projects/DAE/environment.yml
index 8508e17162..b2541efb3f 100644
--- a/pySDC/projects/DAE/environment.yml
+++ b/pySDC/projects/DAE/environment.yml
@@ -5,7 +5,7 @@ channels:
   - conda-forge
 dependencies:
   - numpy
-  - scipy>=0.17.1
+  - scipy>=0.17.1,<1.15
   - dill>=0.2.6
   - mpich
   - mpi4py>=3.0.0

From c317c791e5021b9af49101bd8b71e0537ecd920f Mon Sep 17 00:00:00 2001
From: Thibaut Lunet <thibaut.lunet@tuhh.de>
Date: Thu, 16 Jan 2025 19:04:23 +0100
Subject: [PATCH 02/11] TL: added fieldsIO implementation to helpers

---
 .gitignore                                |   1 +
 pySDC/helpers/blocks.py                   | 137 ++++++++
 pySDC/helpers/fieldsIO.py                 | 400 ++++++++++++++++++++++
 pySDC/tests/test_helpers/test_fieldsIO.py | 381 +++++++++++++++++++++
 4 files changed, 919 insertions(+)
 create mode 100644 pySDC/helpers/blocks.py
 create mode 100644 pySDC/helpers/fieldsIO.py
 create mode 100644 pySDC/tests/test_helpers/test_fieldsIO.py

diff --git a/.gitignore b/.gitignore
index 0f586897a1..845782e782 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,6 +10,7 @@ step_*.png
 *.swp
 *_data.json
 !_dataRef.json
+*.pysdc
 
 # Created by https://www.gitignore.io
 
diff --git a/pySDC/helpers/blocks.py b/pySDC/helpers/blocks.py
new file mode 100644
index 0000000000..ebd38f6c63
--- /dev/null
+++ b/pySDC/helpers/blocks.py
@@ -0,0 +1,137 @@
+class BlockDecomposition(object):
+    """
+    Class decomposing a cartesian space domain (1D to 3D) into a given number of processors.
+
+    Parameters
+    ----------
+    nProcs : int
+        Total number of processors for space block decomposition.
+    gridSizes : list[int]
+        Number of grid points in each dimension
+    algo : str, optional
+        Algorithm used for hte block decomposition :
+
+        - Hybrid : approach minimizing interface communication, inspired from
+          the `[Hybrid CFD solver] <https://web.stanford.edu/group/ctr/ResBriefs07/5_larsson1_pp47_58.pdf>`_.
+        - ChatGPT : quickly generated using `[ChatGPT] <https://chatgpt.com>`_.
+        The default is "Hybrid".
+    gRank : int, optional
+        If provided, the global rank that will determine the local block distribution. Default is None.
+    order : str, optional
+        The order used when computing the rank block distribution. Default is `C`.
+    """
+
+    def __init__(self, nProcs, gridSizes, algo="Hybrid", gRank=None, order="C"):
+        dim = len(gridSizes)
+        assert dim in [1, 2, 3], "block decomposition only works for 1D, 2D or 3D domains"
+
+        if algo == "ChatGPT":
+
+            nBlocks = [1] * dim
+            for i in range(2, int(nProcs**0.5) + 1):
+                while nProcs % i == 0:
+                    nBlocks[0] *= i
+                    nProcs //= i
+                    nBlocks.sort()
+
+            if nProcs > 1:
+                nBlocks[0] *= nProcs
+
+            nBlocks.sort()
+            while len(nBlocks) < dim:
+                smallest = nBlocks.pop(0)
+                nBlocks += [1, smallest]
+                nBlocks.sort()
+
+            while len(nBlocks) > dim:
+                smallest = nBlocks.pop(0)
+                next_smallest = nBlocks.pop(0)
+                nBlocks.append(smallest * next_smallest)
+                nBlocks.sort()
+
+        elif algo == "Hybrid":
+            rest = nProcs
+            facs = {
+                1: [1],
+                2: [2, 1],
+                3: [2, 3, 1],
+            }[dim]
+            exps = [0] * dim
+            for n in range(dim - 1):
+                while (rest % facs[n]) == 0:
+                    exps[n] = exps[n] + 1
+                    rest = rest // facs[n]
+            if rest > 1:
+                facs[dim - 1] = rest
+                exps[dim - 1] = 1
+
+            nBlocks = [1] * dim
+            for n in range(dim - 1, -1, -1):
+                while exps[n] > 0:
+                    dummymax = -1
+                    dmax = 0
+                    for d, nPts in enumerate(gridSizes):
+                        dummy = (nPts + nBlocks[d] - 1) // nBlocks[d]
+                        if dummy >= dummymax:
+                            dummymax = dummy
+                            dmax = d
+                    nBlocks[dmax] = nBlocks[dmax] * facs[n]
+                    exps[n] = exps[n] - 1
+
+        else:
+            raise NotImplementedError(f"algo={algo}")
+
+        # Store attributes
+        self.dim = dim
+        self.nBlocks = nBlocks
+        self.gridSizes = gridSizes
+
+        # Used for rank block distribution
+        self.gRank = gRank
+        self.order = order
+
+    @property
+    def ranks(self):
+        gRank, order = self.gRank, self.order
+        assert gRank is not None, "gRank attribute need to be set"
+        dim, nBlocks = self.dim, self.nBlocks
+        if dim == 1:
+            return (gRank,)
+        elif dim == 2:
+            div = nBlocks[-1] if order == "C" else nBlocks[0]
+            return (gRank // div, gRank % div)
+        else:
+            raise NotImplementedError(f"dim={dim}")
+
+    @property
+    def localBounds(self):
+        iLocList, nLocList = [], []
+        for rank, nPoints, nBlocks in zip(self.ranks, self.gridSizes, self.nBlocks):
+            n0 = nPoints // nBlocks
+            nRest = nPoints - nBlocks * n0
+            nLoc = n0 + 1 * (rank < nRest)
+            iLoc = rank * n0 + nRest * (rank >= nRest) + rank * (rank < nRest)
+
+            iLocList.append(iLoc)
+            nLocList.append(nLoc)
+        return iLocList, nLocList
+
+
+if __name__ == "__main__":
+    from mpi4py import MPI
+    from time import sleep
+
+    comm: MPI.Intracomm = MPI.COMM_WORLD
+    MPI_SIZE = comm.Get_size()
+    MPI_RANK = comm.Get_rank()
+
+    blocks = BlockDecomposition(MPI_SIZE, [256, 64], gRank=MPI_RANK)
+    if MPI_RANK == 0:
+        print(f"nBlocks : {blocks.nBlocks}")
+
+    ranks = blocks.ranks
+    bounds = blocks.localBounds
+
+    comm.Barrier()
+    sleep(0.01 * MPI_RANK)
+    print(f"[Rank {MPI_RANK}] pRankX={ranks}, bounds={bounds}")
diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py
new file mode 100644
index 0000000000..4c516d8947
--- /dev/null
+++ b/pySDC/helpers/fieldsIO.py
@@ -0,0 +1,400 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Base generic script for fields IO
+"""
+import os
+import numpy as np
+from typing import Type, TypeVar
+
+try:
+    from mpi4py import MPI
+except ImportError:
+    pass
+
+T = TypeVar("T")
+
+# Supported data types
+DTYPES = {
+    0: np.float64,  # double precision
+    1: np.complex128,
+    2: np.float128,  # quadruple precision
+    3: np.complex256,
+    4: np.float32,  # single precision
+    5: np.complex64,
+}
+DTYPES_AVAIL = {val: key for key, val in DTYPES.items()}
+
+# Header dtype
+H_DTYPE = np.int8
+T_DTYPE = np.float64
+
+
+class FieldsIO:
+
+    STRUCTS = {}  # supported structures, modified dynamically
+    sID = None  # structure ID of the FieldsIO class, modified dynamically
+
+    tSize = T_DTYPE().itemsize
+
+    def __init__(self, dtype, fileName):
+        assert dtype in DTYPES_AVAIL, f"dtype not available ({dtype})"
+        self.dtype = dtype
+        self.fileName = fileName
+        self.initialized = False
+
+        # Initialized by the setHeader abstract method
+        self.header = None
+        self.nItems = None  # number of values (dof) stored into one field
+
+    def __str__(self):
+        return f"FieldsIO[{self.__class__.__name__}|{self.dtype.__name__}|file:{self.fileName}]<{hex(id(self))}>"
+
+    def __repr__(self):
+        return self.__str__()
+
+    @classmethod
+    def fromFile(cls, fileName):
+        assert os.path.isfile(fileName), f"not a file ({fileName})"
+        with open(fileName, "rb") as f:
+            STRUCT, DTYPE = np.fromfile(f, dtype=H_DTYPE, count=2)
+            fieldsIO: FieldsIO = cls.STRUCTS[STRUCT](DTYPES[DTYPE], fileName)
+            fieldsIO.readHeader(f)
+            fieldsIO.initialized = True
+        return fieldsIO
+
+    @property
+    def hBase(self) -> np.ndarray:
+        return np.array([self.sID, DTYPES_AVAIL[self.dtype]], dtype=H_DTYPE)
+
+    @classmethod
+    def register(cls, sID):
+        def wrapper(registered: Type[T]) -> Type[T]:
+            assert (
+                sID not in cls.STRUCTS
+            ), f"struct ID already taken by {cls.STRUCTS[sID]}, cannot use it for {registered}"
+            cls.STRUCTS[sID] = registered
+            registered.sID = sID
+            return registered
+
+        return wrapper
+
+    def initialize(self):
+        assert self.header is not None, "header must be set before initializing FieldsIO"
+        assert not self.initialized, "FieldsIO already initialized"
+
+        with open(self.fileName, "w+b") as f:
+            self.hBase.tofile(f)
+            for array in self.hInfos:
+                array.tofile(f)
+        self.initialized = True
+
+    def setHeader(self, **params):
+        """Set the header before creating a new file to store the fields"""
+        raise NotImplementedError()
+
+    @property
+    def hInfos(self) -> list[np.ndarray]:
+        raise NotImplementedError()
+
+    @property
+    def hSize(self):
+        """Size of the full header (in bytes)"""
+        return self.hBase.nbytes + sum(hInfo.nbytes for hInfo in self.hInfos)
+
+    @property
+    def itemSize(self):
+        return self.dtype().itemsize
+
+    @property
+    def fSize(self):
+        """Full size of a field (in bytes)"""
+        return self.nItems * self.itemSize
+
+    @property
+    def fileSize(self):
+        return os.path.getsize(self.fileName)
+
+    def readHeader(self, f):
+        """Read the header from the file storing the fields"""
+        raise NotImplementedError()
+
+    def addField(self, time, field):
+        assert self.initialized, "cannot add field to a non initialized FieldsIO"
+        field = np.asarray(field)
+        assert field.dtype == self.dtype, f"expected {self.dtype} dtype, got {field.dtype}"
+        assert field.size == self.nItems, f"expected {self.nItems} values, got {field.size}"
+        with open(self.fileName, "ab") as f:
+            np.array(time, dtype=T_DTYPE).tofile(f)
+            field.tofile(f)
+
+    @property
+    def nFields(self):
+        return int((self.fileSize - self.hSize) // (self.tSize + self.fSize))
+
+    def check(self, idx):
+        nFields = self.nFields
+        if idx < 0:
+            idx = nFields + idx
+        assert idx < nFields, f"cannot read index {idx} from {nFields} fields"
+        assert idx >= 0, f"cannot read index {idx-nFields} from {nFields} fields"
+        return idx
+
+    @property
+    def times(self):
+        times = []
+        with open(self.fileName, "rb") as f:
+            f.seek(self.hSize)
+            for i in range(self.nFields):
+                t = np.fromfile(f, dtype=T_DTYPE, count=1, offset=0 if i == 0 else self.fSize)[0]
+                times.append(float(t))
+        return times
+
+    def time(self, idx):
+        idx = self.check(idx)
+        offset = self.hSize + idx * (self.tSize + self.fSize)
+        with open(self.fileName, "rb") as f:
+            t = np.fromfile(f, dtype=T_DTYPE, count=1, offset=offset)[0]
+        return float(t)
+
+    def readField(self, idx):
+        idx = self.check(idx)
+        offset = self.hSize + idx * (self.tSize + self.fSize)
+        with open(self.fileName, "rb") as f:
+            f.seek(offset)
+            t = float(np.fromfile(f, dtype=T_DTYPE, count=1)[0])
+            field = np.fromfile(f, dtype=self.dtype, count=self.nItems)
+        self.reshape(field)
+        return t, field
+
+    def reshape(self, field):
+        pass
+
+
+@FieldsIO.register(sID=0)
+class Scal0D(FieldsIO):
+
+    def setHeader(self, nVar):
+        self.header = {"nVar": int(nVar)}
+        self.nItems = self.nVar
+
+    @property
+    def nVar(self):
+        return self.header["nVar"]
+
+    @property
+    def hInfos(self):
+        return [np.array([self.nVar], dtype=np.int64)]
+
+    def readHeader(self, f):
+        (nVar,) = np.fromfile(f, dtype=np.int64, count=1)
+        self.setHeader(nVar)
+
+
+@FieldsIO.register(sID=1)
+class Cart1D(Scal0D):
+
+    def setupGrids(self, **grids):
+        grids = {name: np.asarray(grid, dtype=np.float64) for name, grid in grids.items() if name.startswith("grid")}
+        for name, grid in grids.items():
+            assert grid.ndim == 1, f"{name} must be one dimensional"
+        return grids
+
+    def setHeader(self, nVar, gridX):
+        grids = self.setupGrids(gridX=gridX)
+        self.header = {"nVar": int(nVar), **grids}
+        self.nItems = nVar * self.nX
+
+    @property
+    def nX(self):
+        return self.header["gridX"].size
+
+    def reshape(self, fields: np.ndarray):
+        fields.shape = (self.nVar, self.nX)
+
+    @property
+    def hInfos(self):
+        return [np.array([self.nVar, self.nX], dtype=np.int64), np.array(self.header["gridX"], dtype=np.float64)]
+
+    def readHeader(self, f):
+        nVar, nX = np.fromfile(f, dtype=np.int64, count=2)
+        gridX = np.fromfile(f, dtype=np.float64, count=nX)
+        self.setHeader(nVar, gridX)
+
+    # -------------------------------------------------------------------------
+    # MPI-parallel implementation
+    # -------------------------------------------------------------------------
+    comm: MPI.Intracomm = None
+
+    @classmethod
+    def setupMPI(cls, comm: MPI.Intracomm, iLocX, nLocX):
+        cls.comm = comm
+        cls.iLocX = iLocX
+        cls.nLocX = nLocX
+        cls.mpiFile = None
+
+    @property
+    def MPI_ON(self):
+        if self.comm is None:
+            return False
+        return self.comm.Get_size() > 1
+
+    @property
+    def MPI_ROOT(self):
+        if self.comm is None:
+            return True
+        return self.comm.Get_rank() == 0
+
+    def MPI_FILE_OPEN(self, mode):
+        amode = {
+            "r": MPI.MODE_RDONLY,
+            "a": MPI.MODE_WRONLY | MPI.MODE_APPEND,
+        }[mode]
+        self.mpiFile = MPI.File.Open(self.comm, self.fileName, amode)
+
+    def MPI_WRITE(self, data):
+        self.mpiFile.Write(data)
+
+    def MPI_WRITE_AT(self, offset, data: np.ndarray):
+        self.mpiFile.Write_at(offset, data)
+
+    def MPI_READ_AT(self, offset, data):
+        self.mpiFile.Read_at(offset, data)
+
+    def MPI_FILE_CLOSE(self):
+        self.mpiFile.Close()
+        self.mpiFile = None
+
+    def initialize(self):
+        if self.MPI_ROOT:
+            super().initialize()
+        if self.MPI_ON:
+            self.comm.Barrier()
+            self.initialized = True
+
+    def addField(self, time, field):
+        if not self.MPI_ON:
+            return super().addField(time, field)
+
+        assert self.initialized, "cannot add field to a non initialized FieldsIO"
+
+        field = np.asarray(field)
+        assert field.dtype == self.dtype, f"expected {self.dtype} dtype, got {field.dtype}"
+        assert field.shape == (self.nVar, self.nLocX), f"expected {(self.nVar, self.nLocX)} shape, got {field.shape}"
+
+        offset0 = self.fileSize
+        self.MPI_FILE_OPEN(mode="a")
+        if self.MPI_ROOT:
+            self.MPI_WRITE(np.array(time, dtype=T_DTYPE))
+        offset0 += self.tSize
+
+        for iVar in range(self.nVar):
+            offset = offset0 + (iVar * self.nX + self.iLocX) * self.itemSize
+            self.MPI_WRITE_AT(offset, field[iVar])
+        self.MPI_FILE_CLOSE()
+
+    def readField(self, idx):
+        if not self.MPI_ON:
+            return super().readField(idx)
+        idx = self.check(idx)
+
+        offset0 = self.hSize + idx * (self.fSize + self.tSize)
+        with open(self.fileName, "rb") as f:
+            t = float(np.fromfile(f, dtype=T_DTYPE, count=1, offset=offset0)[0])
+        offset0 += self.tSize
+
+        field = np.empty((self.nVar, self.nLocX), dtype=self.dtype)
+
+        self.MPI_FILE_OPEN(mode="r")
+        for iVar in range(self.nVar):
+            offset = offset0 + (iVar * self.nX + self.iLocX) * self.itemSize
+            self.MPI_READ_AT(offset, field[iVar])
+        self.MPI_FILE_CLOSE()
+
+        return t, field
+
+
+@FieldsIO.register(sID=2)
+class Cart2D(Cart1D):
+
+    def setHeader(self, nVar, gridX, gridY):
+        grids = self.setupGrids(gridX=gridX, gridY=gridY)
+        self.header = {"nVar": int(nVar), **grids}
+        self.nItems = nVar * self.nX * self.nY
+
+    @property
+    def nY(self):
+        return self.header["gridY"].size
+
+    def reshape(self, fields: np.ndarray):
+        fields.shape = (self.nVar, self.nX, self.nY)
+
+    @property
+    def hInfos(self):
+        return [
+            np.array([self.nVar, self.nX, self.nY], dtype=np.int64),
+            np.array(self.header["gridX"], dtype=np.float64),
+            np.array(self.header["gridY"], dtype=np.float64),
+        ]
+
+    def readHeader(self, f):
+        nVar, nX, nY = np.fromfile(f, dtype=np.int64, count=3)
+        gridX = np.fromfile(f, dtype=np.float64, count=nX)
+        gridY = np.fromfile(f, dtype=np.float64, count=nY)
+        self.setHeader(nVar, gridX, gridY)
+
+    # -------------------------------------------------------------------------
+    # MPI-parallel implementation
+    # -------------------------------------------------------------------------
+    @classmethod
+    def setupMPI(cls, comm: MPI.Intracomm, iLocX, nLocX, iLocY, nLocY):
+        super().setupMPI(comm, iLocX, nLocX)
+        cls.iLocY = iLocY
+        cls.nLocY = nLocY
+
+    def addField(self, time, field):
+        if not self.MPI_ON:
+            return super().addField(time, field)
+
+        assert self.initialized, "cannot add field to a non initialized FieldsIO"
+
+        field = np.asarray(field)
+        assert field.dtype == self.dtype, f"expected {self.dtype} dtype, got {field.dtype}"
+        assert field.shape == (
+            self.nVar,
+            self.nLocX,
+            self.nLocY,
+        ), f"expected {(self.nVar, self.nLocX, self.nLocY)} shape, got {field.shape}"
+
+        offset0 = self.fileSize
+        self.MPI_FILE_OPEN(mode="a")
+        if self.MPI_ROOT:
+            self.MPI_WRITE(np.array(time, dtype=T_DTYPE))
+        offset0 += self.tSize
+
+        for iVar in range(self.nVar):
+            for iX in range(self.nLocX):
+                offset = offset0 + (iVar * self.nX * self.nY + (self.iLocX + iX) * self.nY + self.iLocY) * self.itemSize
+                self.MPI_WRITE_AT(offset, field[iVar, iX])
+        self.MPI_FILE_CLOSE()
+
+    def readField(self, idx):
+        if not self.MPI_ON:
+            return super().readField(idx)
+        idx = self.check(idx)
+
+        offset0 = self.hSize + idx * (self.tSize + self.fSize)
+        with open(self.fileName, "rb") as f:
+            t = float(np.fromfile(f, dtype=T_DTYPE, count=1, offset=offset0)[0])
+        offset0 += self.tSize
+
+        field = np.empty((self.nVar, self.nLocX, self.nLocY), dtype=self.dtype)
+
+        self.MPI_FILE_OPEN(mode="r")
+        for iVar in range(self.nVar):
+            for iX in range(self.nLocX):
+                offset = offset0 + (iVar * self.nX * self.nY + (self.iLocX + iX) * self.nY + self.iLocY) * self.itemSize
+                self.MPI_READ_AT(offset, field[iVar, iX])
+        self.MPI_FILE_CLOSE()
+
+        return t, field
diff --git a/pySDC/tests/test_helpers/test_fieldsIO.py b/pySDC/tests/test_helpers/test_fieldsIO.py
new file mode 100644
index 0000000000..9e5d77efa4
--- /dev/null
+++ b/pySDC/tests/test_helpers/test_fieldsIO.py
@@ -0,0 +1,381 @@
+import pytest
+import numpy as np
+
+from pySDC.helpers.fieldsIO import DTYPES
+
+
+@pytest.mark.parametrize("dtypeIdx", DTYPES.keys())
+@pytest.mark.parametrize("nDim", range(3))
+def testHeader(nDim, dtypeIdx):
+    from pySDC.helpers.fieldsIO import FieldsIO, Scal0D, Cart1D, Cart2D
+
+    fileName = "testHeader.pysdc"
+    dtype = DTYPES[dtypeIdx]
+
+    gridX = np.linspace(0, 1, num=256, endpoint=False)
+    gridY = np.linspace(0, 1, num=64, endpoint=False)
+
+    if nDim == 0:
+        Class = Scal0D
+        args = {"nVar": 20}
+    elif nDim == 1:
+        Class = Cart1D
+        args = {"nVar": 10, "gridX": gridX}
+    elif nDim == 2:
+        Class = Cart2D
+        args = {"nVar": 10, "gridX": gridX, "gridY": gridY}
+
+    f1 = Class(dtype, fileName)
+    assert f1.__str__() == f1.__repr__(), "__repr__ and __str__ do not return the same result"
+    try:
+        f1.initialize()
+    except AssertionError:
+        pass
+    else:
+        raise AssertionError(f"{f1} should not be initialized without AssertionError before header is set")
+
+    f1.setHeader(**args)
+    assert f1.header is not None, f"{f1} has still None for header after setHeader"
+    assert f1.nItems is not None, f"{f1} has still None for nItems after setHeader"
+    assert f1.nItems > 0, f"{f1} has nItems={f1.nItems} after setHeader"
+    try:
+        f1.addField(0, np.zeros(f1.nItems, dtype=f1.dtype))
+    except AssertionError:
+        pass
+    else:
+        raise AssertionError(f"{f1} should not be initialized without error before header is set")
+
+    f1.initialize()
+    assert f1.initialized, f"{f1} is not initialized after calling initialize()"
+    assert f1.fileSize == f1.hSize, f"{f1} has file size different than its header size"
+
+    f2 = FieldsIO.fromFile(fileName)
+    assert f2.initialized, f"f2 ({f2}) not initialized after instantiating from file"
+    assert type(f2) == type(f1), f"f2 ({f2}) not of the same type as f1 ({f1})"
+    assert f2.dtype == f1.dtype, f"f2 ({f2}) has not the same dtype as f1 ({f1})"
+
+    for key, val in f1.header.items():
+        assert key in f2.header, f"could not read {key} key in written {f2}"
+        assert np.allclose(val, f2.header[key]), f"header's discrepancy for {key} in written {f2}"
+
+
+@pytest.mark.parametrize("dtypeIdx", DTYPES.keys())
+@pytest.mark.parametrize("nSteps", [1, 2, 10, 100])
+@pytest.mark.parametrize("nVar", [1, 2, 5])
+def testScal0D(nVar, nSteps, dtypeIdx):
+    from pySDC.helpers.fieldsIO import FieldsIO, Scal0D
+
+    fileName = "testScal0D.pysdc"
+    dtype = DTYPES[dtypeIdx]
+
+    f1 = Scal0D(dtype, fileName)
+    f1.setHeader(nVar=nVar)
+
+    assert f1.nItems == nVar, f"{f1} do not have nItems == nVar"
+    f1.initialize()
+
+    u0 = np.random.rand(nVar).astype(f1.dtype)
+    times = np.arange(nSteps) / nSteps
+
+    for t in times:
+        ut = (u0 * t).astype(f1.dtype)
+        f1.addField(t, ut)
+
+    assert f1.nFields == nSteps, f"{f1} do not have nFields == nSteps"
+    assert np.allclose(f1.times, times), f"{f1} has wrong times stored in file"
+
+    f2 = FieldsIO.fromFile(fileName)
+
+    assert f1.nFields == f2.nFields, f"f2 ({f2}) has different nFields than f1 ({f1})"
+    assert f1.times == f2.times, f"f2 ({f2}) has different times than f1 ({f1})"
+    assert (f1.time(-1) == f2.times[-1]) and (
+        f1.times[-1] == f2.time(-1)
+    ), f"f2 ({f2}) has different last time than f1 ({f1})"
+
+    for idx, t in enumerate(times):
+        u1 = u0 * t
+        t2, u2 = f2.readField(idx)
+        assert t2 == t, f"{idx}'s fields in {f1} has incorrect time"
+        assert u2.shape == u1.shape, f"{idx}'s fields in {f1} has incorrect shape"
+        assert np.allclose(u2, u1), f"{idx}'s fields in {f1} has incorrect values"
+
+
+@pytest.mark.parametrize("dtypeIdx", DTYPES.keys())
+@pytest.mark.parametrize("nSteps", [1, 2, 5, 10])
+@pytest.mark.parametrize("nX", [5, 10, 16, 32, 64])
+@pytest.mark.parametrize("nVar", [1, 2, 5])
+def testCart1D(nVar, nX, nSteps, dtypeIdx):
+    from pySDC.helpers.fieldsIO import FieldsIO, Cart1D, DTYPES
+
+    fileName = "testCart1D.pysdc"
+    dtype = DTYPES[dtypeIdx]
+
+    gridX = np.linspace(0, 1, num=nX, endpoint=False)
+    nX = gridX.size
+
+    f1 = Cart1D(dtype, fileName)
+    f1.setHeader(nVar=nVar, gridX=gridX)
+
+    assert f1.nItems == nVar * nX, f"{f1} do not have nItems == nVar*nX"
+    assert f1.nX == nX, f"{f1} has incorrect nX"
+    f1.initialize()
+
+    u0 = np.random.rand(nVar, nX).astype(f1.dtype)
+    times = np.arange(nSteps) / nSteps
+
+    for t in times:
+        ut = (u0 * t).astype(f1.dtype)
+        f1.addField(t, ut)
+
+    assert f1.nFields == nSteps, f"{f1} do not have nFields == nSteps"
+    assert np.allclose(f1.times, times), f"{f1} has wrong times stored in file"
+
+    f2 = FieldsIO.fromFile(fileName)
+
+    assert f1.nFields == f2.nFields, f"f2 ({f2}) has different nFields than f1 ({f1})"
+    assert f1.times == f2.times, f"f2 ({f2}) has different times than f1 ({f1})"
+    assert (f1.time(-1) == f2.times[-1]) and (
+        f1.times[-1] == f2.time(-1)
+    ), f"f2 ({f2}) has different last time than f1 ({f1})"
+
+    for idx, t in enumerate(times):
+        u1 = u0 * t
+        t2, u2 = f2.readField(idx)
+        assert t2 == t, f"{idx}'s fields in {f1} has incorrect time"
+        assert u2.shape == u1.shape, f"{idx}'s fields in {f1} has incorrect shape"
+        assert np.allclose(u2, u1), f"{idx}'s fields in {f1} has incorrect values"
+
+
+@pytest.mark.parametrize("dtypeIdx", DTYPES.keys())
+@pytest.mark.parametrize("nSteps", [1, 2, 5, 10])
+@pytest.mark.parametrize("nY", [5, 10, 16])
+@pytest.mark.parametrize("nX", [5, 10, 16])
+@pytest.mark.parametrize("nVar", [1, 2, 5])
+def testCart2D(nVar, nX, nY, nSteps, dtypeIdx):
+    from pySDC.helpers.fieldsIO import FieldsIO, Cart2D, DTYPES
+
+    fileName = "testCart2D.pysdc"
+    dtype = DTYPES[dtypeIdx]
+
+    gridX = np.linspace(0, 1, num=nX, endpoint=False)
+    gridY = np.linspace(0, 1, num=nY, endpoint=False)
+
+    f1 = Cart2D(dtype, fileName)
+    f1.setHeader(nVar=nVar, gridX=gridX, gridY=gridY)
+
+    assert f1.nItems == nVar * nX * nY, f"{f1} do not have nItems == nVar*nX"
+    assert f1.nX == nX, f"{f1} has incorrect nX"
+    assert f1.nY == nY, f"{f1} has incorrect nY"
+    f1.initialize()
+
+    u0 = np.random.rand(nVar, nX, nY).astype(f1.dtype)
+    times = np.arange(nSteps) / nSteps
+
+    for t in times:
+        ut = (u0 * t).astype(f1.dtype)
+        f1.addField(t, ut)
+
+    assert f1.nFields == nSteps, f"{f1} do not have nFields == nSteps"
+    assert np.allclose(f1.times, times), f"{f1} has wrong times stored in file"
+
+    f2 = FieldsIO.fromFile(fileName)
+
+    assert f1.nFields == f2.nFields, f"f2 ({f2}) has different nFields than f1 ({f1})"
+    assert f1.times == f2.times, f"f2 ({f2}) has different times than f1 ({f1})"
+    assert (f1.time(-1) == f2.times[-1]) and (
+        f1.times[-1] == f2.time(-1)
+    ), f"f2 ({f2}) has different last time than f1 ({f1})"
+
+    for idx, t in enumerate(times):
+        u1 = u0 * t
+        t2, u2 = f2.readField(idx)
+        assert t2 == t, f"{idx}'s fields in {f1} has incorrect time"
+        assert u2.shape == u1.shape, f"{idx}'s fields in {f1} has incorrect shape"
+        assert np.allclose(u2, u1), f"{idx}'s fields in {f1} has incorrect values"
+
+
+def initGrid(nVar, nX, nY=None):
+    nDim = 1
+    if nY is not None:
+        nDim += 1
+    x = np.linspace(0, 1, num=nX, endpoint=False)
+    grids = (x,)
+    gridSizes = (nX,)
+    u0 = np.array(np.arange(nVar) + 1)[:, None] * x[None, :]
+
+    if nDim > 1:
+        y = np.linspace(0, 1, num=nY, endpoint=False)
+        grids += (y,)
+        gridSizes += (nY,)
+        u0 = u0[:, :, None] * y[None, None, :]
+
+    return grids, gridSizes, u0
+
+
+def writeFields_MPI(fileName, nDim, dtypeIdx, algo, nSteps, nVar, nX, nY=None):
+    grids, gridSizes, u0 = initGrid(nVar, nX, nY)
+
+    from mpi4py import MPI
+
+    comm = MPI.COMM_WORLD
+    MPI_SIZE = comm.Get_size()
+    MPI_RANK = comm.Get_rank()
+
+    from pySDC.helpers.blocks import BlockDecomposition
+
+    blocks = BlockDecomposition(MPI_SIZE, gridSizes, algo, MPI_RANK)
+
+    from time import sleep
+
+    from pySDC.helpers.fieldsIO import Cart1D, Cart2D
+
+    if nDim == 1:
+        (iLocX,), (nLocX,) = blocks.localBounds
+        (pRankX,) = blocks.ranks
+        Cart1D.setupMPI(comm, iLocX, nLocX)
+        u0 = u0[:, iLocX : iLocX + nLocX]
+
+        MPI.COMM_WORLD.Barrier()
+        sleep(0.01 * MPI_RANK)
+        print(f"[Rank {MPI_RANK}] pRankX={pRankX} ({iLocX}, {nLocX})")
+        MPI.COMM_WORLD.Barrier()
+
+        f1 = Cart1D(DTYPES[dtypeIdx], fileName)
+        f1.setHeader(nVar=nVar, gridX=grids[0])
+
+    if nDim == 2:
+        (iLocX, iLocY), (nLocX, nLocY) = blocks.localBounds
+        Cart2D.setupMPI(comm, iLocX, nLocX, iLocY, nLocY)
+        u0 = u0[:, iLocX : iLocX + nLocX, iLocY : iLocY + nLocY]
+
+        f1 = Cart2D(DTYPES[dtypeIdx], fileName)
+        f1.setHeader(nVar=nVar, gridX=grids[0], gridY=grids[1])
+
+    u0 = np.asarray(u0, dtype=f1.dtype)
+    f1.initialize()
+
+    times = np.arange(nSteps) / nSteps
+    for t in times:
+        ut = (u0 * t).astype(f1.dtype)
+        f1.addField(t, ut)
+
+    return u0
+
+
+def compareFields_MPI(fileName, u0, nSteps):
+    from pySDC.helpers.fieldsIO import FieldsIO
+
+    f2 = FieldsIO.fromFile(fileName)
+
+    times = np.arange(nSteps) / nSteps
+    for idx, t in enumerate(times):
+        u1 = u0 * t
+        t2, u2 = f2.readField(idx)
+        assert t2 == t, f"{idx}'s fields in {f2} has incorrect time"
+        assert u2.shape == u1.shape, f"{idx}'s fields in {f2} has incorrect shape"
+        assert np.allclose(u2, u1), f"{idx}'s fields in {f2} has incorrect values"
+
+
+@pytest.mark.mpi4py
+@pytest.mark.parametrize("nX", [61, 16, 32])
+@pytest.mark.parametrize("nVar", [1, 4])
+@pytest.mark.parametrize("nSteps", [1, 10])
+@pytest.mark.parametrize("algo", ["ChatGPT", "Hybrid"])
+@pytest.mark.parametrize("dtypeIdx", [0, 1])
+@pytest.mark.parametrize("nProcs", [1, 2, 4])
+def testCart1D_MPI(nProcs, dtypeIdx, algo, nSteps, nVar, nX):
+
+    import subprocess
+
+    fileName = "testCart1D_MPI.pysdc"
+
+    cmd = f"mpirun -np {nProcs} python {__file__} --fileName {fileName} --nDim 1 "
+    cmd += f"--dtypeIdx {dtypeIdx} --algo {algo} --nSteps {nSteps} --nVar {nVar} --nX {nX}"
+
+    p = subprocess.Popen(cmd.split(), cwd=".")
+    p.wait()
+    assert p.returncode == 0, f"MPI write with {nProcs} did not return code 0, but {p.returncode}"
+
+    from pySDC.helpers.fieldsIO import FieldsIO, Cart1D
+
+    f2: Cart1D = FieldsIO.fromFile(fileName)
+
+    assert type(f2) == Cart1D, f"incorrect type in MPI written fields {f2}"
+    assert f2.nFields == nSteps, f"incorrect nFields in MPI written fields {f2}"
+    assert f2.nVar == nVar, f"incorrect nVar in MPI written fields {f2}"
+    assert f2.nX == nX, f"incorrect nX in MPI written fields {f2}"
+
+    grids, _, u0 = initGrid(nVar, nX)
+    assert np.allclose(f2.header['gridX'], grids[0]), f"incorrect gridX in MPI written fields {f2}"
+
+    times = np.arange(nSteps) / nSteps
+    for idx, t in enumerate(times):
+        u1 = u0 * t
+        t2, u2 = f2.readField(idx)
+        assert t2 == t, f"{idx}'s fields in {f2} has incorrect time"
+        assert u2.shape == u1.shape, f"{idx}'s fields in {f2} has incorrect shape"
+        assert np.allclose(u2, u1), f"{idx}'s fields in {f2} has incorrect values"
+
+
+@pytest.mark.mpi4py
+@pytest.mark.parametrize("nY", [61, 16, 32])
+@pytest.mark.parametrize("nX", [61, 16, 32])
+@pytest.mark.parametrize("nVar", [1, 4])
+@pytest.mark.parametrize("nSteps", [1, 10])
+@pytest.mark.parametrize("algo", ["ChatGPT", "Hybrid"])
+@pytest.mark.parametrize("dtypeIdx", [0, 1])
+@pytest.mark.parametrize("nProcs", [1, 2, 4])
+def testCart2D_MPI(nProcs, dtypeIdx, algo, nSteps, nVar, nX, nY):
+
+    import subprocess
+
+    fileName = "testCart2D_MPI.pysdc"
+
+    cmd = f"mpirun -np {nProcs} python {__file__} --fileName {fileName} --nDim 2 "
+    cmd += f"--dtypeIdx {dtypeIdx} --algo {algo} --nSteps {nSteps} --nVar {nVar} --nX {nX} --nY {nY}"
+
+    p = subprocess.Popen(cmd.split(), cwd=".")
+    p.wait()
+    assert p.returncode == 0, f"MPI write with {nProcs} did not return code 0, but {p.returncode}"
+
+    from pySDC.helpers.fieldsIO import FieldsIO, Cart2D
+
+    f2: Cart2D = FieldsIO.fromFile(fileName)
+
+    assert type(f2) == Cart2D, f"incorrect type in MPI written fields {f2}"
+    assert f2.nFields == nSteps, f"incorrect nFields in MPI written fields {f2}"
+    assert f2.nVar == nVar, f"incorrect nVar in MPI written fields {f2}"
+    assert f2.nX == nX, f"incorrect nX in MPI written fields {f2}"
+    assert f2.nY == nY, f"incorrect nY in MPI written fields {f2}"
+
+    grids, _, u0 = initGrid(nVar, nX, nY)
+    assert np.allclose(f2.header['gridX'], grids[0]), f"incorrect gridX in MPI written fields {f2}"
+    assert np.allclose(f2.header['gridY'], grids[1]), f"incorrect gridY in MPI written fields {f2}"
+
+    times = np.arange(nSteps) / nSteps
+    for idx, t in enumerate(times):
+        u1 = u0 * t
+        t2, u2 = f2.readField(idx)
+        assert t2 == t, f"{idx}'s fields in {f2} has incorrect time"
+        assert u2.shape == u1.shape, f"{idx}'s fields in {f2} has incorrect shape"
+        assert np.allclose(u2, u1), f"{idx}'s fields in {f2} has incorrect values"
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--fileName', type=str, help='fileName of the file')
+    parser.add_argument('--nDim', type=int, help='space dimension', choices=[1, 2])
+    parser.add_argument('--dtypeIdx', type=int, help="dtype index", choices=DTYPES.keys())
+    parser.add_argument(
+        '--algo', type=str, help="algorithm used for block decomposition", choices=["ChatGPT", "Hybrid"]
+    )
+    parser.add_argument('--nSteps', type=int, help="number of field variables")
+    parser.add_argument('--nVar', type=int, help="number of field variables")
+    parser.add_argument('--nX', type=int, help="number of grid points in x dimension")
+    parser.add_argument('--nY', type=int, help="number of grid points in y dimension")
+    args = parser.parse_args()
+
+    u0 = writeFields_MPI(**args.__dict__)
+    compareFields_MPI(args.fileName, u0, args.nSteps)

From 291512b9d0058305175560e296c15d5489b7c552 Mon Sep 17 00:00:00 2001
From: Thibaut Lunet <thibaut.lunet@tuhh.de>
Date: Thu, 16 Jan 2025 19:15:04 +0100
Subject: [PATCH 03/11] TL: fixed the fieldsIO when mpi4py not available

---
 pySDC/helpers/fieldsIO.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py
index 4c516d8947..f5883695ca 100644
--- a/pySDC/helpers/fieldsIO.py
+++ b/pySDC/helpers/fieldsIO.py
@@ -10,7 +10,8 @@
 try:
     from mpi4py import MPI
 except ImportError:
-    pass
+    class MPI:
+        COMM_WORLD = None
 
 T = TypeVar("T")
 

From d517d7e83a36be6c78d31bb010dd9c5da9563ec3 Mon Sep 17 00:00:00 2001
From: Thibaut Lunet <thibaut.lunet@tuhh.de>
Date: Thu, 16 Jan 2025 19:18:20 +0100
Subject: [PATCH 04/11] TL: SO IMPORTANT CHANGE THANKS BLACK !!!

---
 pySDC/helpers/fieldsIO.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py
index f5883695ca..c9dfa0086a 100644
--- a/pySDC/helpers/fieldsIO.py
+++ b/pySDC/helpers/fieldsIO.py
@@ -10,9 +10,11 @@
 try:
     from mpi4py import MPI
 except ImportError:
+
     class MPI:
         COMM_WORLD = None
 
+
 T = TypeVar("T")
 
 # Supported data types

From 3886bdab44cb287d89539c5c1ae507a824715eab Mon Sep 17 00:00:00 2001
From: Thibaut Lunet <thibaut.lunet@tuhh.de>
Date: Fri, 17 Jan 2025 09:18:57 +0100
Subject: [PATCH 05/11] TL: cleaning and docstrings

---
 pySDC/helpers/fieldsIO.py                 | 387 ++++++++++++++++++++--
 pySDC/tests/test_helpers/test_fieldsIO.py |  13 +-
 2 files changed, 357 insertions(+), 43 deletions(-)

diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py
index c9dfa0086a..d7b32d9804 100644
--- a/pySDC/helpers/fieldsIO.py
+++ b/pySDC/helpers/fieldsIO.py
@@ -1,7 +1,43 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 """
-Base generic script for fields IO
+Generic utility class to write and read cartesian grid field solutions into binary files.
+It implements the base file handler class :class:`FieldsIO`, that is specialized into :
+
+- :class:`Scal0D` : for 0D fields (scalar) with a given number of variables
+- :class:`Cart1D` : for 1D fields with a given number of variables
+- :class:`Cart2D` : for 2D fields on a cartesian grid with a given number of variables
+
+While each file handler need to be setup with specific parameters (grid, ...),
+each written file can be read using the same interface implemented in the
+base abstract class.
+
+Example
+-------
+>>> import numpy as np
+>>> from pySDC.helpers.fieldsIO import Cart1D, FieldsIO
+>>>
+>>> # Write some fields in files
+>>> x = np.linspace(0, 1, 101)
+>>> fOut = Cart2D(np.float64, "file.pysdc")
+>>> fOut.setHeader(nVar=2, gridX=x)
+>>> fOut.initialize()
+>>> times = [0, 1, 2]
+>>> u0 = np.array([-1, 1])[:, None]*x[None, :]
+>>> for t in times:
+>>>    fOut.addField(t, t*u0)
+>>>
+>>> # Read the file using a the generic interface
+>>> fIn = FieldsIO.fromFile("file.pysdc")
+>>> times = fIn.times
+>>> assert len(times) == fIn.nFields
+>>> tEnd, uEnd = fIn.readField(-1)
+>>> assert tEnd == times[-1]
+
+Notes
+-----
+🚀 :class:`Cart1D` and :class:`Cart2D` are compatible with a MPI-based cartesian decomposition.
+See :class:`pySDC.tests.test_helpers.test_fieldsIO.writeFields_MPI` for an illustrative example.
 """
 import os
 import numpy as np
@@ -34,6 +70,7 @@ class MPI:
 
 
 class FieldsIO:
+    """Abstract IO file handler"""
 
     STRUCTS = {}  # supported structures, modified dynamically
     sID = None  # structure ID of the FieldsIO class, modified dynamically
@@ -41,6 +78,14 @@ class FieldsIO:
     tSize = T_DTYPE().itemsize
 
     def __init__(self, dtype, fileName):
+        """
+        Parameters
+        ----------
+        dtype : np.dtype
+            The data type of the fields values.
+        fileName : str
+            File.
+        """
         assert dtype in DTYPES_AVAIL, f"dtype not available ({dtype})"
         self.dtype = dtype
         self.fileName = fileName
@@ -58,6 +103,20 @@ def __repr__(self):
 
     @classmethod
     def fromFile(cls, fileName):
+        """
+        Read a file storing fields, and return the `FieldsIO` of the appropriate
+        field type (structure).
+
+        Parameters
+        ----------
+        fileName : str
+            Name of the binary file.
+
+        Returns
+        -------
+        fieldsIO : :class:`FieldsIO`
+            The specialized `FieldsIO` adapted to the file.
+        """
         assert os.path.isfile(fileName), f"not a file ({fileName})"
         with open(fileName, "rb") as f:
             STRUCT, DTYPE = np.fromfile(f, dtype=H_DTYPE, count=2)
@@ -68,10 +127,29 @@ def fromFile(cls, fileName):
 
     @property
     def hBase(self) -> np.ndarray:
+        """Base header into numpy array format"""
         return np.array([self.sID, DTYPES_AVAIL[self.dtype]], dtype=H_DTYPE)
 
     @classmethod
     def register(cls, sID):
+        """
+        Decorator used to register a new class FieldsIO specialized class
+
+        Parameters
+        ----------
+        sID : int
+            Unique identifyer for the file, used in the binary file.
+            Since it's encoded on a 8-bytes signed integer,
+            it must be between -128 and 127
+
+        Example
+        -------
+        >>> # New specialized FieldsIO class
+        >>> @FieldsIO.register(sID=31)
+        >>> class HexaMesh2D(FieldsIO):
+        >>>     pass # ... implementation
+        """
+
         def wrapper(registered: Type[T]) -> Type[T]:
             assert (
                 sID not in cls.STRUCTS
@@ -83,6 +161,7 @@ def wrapper(registered: Type[T]) -> Type[T]:
         return wrapper
 
     def initialize(self):
+        """Initialize the file handler : create the file with header, removing any existing file with the same name"""
         assert self.header is not None, "header must be set before initializing FieldsIO"
         assert not self.initialized, "FieldsIO already initialized"
 
@@ -93,11 +172,23 @@ def initialize(self):
         self.initialized = True
 
     def setHeader(self, **params):
-        """Set the header before creating a new file to store the fields"""
+        """(Abstract) Set the header before creating a new file to store the fields"""
         raise NotImplementedError()
 
     @property
     def hInfos(self) -> list[np.ndarray]:
+        """(Abstract) Array representing the grid structure to be written in the binary file."""
+        raise NotImplementedError()
+
+    def readHeader(self, f):
+        """
+        (Abstract) Read the header from the file storing the fields.
+
+        Parameters
+        ----------
+        f : `_io.TextIOWrapper`
+            File to read the header from.
+        """
         raise NotImplementedError()
 
     @property
@@ -107,6 +198,7 @@ def hSize(self):
 
     @property
     def itemSize(self):
+        """Size of one field value (in bytes)"""
         return self.dtype().itemsize
 
     @property
@@ -116,13 +208,20 @@ def fSize(self):
 
     @property
     def fileSize(self):
+        """Current size of the file (in bytes)"""
         return os.path.getsize(self.fileName)
 
-    def readHeader(self, f):
-        """Read the header from the file storing the fields"""
-        raise NotImplementedError()
-
     def addField(self, time, field):
+        """
+        Append one field solution at the end of the file with one given time.
+
+        Parameters
+        ----------
+        time : float-like
+            The associated time of the field solution.
+        field : np.ndarray
+            The field values.
+        """
         assert self.initialized, "cannot add field to a non initialized FieldsIO"
         field = np.asarray(field)
         assert field.dtype == self.dtype, f"expected {self.dtype} dtype, got {field.dtype}"
@@ -133,9 +232,11 @@ def addField(self, time, field):
 
     @property
     def nFields(self):
+        """Number of fields currently stored in the binary file"""
         return int((self.fileSize - self.hSize) // (self.tSize + self.fSize))
 
     def check(self, idx):
+        """Utility method to check a positional index for a stored field"""
         nFields = self.nFields
         if idx < 0:
             idx = nFields + idx
@@ -145,6 +246,7 @@ def check(self, idx):
 
     @property
     def times(self):
+        """Vector of all times stored in the binary file"""
         times = []
         with open(self.fileName, "rb") as f:
             f.seek(self.hSize)
@@ -154,6 +256,7 @@ def times(self):
         return times
 
     def time(self, idx):
+        """Time stored at a given field index"""
         idx = self.check(idx)
         offset = self.hSize + idx * (self.tSize + self.fSize)
         with open(self.fileName, "rb") as f:
@@ -161,6 +264,22 @@ def time(self, idx):
         return float(t)
 
     def readField(self, idx):
+        """
+        Read one field stored in the binary file, corresponding to the given
+        time index.
+
+        Parameters
+        ----------
+        idx : int
+            Positional index of the field.
+
+        Returns
+        -------
+        t : float
+            Stored time for this field.
+        field : np.ndarray
+            Read fields in a numpy array.
+        """
         idx = self.check(idx)
         offset = self.hSize + idx * (self.tSize + self.fSize)
         with open(self.fileName, "rb") as f:
@@ -171,59 +290,118 @@ def readField(self, idx):
         return t, field
 
     def reshape(self, field):
+        """Eventually reshape the field to correspond to the grid structure"""
         pass
 
 
 @FieldsIO.register(sID=0)
 class Scal0D(FieldsIO):
+    """FieldsIO handler storing a given number of scalar"""
 
+    # -------------------------------------------------------------------------
+    # Overriden methods
+    # -------------------------------------------------------------------------
     def setHeader(self, nVar):
+        """
+        Set the descriptive grid structure to be stored in the file header.
+
+        Parameters
+        ----------
+        nVar : int
+            Number of scalar variable stored.
+        """
         self.header = {"nVar": int(nVar)}
         self.nItems = self.nVar
 
     @property
-    def nVar(self):
-        return self.header["nVar"]
-
-    @property
-    def hInfos(self):
+    def hinfos(self):
+        """Array representing the grid structure to be written in the binary file."""
         return [np.array([self.nVar], dtype=np.int64)]
 
     def readHeader(self, f):
+        """
+        Read the header from the binary file storing the fields.
+
+        Parameters
+        ----------
+        f : `_io.TextIOWrapper`
+            File to read the header from.
+        """
         (nVar,) = np.fromfile(f, dtype=np.int64, count=1)
         self.setHeader(nVar)
 
+    # -------------------------------------------------------------------------
+    # Class specifics
+    # -------------------------------------------------------------------------
+    @property
+    def nVar(self):
+        """Number of variables in a fields, as described in the header"""
+        return self.header["nVar"]
+
 
 @FieldsIO.register(sID=1)
 class Cart1D(Scal0D):
+    """FieldsIO handler storing a given number of 2D scalar variables"""
 
-    def setupGrids(self, **grids):
-        grids = {name: np.asarray(grid, dtype=np.float64) for name, grid in grids.items() if name.startswith("grid")}
-        for name, grid in grids.items():
-            assert grid.ndim == 1, f"{name} must be one dimensional"
-        return grids
-
+    # -------------------------------------------------------------------------
+    # Overriden methods
+    # -------------------------------------------------------------------------
     def setHeader(self, nVar, gridX):
+        """
+        Set the descriptive grid structure to be stored in the file header.
+
+        Parameters
+        ----------
+        nVar : int
+            Number of 1D variables stored.
+        gridX : np.1darray
+            The grid coordinates in X direction.
+
+        Note
+        ----
+        When used in MPI decomposition mode, `gridX` **must** be the global grid.
+        """
         grids = self.setupGrids(gridX=gridX)
         self.header = {"nVar": int(nVar), **grids}
         self.nItems = nVar * self.nX
 
-    @property
-    def nX(self):
-        return self.header["gridX"].size
-
-    def reshape(self, fields: np.ndarray):
-        fields.shape = (self.nVar, self.nX)
-
     @property
     def hInfos(self):
+        """Array representing the grid structure to be written in the binary file."""
         return [np.array([self.nVar, self.nX], dtype=np.int64), np.array(self.header["gridX"], dtype=np.float64)]
 
     def readHeader(self, f):
+        """
+        Read the header from the binary file storing the fields.
+
+        Parameters
+        ----------
+        f : `_io.TextIOWrapper`
+            File to read the header from.
+        """
         nVar, nX = np.fromfile(f, dtype=np.int64, count=2)
         gridX = np.fromfile(f, dtype=np.float64, count=nX)
         self.setHeader(nVar, gridX)
 
+    def reshape(self, fields: np.ndarray):
+        fields.shape = (self.nVar, self.nX)
+
+    # -------------------------------------------------------------------------
+    # Class specifics
+    # -------------------------------------------------------------------------
+    @property
+    def nX(self):
+        """Number of points in x direction"""
+        return self.header["gridX"].size
+
+    @staticmethod
+    def setupGrids(**grids):
+        """Utility function to setup grids in multuple dimensions, given the keyword arguments"""
+        grids = {name: np.asarray(grid, dtype=np.float64) for name, grid in grids.items() if name.startswith("grid")}
+        for name, grid in grids.items():
+            assert grid.ndim == 1, f"{name} must be one dimensional"
+        return grids
+
     # -------------------------------------------------------------------------
     # MPI-parallel implementation
     # -------------------------------------------------------------------------
@@ -231,6 +409,19 @@ def readHeader(self, f):
 
     @classmethod
     def setupMPI(cls, comm: MPI.Intracomm, iLocX, nLocX):
+        """
+        Setup the MPI mode for the files IO, considering a decomposition
+        of the 1D grid into contiuous subintervals.
+
+        Parameters
+        ----------
+        comm : MPI.Intracomm
+            The space decomposition communicator.
+        iLocX : int
+            Starting index of the local sub-domain in the global `gridX`.
+        nLocX : int
+            Number of points in the local sub-domain.
+        """
         cls.comm = comm
         cls.iLocX = iLocX
         cls.nLocX = nLocX
@@ -238,17 +429,20 @@ def setupMPI(cls, comm: MPI.Intracomm, iLocX, nLocX):
 
     @property
     def MPI_ON(self):
+        """Wether or not MPI is activated"""
         if self.comm is None:
             return False
         return self.comm.Get_size() > 1
 
     @property
     def MPI_ROOT(self):
+        """Wether or not the process is MPI Root"""
         if self.comm is None:
             return True
         return self.comm.Get_rank() == 0
 
     def MPI_FILE_OPEN(self, mode):
+        """Open the binary file in MPI mode"""
         amode = {
             "r": MPI.MODE_RDONLY,
             "a": MPI.MODE_WRONLY | MPI.MODE_APPEND,
@@ -256,26 +450,66 @@ def MPI_FILE_OPEN(self, mode):
         self.mpiFile = MPI.File.Open(self.comm, self.fileName, amode)
 
     def MPI_WRITE(self, data):
+        """Write data (np.ndarray) in the binary file in MPI mode, at the current file cursor position."""
         self.mpiFile.Write(data)
 
     def MPI_WRITE_AT(self, offset, data: np.ndarray):
+        """
+        Write data in the binary file in MPI mode, with a given offset
+        **relative to the beginning of the file**.
+
+        Parameters
+        ----------
+        offset : int
+            Offset to write at, relative to the beginning of the file, in bytes.
+        data : np.ndarray
+            Data to be written in the binary file.
+        """
         self.mpiFile.Write_at(offset, data)
 
     def MPI_READ_AT(self, offset, data):
+        """
+        Read data from the binary file in MPI mode, with a given offset
+        **relative to the beginning of the file**.
+
+        Parameters
+        ----------
+        offset : int
+            Offset to read at, relative to the beginning of the file, in bytes.
+        data : np.ndarray
+            Array on which to read the data from the binary file.
+        """
         self.mpiFile.Read_at(offset, data)
 
     def MPI_FILE_CLOSE(self):
+        """Close the binary file in MPI mode"""
         self.mpiFile.Close()
         self.mpiFile = None
 
     def initialize(self):
+        """Initialize the binary file (write header) in MPI mode"""
         if self.MPI_ROOT:
             super().initialize()
         if self.MPI_ON:
-            self.comm.Barrier()
+            self.comm.Barrier()  # Important, should not be removed !
             self.initialized = True
 
     def addField(self, time, field):
+        """
+        Append one field solution at the end of the file with one given time,
+        eventually using MPI.
+
+        Parameters
+        ----------
+        time : float-like
+            The associated time of the field solution.
+        field : np.ndarray
+            The (local) field values.
+
+        Note
+        ----
+        If a MPI decomposition is used, field **must be** the local field values.
+        """
         if not self.MPI_ON:
             return super().addField(time, field)
 
@@ -297,6 +531,26 @@ def addField(self, time, field):
         self.MPI_FILE_CLOSE()
 
     def readField(self, idx):
+        """
+        Read one field stored in the binary file, corresponding to the given
+        time index, eventually in MPI mode.
+
+        Parameters
+        ----------
+        idx : int
+            Positional index of the field.
+
+        Returns
+        -------
+        t : float
+            Stored time for this field.
+        field : np.ndarray
+            Read (local) fields in a numpy array.
+
+        Note
+        ----
+        If a MPI decomposition is used, it reads and returns the local fields values only.
+        """
         if not self.MPI_ON:
             return super().readField(idx)
         idx = self.check(idx)
@@ -319,21 +573,35 @@ def readField(self, idx):
 
 @FieldsIO.register(sID=2)
 class Cart2D(Cart1D):
+    """FieldsIO handler storing a given number of 2D scalar variables"""
 
+    # -------------------------------------------------------------------------
+    # Overriden methods
+    # -------------------------------------------------------------------------
     def setHeader(self, nVar, gridX, gridY):
+        """
+        Set the descriptive grid structure to be stored in the file header.
+
+        Parameters
+        ----------
+        nVar : int
+            Number of 1D variables stored.
+        gridX : np.1darray
+            The grid coordinates in x direction.
+        gridY : np.1darray
+            The grid coordinates in y direction.
+
+        Note
+        ----
+        When used in MPI decomposition mode, `gridX` and `gridY` **must** be the global grid.
+        """
         grids = self.setupGrids(gridX=gridX, gridY=gridY)
         self.header = {"nVar": int(nVar), **grids}
         self.nItems = nVar * self.nX * self.nY
 
-    @property
-    def nY(self):
-        return self.header["gridY"].size
-
-    def reshape(self, fields: np.ndarray):
-        fields.shape = (self.nVar, self.nX, self.nY)
-
     @property
     def hInfos(self):
+        """Array representing the grid structure to be written in the binary file."""
         return [
             np.array([self.nVar, self.nX, self.nY], dtype=np.int64),
             np.array(self.header["gridX"], dtype=np.float64),
@@ -341,11 +609,31 @@ def hInfos(self):
         ]
 
     def readHeader(self, f):
+        """
+        Read the header from the binary file storing the fields.
+
+        Parameters
+        ----------
+        f : `_io.TextIOWrapper`
+            File to read the header from.
+        """
         nVar, nX, nY = np.fromfile(f, dtype=np.int64, count=3)
         gridX = np.fromfile(f, dtype=np.float64, count=nX)
         gridY = np.fromfile(f, dtype=np.float64, count=nY)
         self.setHeader(nVar, gridX, gridY)
 
+    def reshape(self, fields: np.ndarray):
+        """Reshape the fields to a [nVar, nX, nY] array (inplace operation)"""
+        fields.shape = (self.nVar, self.nX, self.nY)
+
+    # -------------------------------------------------------------------------
+    # Class specifics
+    # -------------------------------------------------------------------------
+    @property
+    def nY(self):
+        """Number of points in y direction"""
+        return self.header["gridY"].size
+
     # -------------------------------------------------------------------------
     # MPI-parallel implementation
     # -------------------------------------------------------------------------
@@ -356,6 +644,21 @@ def setupMPI(cls, comm: MPI.Intracomm, iLocX, nLocX, iLocY, nLocY):
         cls.nLocY = nLocY
 
     def addField(self, time, field):
+        """
+        Append one field solution at the end of the file with one given time,
+        eventually using MPI.
+
+        Parameters
+        ----------
+        time : float-like
+            The associated time of the field solution.
+        field : np.ndarray
+            The (local) field values.
+
+        Note
+        ----
+        If a MPI decomposition is used, field **must be** the local field values.
+        """
         if not self.MPI_ON:
             return super().addField(time, field)
 
@@ -382,6 +685,26 @@ def addField(self, time, field):
         self.MPI_FILE_CLOSE()
 
     def readField(self, idx):
+        """
+        Read one field stored in the binary file, corresponding to the given
+        time index, eventually in MPI mode.
+
+        Parameters
+        ----------
+        idx : int
+            Positional index of the field.
+
+        Returns
+        -------
+        t : float
+            Stored time for this field.
+        field : np.ndarray
+            Read (local) fields in a numpy array.
+
+        Note
+        ----
+        If a MPI decomposition is used, it reads and returns the local fields values only.
+        """
         if not self.MPI_ON:
             return super().readField(idx)
         idx = self.check(idx)
diff --git a/pySDC/tests/test_helpers/test_fieldsIO.py b/pySDC/tests/test_helpers/test_fieldsIO.py
index 9e5d77efa4..e492284296 100644
--- a/pySDC/tests/test_helpers/test_fieldsIO.py
+++ b/pySDC/tests/test_helpers/test_fieldsIO.py
@@ -216,30 +216,21 @@ def writeFields_MPI(fileName, nDim, dtypeIdx, algo, nSteps, nVar, nX, nY=None):
     grids, gridSizes, u0 = initGrid(nVar, nX, nY)
 
     from mpi4py import MPI
+    from pySDC.helpers.blocks import BlockDecomposition
+    from pySDC.helpers.fieldsIO import Cart1D, Cart2D
 
     comm = MPI.COMM_WORLD
     MPI_SIZE = comm.Get_size()
     MPI_RANK = comm.Get_rank()
 
-    from pySDC.helpers.blocks import BlockDecomposition
-
     blocks = BlockDecomposition(MPI_SIZE, gridSizes, algo, MPI_RANK)
 
-    from time import sleep
-
-    from pySDC.helpers.fieldsIO import Cart1D, Cart2D
-
     if nDim == 1:
         (iLocX,), (nLocX,) = blocks.localBounds
         (pRankX,) = blocks.ranks
         Cart1D.setupMPI(comm, iLocX, nLocX)
         u0 = u0[:, iLocX : iLocX + nLocX]
 
-        MPI.COMM_WORLD.Barrier()
-        sleep(0.01 * MPI_RANK)
-        print(f"[Rank {MPI_RANK}] pRankX={pRankX} ({iLocX}, {nLocX})")
-        MPI.COMM_WORLD.Barrier()
-
         f1 = Cart1D(DTYPES[dtypeIdx], fileName)
         f1.setHeader(nVar=nVar, gridX=grids[0])
 

From f08b9f4ce51ca17c25c52434d09946954f0a98b7 Mon Sep 17 00:00:00 2001
From: Thibaut Lunet <thibaut.lunet@tuhh.de>
Date: Fri, 17 Jan 2025 09:21:09 +0100
Subject: [PATCH 06/11] TL: added a small comment

---
 pySDC/helpers/blocks.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pySDC/helpers/blocks.py b/pySDC/helpers/blocks.py
index ebd38f6c63..c7c019c78e 100644
--- a/pySDC/helpers/blocks.py
+++ b/pySDC/helpers/blocks.py
@@ -118,6 +118,7 @@ def localBounds(self):
 
 
 if __name__ == "__main__":
+    # Base usage of this module for a 2D decomposition
     from mpi4py import MPI
     from time import sleep
 

From d63e93e6b8ede1e4638a0e971c9b7ebc96e5c121 Mon Sep 17 00:00:00 2001
From: Thibaut Lunet <thibaut.lunet@tuhh.de>
Date: Fri, 17 Jan 2025 09:47:45 +0100
Subject: [PATCH 07/11] TL: tentative to solve static typing issue with no
 mpi4py

---
 pySDC/helpers/fieldsIO.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py
index d7b32d9804..7329331174 100644
--- a/pySDC/helpers/fieldsIO.py
+++ b/pySDC/helpers/fieldsIO.py
@@ -43,15 +43,18 @@
 import numpy as np
 from typing import Type, TypeVar
 
+T = TypeVar("T")
+
 try:
     from mpi4py import MPI
 except ImportError:
 
     class MPI:
         COMM_WORLD = None
+        Intracomm = T
+
 
 
-T = TypeVar("T")
 
 # Supported data types
 DTYPES = {

From 1f5a201f0bcf6191ab3b282512463d821c9a4ab2 Mon Sep 17 00:00:00 2001
From: Thibaut Lunet <thibaut.lunet@tuhh.de>
Date: Fri, 17 Jan 2025 09:48:27 +0100
Subject: [PATCH 08/11] TL: forgot black, ofc

---
 pySDC/helpers/fieldsIO.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py
index 7329331174..8027dda484 100644
--- a/pySDC/helpers/fieldsIO.py
+++ b/pySDC/helpers/fieldsIO.py
@@ -54,8 +54,6 @@ class MPI:
         Intracomm = T
 
 
-
-
 # Supported data types
 DTYPES = {
     0: np.float64,  # double precision

From cf130ef666b3595e13f650fc2a838de4b8ced6a9 Mon Sep 17 00:00:00 2001
From: Thibaut Lunet <thibaut.lunet@tuhh.de>
Date: Fri, 17 Jan 2025 10:03:39 +0100
Subject: [PATCH 09/11] =?UTF-8?q?TL:=20introducing=20modern=20python=20tac?=
 =?UTF-8?q?tics=20in=20ci=20=F0=9F=A4=93?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .github/workflows/ci_pipeline.yml | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/.github/workflows/ci_pipeline.yml b/.github/workflows/ci_pipeline.yml
index dee29be4f4..5821c388d3 100644
--- a/.github/workflows/ci_pipeline.yml
+++ b/.github/workflows/ci_pipeline.yml
@@ -55,6 +55,10 @@ jobs:
       - name: Install additional packages as needed
         run: |
           micromamba install -y --file etc/environment-tests.yml --freeze-installed
+      - name: Install pySDC as a package in the current environment
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e --no-deps .
       - name: Run pytest for CPU stuff
         run: |
           echo "print('Loading sitecustomize.py...')

From ddf2521137018e13c1d33141677719983aec3124 Mon Sep 17 00:00:00 2001
From: Thibaut Lunet <thibaut.lunet@tuhh.de>
Date: Fri, 17 Jan 2025 10:10:31 +0100
Subject: [PATCH 10/11] TL: attempt to solve the fenics tests

---
 pySDC/helpers/fieldsIO.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py
index 8027dda484..7d5f549e9a 100644
--- a/pySDC/helpers/fieldsIO.py
+++ b/pySDC/helpers/fieldsIO.py
@@ -46,6 +46,10 @@
 T = TypeVar("T")
 
 try:
+    try:
+        import dolfin as df  # noqa: F841 (for some reason, dolfin always needs to be imported before mpi4py)
+    except ImportError:
+        pass
     from mpi4py import MPI
 except ImportError:
 

From 1dc21db15126aca9ca2d37c10a97a8e453837b51 Mon Sep 17 00:00:00 2001
From: Thibaut Lunet <thibaut.lunet@tuhh.de>
Date: Fri, 17 Jan 2025 10:14:23 +0100
Subject: [PATCH 11/11] =?UTF-8?q?TL:=20pip=20is=20a=20little=20trickster?=
 =?UTF-8?q?=20=F0=9F=98=85?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .github/workflows/ci_pipeline.yml | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/.github/workflows/ci_pipeline.yml b/.github/workflows/ci_pipeline.yml
index 5821c388d3..4ba3557084 100644
--- a/.github/workflows/ci_pipeline.yml
+++ b/.github/workflows/ci_pipeline.yml
@@ -57,8 +57,7 @@ jobs:
           micromamba install -y --file etc/environment-tests.yml --freeze-installed
       - name: Install pySDC as a package in the current environment
         run: |
-          python -m pip install --upgrade pip
-          pip install -e --no-deps .
+          pip install --no-deps -e .
       - name: Run pytest for CPU stuff
         run: |
           echo "print('Loading sitecustomize.py...')