From d772718177ae86a4add07521e921f66e52ef234f Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 00:06:24 +0100 Subject: [PATCH 01/16] TL: simplified fieldsIO to two generic classes (Scalar + Rectilinear) --- pySDC/helpers/blocks.py | 16 +- pySDC/helpers/fieldsIO.py | 277 ++++++-------------- pySDC/tests/test_helpers/test_fieldsIO.py | 298 +++++++--------------- 3 files changed, 173 insertions(+), 418 deletions(-) diff --git a/pySDC/helpers/blocks.py b/pySDC/helpers/blocks.py index 80482e712..d6c846aa9 100644 --- a/pySDC/helpers/blocks.py +++ b/pySDC/helpers/blocks.py @@ -1,3 +1,6 @@ +import numpy as np + + class BlockDecomposition(object): """ Class decomposing a cartesian space domain (1D to 3D) into a given number of processors. @@ -92,16 +95,9 @@ def __init__(self, nProcs, gridSizes, algo="Hybrid", gRank=None, order="C"): @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}") + assert self.gRank is not None, "gRank attribute needs to be set" + cart = np.arange(np.prod(self.nBlocks)).reshape(self.nBlocks, order=self.order) + return list(np.argwhere(cart == self.gRank)[0]) @property def localBounds(self): diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py index 82675ac65..ea8c403cd 100644 --- a/pySDC/helpers/fieldsIO.py +++ b/pySDC/helpers/fieldsIO.py @@ -4,9 +4,8 @@ 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 +- :class:`Scalar` : for 0D fields (scalar) with a given number of variables +- :class:`Rectilinear` : for fields on N-dimensional rectilinear grids 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 @@ -15,12 +14,12 @@ Example ------- >>> import numpy as np ->>> from pySDC.helpers.fieldsIO import Cart1D, FieldsIO +>>> from pySDC.helpers.fieldsIO import Rectilinear, FieldsIO >>> >>> # Write some fields in files >>> x = np.linspace(0, 1, 101) ->>> fOut = Cart2D(np.float64, "file.pysdc") ->>> fOut.setHeader(nVar=2, coordX=x) +>>> fOut = Rectilinear(np.float64, "file.pysdc") +>>> fOut.setHeader(nVar=2, coords=x) >>> fOut.initialize() >>> times = [0, 1, 2] >>> u0 = np.array([-1, 1])[:, None]*x[None, :] @@ -36,20 +35,19 @@ Notes ----- -🚀 :class:`Cart1D` and :class:`Cart2D` are compatible with a MPI-based cartesian decomposition. +🚀 :class:`Rectilinear` is compatible with a MPI-based cartesian decomposition. See :class:`pySDC.tests.test_helpers.test_fieldsIO.writeFields_MPI` for an illustrative example. Warning ------- -To use MPI collective writing, you need to call first the class methods :class:`Cart1D.initMPI` -or :class:`Cart2D.initMPI` from the associated class (cf their docstring). -Also, their associated `setHeader` methods **must be given the global grids coordinates**, -wether code is run in parallel or not. +To use MPI collective writing, you need to call first the class methods :class:`Rectilinear.initMPI` (cf their docstring). +Also, `Rectilinear.setHeader` **must be given the global grids coordinates**, wether the code is run in parallel or not. """ import os import numpy as np from typing import Type, TypeVar import logging +import itertools T = TypeVar("T") @@ -330,11 +328,11 @@ def reshape(self, field): @FieldsIO.register(sID=0) -class Scal0D(FieldsIO): +class Scalar(FieldsIO): """FieldsIO handler storing a given number of scalar""" # ------------------------------------------------------------------------- - # Overriden methods + # Overridden methods # ------------------------------------------------------------------------- def setHeader(self, nVar): """ @@ -375,13 +373,21 @@ def nVar(self): @FieldsIO.register(sID=1) -class Cart1D(Scal0D): - """FieldsIO handler storing a given number of 2D scalar variables""" +class Rectilinear(Scalar): + """FieldsIO handler storing a given number of scalar variables on a N-dimensional rectilinear grid""" + + @staticmethod + def setupCoords(*coords): + """Utility function to setup grids in multiple dimensions, given the keyword arguments""" + coords = [np.asarray(coord, dtype=np.float64) for coord in coords] + for axis, coord in enumerate(coords): + assert coord.ndim == 1, f"coord for {axis=} must be one dimensional" + return coords # ------------------------------------------------------------------------- - # Overriden methods + # Overridden methods # ------------------------------------------------------------------------- - def setHeader(self, nVar, coordX): + def setHeader(self, nVar, coords): """ Set the descriptive grid structure to be stored in the file header. @@ -389,21 +395,25 @@ def setHeader(self, nVar, coordX): ---------- nVar : int Number of 1D variables stored. - coordX : np.1darray - The grid coordinates in X direction. + coords : np.1darray or list[np.1darray] + The grid coordinates in each dimensions. Note ---- - When used in MPI decomposition mode, `coordX` **must** be the global grid. + When used in MPI decomposition mode, all coordinate **must** be the global grid. """ - coords = self.setupCoords(coordX=coordX) - self.header = {"nVar": int(nVar), **coords} - self.nItems = nVar * self.nX + if not isinstance(coords, (tuple, list)): + coords = [coords] + coords = self.setupCoords(*coords) + self.header = {"nVar": int(nVar), "coords": coords} + self.nItems = nVar * self.nDoF @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["coordX"], dtype=np.float64)] + return [np.array([self.nVar, self.dim, *self.nX], dtype=np.int32)] + [ + np.array(coord, dtype=np.float64) for coord in self.header["coords"] + ] def readHeader(self, f): """ @@ -414,28 +424,32 @@ def readHeader(self, f): f : `_io.TextIOWrapper` File to read the header from. """ - nVar, nX = np.fromfile(f, dtype=np.int64, count=2) - coordX = np.fromfile(f, dtype=np.float64, count=nX) - self.setHeader(nVar, coordX) + nVar, dim = np.fromfile(f, dtype=np.int32, count=2) + nX = np.fromfile(f, dtype=np.int32, count=dim) + coords = [np.fromfile(f, dtype=np.float64, count=n) for n in nX] + self.setHeader(nVar, coords) def reshape(self, fields: np.ndarray): - fields.shape = (self.nVar, self.nX) + """Reshape the fields to a N-d array (inplace operation)""" + fields.shape = (self.nVar, *self.nX) # ------------------------------------------------------------------------- # Class specifics # ------------------------------------------------------------------------- @property def nX(self): - """Number of points in x direction""" - return self.header["coordX"].size + """Number of points in y direction""" + return [coord.size for coord in self.header["coords"]] - @staticmethod - def setupCoords(**coords): - """Utility function to setup grids in multuple dimensions, given the keyword arguments""" - coords = {name: np.asarray(coord, dtype=np.float64) for name, coord in coords.items()} - for name, coord in coords.items(): - assert coord.ndim == 1, f"{name} must be one dimensional" - return coords + @property + def dim(self): + """Number of grid dimensions""" + return len(self.nX) + + @property + def nDoF(self): + """Number of degrees of freedom for one variable""" + return np.prod(self.nX) # ------------------------------------------------------------------------- # MPI-parallel implementation @@ -443,7 +457,7 @@ def setupCoords(**coords): comm: MPI.Intracomm = None @classmethod - def setupMPI(cls, comm: MPI.Intracomm, iLocX, nLocX): + def setupMPI(cls, comm: MPI.Intracomm, iLoc, nLoc): """ Setup the MPI mode for the files IO, considering a decomposition of the 1D grid into contiuous subintervals. @@ -452,14 +466,14 @@ def setupMPI(cls, comm: MPI.Intracomm, iLocX, nLocX): ---------- comm : MPI.Intracomm The space decomposition communicator. - iLocX : int + iLoc : list[int] Starting index of the local sub-domain in the global `coordX`. - nLocX : int + nLoc : list[int] Number of points in the local sub-domain. """ cls.comm = comm - cls.iLocX = iLocX - cls.nLocX = nLocX + cls.iLoc = iLoc + cls.nLoc = nLoc cls.mpiFile = None @property @@ -537,155 +551,6 @@ def initialize(self): 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, - possibly 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) - - 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): - """ - Read one field stored in the binary file, corresponding to the given - time index, possibly 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.formatIndex(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): - """FieldsIO handler storing a given number of 2D scalar variables""" - - # ------------------------------------------------------------------------- - # Overriden methods - # ------------------------------------------------------------------------- - def setHeader(self, nVar, coordX, coordY): - """ - Set the descriptive grid structure to be stored in the file header. - - Parameters - ---------- - nVar : int - Number of 1D variables stored. - coordX : np.1darray - The grid coordinates in x direction. - coordY : np.1darray - The grid coordinates in y direction. - - Note - ---- - When used in MPI decomposition mode, `coordX` and `coordX` **must** be the global grid. - """ - coords = self.setupCoords(coordX=coordX, coordY=coordY) - self.header = {"nVar": int(nVar), **coords} - self.nItems = 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["coordX"], dtype=np.float64), - np.array(self.header["coordY"], 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, nY = np.fromfile(f, dtype=np.int64, count=3) - coordX = np.fromfile(f, dtype=np.float64, count=nX) - coordY = np.fromfile(f, dtype=np.float64, count=nY) - self.setHeader(nVar, coordX, coordY) - - 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["coordY"].size - - # ------------------------------------------------------------------------- - # 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): """ Append one field solution at the end of the file with one given time, @@ -711,9 +576,8 @@ def addField(self, time, 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}" + *self.nLoc, + ), f"expected {(self.nVar, *self.nLoc)} shape, got {field.shape}" offset0 = self.fileSize self.MPI_FILE_OPEN(mode="a") @@ -721,12 +585,18 @@ def addField(self, time, field): 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]) + for (iVar, *iX) in itertools.product(range(self.nVar), *[range(nX) for nX in self.nLoc[:-1]]): + offset = offset0 + self.iPos(iVar, iX) * self.itemSize + self.MPI_WRITE_AT(offset, field[iVar, *iX]) self.MPI_FILE_CLOSE() + def iPos(self, iVar, iX): + iPos = iVar * self.nDoF + for axis in range(self.dim - 1): + iPos += (self.iLoc[axis] + iX[axis]) * np.prod(self.nX[axis + 1 :]) + iPos += self.iLoc[-1] + return iPos + def readField(self, idx): """ Read one field stored in the binary file, corresponding to the given @@ -750,20 +620,19 @@ def readField(self, idx): """ if not self.MPI_ON: return super().readField(idx) - idx = self.formatIndex(idx) + idx = self.formatIndex(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) + field = np.empty((self.nVar, *self.nLoc), 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]) + for (iVar, *iX) in itertools.product(range(self.nVar), *[range(nX) for nX in self.nLoc[:-1]]): + offset = offset0 + self.iPos(iVar, iX) * 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 index 5d5be0151..4037c6cce 100644 --- a/pySDC/tests/test_helpers/test_fieldsIO.py +++ b/pySDC/tests/test_helpers/test_fieldsIO.py @@ -1,4 +1,5 @@ import pytest +import itertools import numpy as np from pySDC.helpers.fieldsIO import DTYPES, FieldsIO @@ -7,25 +8,21 @@ @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 +@pytest.mark.parametrize("dim", range(4)) +def testHeader(dim, dtypeIdx): + from pySDC.helpers.fieldsIO import FieldsIO, Scalar, Rectilinear fileName = "testHeader.pysdc" dtype = DTYPES[dtypeIdx] - coordX = np.linspace(0, 1, num=256, endpoint=False) - coordY = np.linspace(0, 1, num=64, endpoint=False) + coords = [np.linspace(0, 1, num=256, endpoint=False) for n in [256, 64, 32]] - if nDim == 0: - Class = Scal0D + if dim == 0: + Class = Scalar args = {"nVar": 20} - elif nDim == 1: - Class = Cart1D - args = {"nVar": 10, "coordX": coordX} - elif nDim == 2: - Class = Cart2D - args = {"nVar": 10, "coordX": coordX, "coordY": coordY} + else: + Class = Rectilinear + args = {"nVar": 10, "coords": coords[:dim]} f1 = Class(dtype, fileName) assert f1.__str__() == f1.__repr__(), "__repr__ and __str__ do not return the same result" @@ -64,13 +61,13 @@ def testHeader(nDim, dtypeIdx): @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 +def testScalar(nVar, nSteps, dtypeIdx): + from pySDC.helpers.fieldsIO import FieldsIO, Scalar - fileName = "testScal0D.pysdc" + fileName = "testScalar.pysdc" dtype = DTYPES[dtypeIdx] - f1 = Scal0D(dtype, fileName) + f1 = Scalar(dtype, fileName) f1.setHeader(nVar=nVar) assert f1.nItems == nVar, f"{f1} do not have nItems == nVar" @@ -104,145 +101,83 @@ def testScal0D(nVar, nSteps, dtypeIdx): @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] - - coordX = np.linspace(0, 1, num=nX, endpoint=False) - nX = coordX.size - - f1 = Cart1D(dtype, fileName) - f1.setHeader(nVar=nVar, coordX=coordX) - - 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 +@pytest.mark.parametrize("dim", [1, 2, 3]) +def testRectilinear(dim, nVar, nSteps, dtypeIdx): + from pySDC.helpers.fieldsIO import FieldsIO, Rectilinear, DTYPES - fileName = "testCart2D.pysdc" + fileName = f"testRectilinear{dim}D.pysdc" dtype = DTYPES[dtypeIdx] - coordX = np.linspace(0, 1, num=nX, endpoint=False) - coordY = np.linspace(0, 1, num=nY, endpoint=False) + for nX in itertools.product(*[[5, 10, 16]] * dim): - f1 = Cart2D(dtype, fileName) - f1.setHeader(nVar=nVar, coordX=coordX, coordY=coordY) + coords = [np.linspace(0, 1, num=n, endpoint=False) for n in nX] - 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() + f1 = Rectilinear(dtype, fileName) + f1.setHeader(nVar=nVar, coords=coords) - u0 = np.random.rand(nVar, nX, nY).astype(f1.dtype) - times = np.arange(nSteps) / nSteps + assert f1.dim == dim, f"{f1} has incorrect dimension" + assert f1.nX == list(nX), f"{f1} has incorrect nX" + assert f1.nDoF == np.prod(nX), f"{f1} has incorrect nDOF" + assert f1.nItems == nVar * np.prod(nX), f"{f1} do not have nItems == nVar*nX**dim" - 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" + f1.initialize() + u0 = np.random.rand(nVar, *nX).astype(f1.dtype) + times = np.arange(nSteps) / nSteps - f2 = FieldsIO.fromFile(fileName) + for t in times: + ut = (u0 * t).astype(f1.dtype) + f1.addField(t, ut) - 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})" + 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" - 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" + 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})" -def initGrid(nVar, nX, nY=None): - nDim = 1 - if nY is not None: - nDim += 1 - x = np.linspace(0, 1, num=nX, endpoint=False) - coords = (x,) - gridSizes = (nX,) - u0 = np.array(np.arange(nVar) + 1)[:, None] * x[None, :] + 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" - if nDim > 1: - y = np.linspace(0, 1, num=nY, endpoint=False) - coords += (y,) - gridSizes += (nY,) - u0 = u0[:, :, None] * y[None, None, :] - return coords, gridSizes, u0 +def initGrid(nVar, gridSizes): + dim = len(gridSizes) + coords = [np.linspace(0, 1, num=n, endpoint=False) for n in gridSizes] + u0 = np.array(np.arange(nVar) + 1)[:, *[None] * dim] + for x in np.meshgrid(*coords, indexing="ij"): + u0 = u0 * x + return coords, u0 -def writeFields_MPI(fileName, nDim, dtypeIdx, algo, nSteps, nVar, nX, nY=None): - coords, gridSizes, u0 = initGrid(nVar, nX, nY) +def writeFields_MPI(fileName, dtypeIdx, algo, nSteps, nVar, nX): + coords, u0 = initGrid(nVar, nX) from mpi4py import MPI from pySDC.helpers.blocks import BlockDecomposition - from pySDC.helpers.fieldsIO import Cart1D, Cart2D + from pySDC.helpers.fieldsIO import Rectilinear comm = MPI.COMM_WORLD MPI_SIZE = comm.Get_size() MPI_RANK = comm.Get_rank() - blocks = BlockDecomposition(MPI_SIZE, gridSizes, algo, MPI_RANK) - - if nDim == 1: - (iLocX,), (nLocX,) = blocks.localBounds - (pRankX,) = blocks.ranks - Cart1D.setupMPI(comm, iLocX, nLocX) - u0 = u0[:, iLocX : iLocX + nLocX] + blocks = BlockDecomposition(MPI_SIZE, nX, algo, MPI_RANK) - f1 = Cart1D(DTYPES[dtypeIdx], fileName) - f1.setHeader(nVar=nVar, coordX=coords[0]) + iLoc, nLoc = blocks.localBounds + Rectilinear.setupMPI(comm, iLoc, nLoc) + s = [slice(i, i + n) for i, n in zip(iLoc, nLoc)] + u0 = u0[:, *s] + print(MPI_RANK, u0.shape) - 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, coordX=coords[0], coordY=coords[1]) + f1 = Rectilinear(DTYPES[dtypeIdx], fileName) + f1.setHeader(nVar=nVar, coords=coords) u0 = np.asarray(u0, dtype=f1.dtype) f1.initialize() @@ -264,94 +199,53 @@ def compareFields_MPI(fileName, u0, 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 t2 == t, f"fields[{idx}] in {f2} has incorrect time ({t2} instead of {t})" 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): +@pytest.mark.parametrize("nProcs", [2, 4]) +@pytest.mark.parametrize("dim", [2, 3]) +def testRectilinear_MPI(dim, nProcs, dtypeIdx, algo, nSteps, nVar): 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}" + fileName = f"testRectilinear{dim}D_MPI.pysdc" - 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}" + for nX in itertools.product(*[[61, 16]] * dim): - from pySDC.helpers.fieldsIO import FieldsIO, Cart1D + cmd = f"mpirun -np {nProcs} python {__file__} --fileName {fileName}" + cmd += f" --dtypeIdx {dtypeIdx} --algo {algo} --nSteps {nSteps} --nVar {nVar} --nX {' '.join([str(n) for n in nX])}" - f2: Cart1D = FieldsIO.fromFile(fileName) + p = subprocess.Popen(cmd.split(), cwd=".") + p.wait() + assert p.returncode == 0, f"MPI write with {nProcs} proc(s) did not return code 0, but {p.returncode}" - 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}" + from pySDC.helpers.fieldsIO import FieldsIO, Rectilinear - coords, _, u0 = initGrid(nVar, nX) - assert np.allclose(f2.header['coordX'], coords[0]), f"incorrect coordX in MPI written fields {f2}" + f2: Rectilinear = 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" + assert type(f2) == Rectilinear, f"incorrect type in MPI written fields {f2}" + assert f2.nFields == nSteps, f"incorrect nFields in MPI written fields {f2} ({f2.nFields} instead of {nSteps})" + assert f2.nVar == nVar, f"incorrect nVar in MPI written fields {f2}" + assert f2.nX == list(nX), f"incorrect nX in MPI written fields {f2}" + coords, u0 = initGrid(nVar, nX) + for i, (cFile, cRef) in enumerate(zip(f2.header['coords'], coords)): + assert np.allclose(cFile, cRef), f"incorrect coords[{i}] in MPI written fields {f2}" -@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['coordX'], grids[0]), f"incorrect coordX in MPI written fields {f2}" - assert np.allclose(f2.header['coordY'], grids[1]), f"incorrect coordY 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" + times = np.arange(nSteps) / nSteps + for idx, t in enumerate(times): + u1 = u0 * t + t2, u2 = f2.readField(idx) + assert t2 == t, f"fields[{idx}] in {f2} has incorrect time ({t2} instead of {t})" + 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__": @@ -359,15 +253,11 @@ def testCart2D_MPI(nProcs, dtypeIdx, algo, nSteps, nVar, nX, nY): 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('--algo', type=str, help="algorithm used for block decomposition") + parser.add_argument('--nSteps', type=int, help="number of time-steps") 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") + parser.add_argument('--nX', type=int, nargs='+', help="number of grid points in each dimensions") args = parser.parse_args() u0 = writeFields_MPI(**args.__dict__) From 0d2a7eb9a113c1536dd9cddf451b39222ddf1f64 Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 00:26:44 +0100 Subject: [PATCH 02/16] TL: trying something --- pySDC/tests/test_helpers/test_fieldsIO.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pySDC/tests/test_helpers/test_fieldsIO.py b/pySDC/tests/test_helpers/test_fieldsIO.py index 4037c6cce..d2a2339f4 100644 --- a/pySDC/tests/test_helpers/test_fieldsIO.py +++ b/pySDC/tests/test_helpers/test_fieldsIO.py @@ -151,7 +151,8 @@ def testRectilinear(dim, nVar, nSteps, dtypeIdx): def initGrid(nVar, gridSizes): dim = len(gridSizes) coords = [np.linspace(0, 1, num=n, endpoint=False) for n in gridSizes] - u0 = np.array(np.arange(nVar) + 1)[:, *[None] * dim] + s = [None] * dim + u0 = np.array(np.arange(nVar) + 1)[:, *s] for x in np.meshgrid(*coords, indexing="ij"): u0 = u0 * x return coords, u0 From 53c0a29f3b0ce75b162c57d6df47b76fa77b907a Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 00:47:48 +0100 Subject: [PATCH 03/16] TL: restriction to python 3.11 --- pySDC/helpers/fieldsIO.py | 1175 +++++++++++---------- pySDC/tests/test_helpers/test_fieldsIO.py | 397 ++++--- 2 files changed, 786 insertions(+), 786 deletions(-) diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py index ea8c403cd..57a379907 100644 --- a/pySDC/helpers/fieldsIO.py +++ b/pySDC/helpers/fieldsIO.py @@ -42,597 +42,600 @@ ------- To use MPI collective writing, you need to call first the class methods :class:`Rectilinear.initMPI` (cf their docstring). Also, `Rectilinear.setHeader` **must be given the global grids coordinates**, wether the code is run in parallel or not. + +Also : this feature is only available with Python 3.11 or higher ... """ -import os -import numpy as np -from typing import Type, TypeVar -import logging -import itertools +import sys + +if sys.version_info >= (3, 11): + import os + import numpy as np + from typing import Type, TypeVar + import logging + import itertools -T = TypeVar("T") + T = TypeVar("T") -try: try: - import dolfin as df # noqa: F841 (for some reason, dolfin always needs to be imported before mpi4py) + 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: - pass - from mpi4py import MPI -except ImportError: - - class MPI: - COMM_WORLD = None - Intracomm = T - - -# Supported data types -DTYPES = { - 0: np.float64, # double precision - 1: np.complex128, -} -try: - DTYPES.update( - { - 2: np.float128, # quadruple precision - 3: np.complex256, - } - ) -except AttributeError: - logging.getLogger('FieldsIO').debug('Warning: Quadruple precision not available on this machine') -try: - DTYPES.update( - { - 4: np.float32, # single precision - 5: np.complex64, - } - ) -except AttributeError: - logging.getLogger('FieldsIO').debug('Warning: Single precision not available on this machine') - -DTYPES_AVAIL = {val: key for key, val in DTYPES.items()} - -# Header dtype -H_DTYPE = np.int8 -T_DTYPE = np.float64 - - -class FieldsIO: - """Abstract IO file handler""" - - STRUCTS = {} # supported structures, modified dynamically - sID = None # structure ID of the FieldsIO class, modified dynamically - - tSize = T_DTYPE().itemsize - - ALLOW_OVERWRITE = False - - 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. Supported on this machine: {list(DTYPES_AVAIL.keys())}" - 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): - """ - 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) - fieldsIO: FieldsIO = cls.STRUCTS[STRUCT](DTYPES[DTYPE], fileName) - fieldsIO.readHeader(f) - fieldsIO.initialized = True - return fieldsIO - - @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]: + + class MPI: + COMM_WORLD = None + Intracomm = T + + # Supported data types + DTYPES = { + 0: np.float64, # double precision + 1: np.complex128, + } + try: + DTYPES.update( + { + 2: np.float128, # quadruple precision + 3: np.complex256, + } + ) + except AttributeError: + logging.getLogger('FieldsIO').debug('Warning: Quadruple precision not available on this machine') + try: + DTYPES.update( + { + 4: np.float32, # single precision + 5: np.complex64, + } + ) + except AttributeError: + logging.getLogger('FieldsIO').debug('Warning: Single precision not available on this machine') + + DTYPES_AVAIL = {val: key for key, val in DTYPES.items()} + + # Header dtype + H_DTYPE = np.int8 + T_DTYPE = np.float64 + + class FieldsIO: + """Abstract IO file handler""" + + STRUCTS = {} # supported structures, modified dynamically + sID = None # structure ID of the FieldsIO class, modified dynamically + + tSize = T_DTYPE().itemsize + + ALLOW_OVERWRITE = False + + def __init__(self, dtype, fileName): + """ + Parameters + ---------- + dtype : np.dtype + The data type of the fields values. + fileName : str + File. + """ 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): - """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" - - if not self.ALLOW_OVERWRITE: - assert not os.path.isfile( - self.fileName - ), "file already exists, use FieldsIO.ALLOW_OVERWRITE = True to allow overwriting" - - 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): - """(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 - 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): - """Size of one field value (in bytes)""" - return self.dtype().itemsize - - @property - def fSize(self): - """Full size of a field (in bytes)""" - return self.nItems * self.itemSize - - @property - def fileSize(self): - """Current size of the file (in bytes)""" - return os.path.getsize(self.fileName) - - 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}" - 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): - """Number of fields currently stored in the binary file""" - return int((self.fileSize - self.hSize) // (self.tSize + self.fSize)) - - def formatIndex(self, idx): - """Utility method to format a fields index to a positional integer (negative starts from last field index, like python lists)""" - 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): - """Vector of all times stored in the binary file""" - 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): - """Time stored at a given field index""" - idx = self.formatIndex(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): - """ - 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.formatIndex(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): - """Eventually reshape the field to correspond to the grid structure""" - pass - - -@FieldsIO.register(sID=0) -class Scalar(FieldsIO): - """FieldsIO handler storing a given number of scalar""" - - # ------------------------------------------------------------------------- - # Overridden 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 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 Rectilinear(Scalar): - """FieldsIO handler storing a given number of scalar variables on a N-dimensional rectilinear grid""" - - @staticmethod - def setupCoords(*coords): - """Utility function to setup grids in multiple dimensions, given the keyword arguments""" - coords = [np.asarray(coord, dtype=np.float64) for coord in coords] - for axis, coord in enumerate(coords): - assert coord.ndim == 1, f"coord for {axis=} must be one dimensional" - return coords - - # ------------------------------------------------------------------------- - # Overridden methods - # ------------------------------------------------------------------------- - def setHeader(self, nVar, coords): - """ - Set the descriptive grid structure to be stored in the file header. - - Parameters - ---------- - nVar : int - Number of 1D variables stored. - coords : np.1darray or list[np.1darray] - The grid coordinates in each dimensions. - - Note - ---- - When used in MPI decomposition mode, all coordinate **must** be the global grid. - """ - if not isinstance(coords, (tuple, list)): - coords = [coords] - coords = self.setupCoords(*coords) - self.header = {"nVar": int(nVar), "coords": coords} - self.nItems = nVar * self.nDoF - - @property - def hInfos(self): - """Array representing the grid structure to be written in the binary file.""" - return [np.array([self.nVar, self.dim, *self.nX], dtype=np.int32)] + [ - np.array(coord, dtype=np.float64) for coord in self.header["coords"] - ] - - 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, dim = np.fromfile(f, dtype=np.int32, count=2) - nX = np.fromfile(f, dtype=np.int32, count=dim) - coords = [np.fromfile(f, dtype=np.float64, count=n) for n in nX] - self.setHeader(nVar, coords) - - def reshape(self, fields: np.ndarray): - """Reshape the fields to a N-d array (inplace operation)""" - fields.shape = (self.nVar, *self.nX) - - # ------------------------------------------------------------------------- - # Class specifics - # ------------------------------------------------------------------------- - @property - def nX(self): - """Number of points in y direction""" - return [coord.size for coord in self.header["coords"]] - - @property - def dim(self): - """Number of grid dimensions""" - return len(self.nX) - - @property - def nDoF(self): - """Number of degrees of freedom for one variable""" - return np.prod(self.nX) - - # ------------------------------------------------------------------------- - # MPI-parallel implementation - # ------------------------------------------------------------------------- - comm: MPI.Intracomm = None - - @classmethod - def setupMPI(cls, comm: MPI.Intracomm, iLoc, nLoc): - """ - 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. - iLoc : list[int] - Starting index of the local sub-domain in the global `coordX`. - nLoc : list[int] - Number of points in the local sub-domain. - """ - cls.comm = comm - cls.iLoc = iLoc - cls.nLoc = nLoc - cls.mpiFile = None - - @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, - }[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: - try: - super().initialize() - except AssertionError as e: - if self.MPI_ON: - print(f"{type(e)}: {e}") - self.comm.Abort() - else: - raise e - - if self.MPI_ON: - self.comm.Barrier() # Important, should not be removed ! + dtype in DTYPES_AVAIL + ), f"{dtype=} not available. Supported on this machine: {list(DTYPES_AVAIL.keys())}" + 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): + """ + 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) + fieldsIO: FieldsIO = cls.STRUCTS[STRUCT](DTYPES[DTYPE], fileName) + fieldsIO.readHeader(f) + fieldsIO.initialized = True + return fieldsIO + + @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 + ), 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): + """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" + + if not self.ALLOW_OVERWRITE: + assert not os.path.isfile( + self.fileName + ), "file already exists, use FieldsIO.ALLOW_OVERWRITE = True to allow overwriting" + + with open(self.fileName, "w+b") as f: + self.hBase.tofile(f) + for array in self.hInfos: + array.tofile(f) self.initialized = True - def addField(self, time, field): - """ - Append one field solution at the end of the file with one given time, - possibly 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) - - 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.nLoc, - ), f"expected {(self.nVar, *self.nLoc)} 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, *iX) in itertools.product(range(self.nVar), *[range(nX) for nX in self.nLoc[:-1]]): - offset = offset0 + self.iPos(iVar, iX) * self.itemSize - self.MPI_WRITE_AT(offset, field[iVar, *iX]) - self.MPI_FILE_CLOSE() - - def iPos(self, iVar, iX): - iPos = iVar * self.nDoF - for axis in range(self.dim - 1): - iPos += (self.iLoc[axis] + iX[axis]) * np.prod(self.nX[axis + 1 :]) - iPos += self.iLoc[-1] - return iPos - - 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.formatIndex(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.nLoc), dtype=self.dtype) - - self.MPI_FILE_OPEN(mode="r") - for (iVar, *iX) in itertools.product(range(self.nVar), *[range(nX) for nX in self.nLoc[:-1]]): - offset = offset0 + self.iPos(iVar, iX) * self.itemSize - self.MPI_READ_AT(offset, field[iVar, *iX]) - self.MPI_FILE_CLOSE() - - return t, field + def setHeader(self, **params): + """(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 + 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): + """Size of one field value (in bytes)""" + return self.dtype().itemsize + + @property + def fSize(self): + """Full size of a field (in bytes)""" + return self.nItems * self.itemSize + + @property + def fileSize(self): + """Current size of the file (in bytes)""" + return os.path.getsize(self.fileName) + + 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}" + 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): + """Number of fields currently stored in the binary file""" + return int((self.fileSize - self.hSize) // (self.tSize + self.fSize)) + + def formatIndex(self, idx): + """Utility method to format a fields index to a positional integer (negative starts from last field index, like python lists)""" + 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): + """Vector of all times stored in the binary file""" + 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): + """Time stored at a given field index""" + idx = self.formatIndex(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): + """ + 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.formatIndex(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): + """Eventually reshape the field to correspond to the grid structure""" + pass + + @FieldsIO.register(sID=0) + class Scalar(FieldsIO): + """FieldsIO handler storing a given number of scalar""" + + # ------------------------------------------------------------------------- + # Overridden 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 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 Rectilinear(Scalar): + """FieldsIO handler storing a given number of scalar variables on a N-dimensional rectilinear grid""" + + @staticmethod + def setupCoords(*coords): + """Utility function to setup grids in multiple dimensions, given the keyword arguments""" + coords = [np.asarray(coord, dtype=np.float64) for coord in coords] + for axis, coord in enumerate(coords): + assert coord.ndim == 1, f"coord for {axis=} must be one dimensional" + return coords + + # ------------------------------------------------------------------------- + # Overridden methods + # ------------------------------------------------------------------------- + def setHeader(self, nVar, coords): + """ + Set the descriptive grid structure to be stored in the file header. + + Parameters + ---------- + nVar : int + Number of 1D variables stored. + coords : np.1darray or list[np.1darray] + The grid coordinates in each dimensions. + + Note + ---- + When used in MPI decomposition mode, all coordinate **must** be the global grid. + """ + if not isinstance(coords, (tuple, list)): + coords = [coords] + coords = self.setupCoords(*coords) + self.header = {"nVar": int(nVar), "coords": coords} + self.nItems = nVar * self.nDoF + + @property + def hInfos(self): + """Array representing the grid structure to be written in the binary file.""" + return [np.array([self.nVar, self.dim, *self.nX], dtype=np.int32)] + [ + np.array(coord, dtype=np.float64) for coord in self.header["coords"] + ] + + 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, dim = np.fromfile(f, dtype=np.int32, count=2) + nX = np.fromfile(f, dtype=np.int32, count=dim) + coords = [np.fromfile(f, dtype=np.float64, count=n) for n in nX] + self.setHeader(nVar, coords) + + def reshape(self, fields: np.ndarray): + """Reshape the fields to a N-d array (inplace operation)""" + fields.shape = (self.nVar, *self.nX) + + # ------------------------------------------------------------------------- + # Class specifics + # ------------------------------------------------------------------------- + @property + def nX(self): + """Number of points in y direction""" + return [coord.size for coord in self.header["coords"]] + + @property + def dim(self): + """Number of grid dimensions""" + return len(self.nX) + + @property + def nDoF(self): + """Number of degrees of freedom for one variable""" + return np.prod(self.nX) + + # ------------------------------------------------------------------------- + # MPI-parallel implementation + # ------------------------------------------------------------------------- + comm: MPI.Intracomm = None + + @classmethod + def setupMPI(cls, comm: MPI.Intracomm, iLoc, nLoc): + """ + 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. + iLoc : list[int] + Starting index of the local sub-domain in the global `coordX`. + nLoc : list[int] + Number of points in the local sub-domain. + """ + cls.comm = comm + cls.iLoc = iLoc + cls.nLoc = nLoc + cls.mpiFile = None + + @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, + }[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: + try: + super().initialize() + except AssertionError as e: + if self.MPI_ON: + print(f"{type(e)}: {e}") + self.comm.Abort() + else: + raise e + + if self.MPI_ON: + 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, + possibly 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) + + 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.nLoc, + ), f"expected {(self.nVar, *self.nLoc)} 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, *iX) in itertools.product(range(self.nVar), *[range(nX) for nX in self.nLoc[:-1]]): + offset = offset0 + self.iPos(iVar, iX) * self.itemSize + self.MPI_WRITE_AT(offset, field[iVar, *iX]) + self.MPI_FILE_CLOSE() + + def iPos(self, iVar, iX): + iPos = iVar * self.nDoF + for axis in range(self.dim - 1): + iPos += (self.iLoc[axis] + iX[axis]) * np.prod(self.nX[axis + 1 :]) + iPos += self.iLoc[-1] + return iPos + + 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.formatIndex(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.nLoc), dtype=self.dtype) + + self.MPI_FILE_OPEN(mode="r") + for (iVar, *iX) in itertools.product(range(self.nVar), *[range(nX) for nX in self.nLoc[:-1]]): + offset = offset0 + self.iPos(iVar, iX) * 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 index d2a2339f4..a59baf675 100644 --- a/pySDC/tests/test_helpers/test_fieldsIO.py +++ b/pySDC/tests/test_helpers/test_fieldsIO.py @@ -1,128 +1,80 @@ -import pytest -import itertools -import numpy as np +import sys + +if sys.version_info >= (3, 11): + import pytest + import itertools + import numpy as np + + from pySDC.helpers.fieldsIO import DTYPES, FieldsIO + + FieldsIO.ALLOW_OVERWRITE = True + + @pytest.mark.parametrize("dtypeIdx", DTYPES.keys()) + @pytest.mark.parametrize("dim", range(4)) + def testHeader(dim, dtypeIdx): + from pySDC.helpers.fieldsIO import FieldsIO, Scalar, Rectilinear + + fileName = "testHeader.pysdc" + dtype = DTYPES[dtypeIdx] + + coords = [np.linspace(0, 1, num=256, endpoint=False) for n in [256, 64, 32]] + + if dim == 0: + Class = Scalar + args = {"nVar": 20} + else: + Class = Rectilinear + args = {"nVar": 10, "coords": coords[:dim]} + + 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") -from pySDC.helpers.fieldsIO import DTYPES, FieldsIO - -FieldsIO.ALLOW_OVERWRITE = True - - -@pytest.mark.parametrize("dtypeIdx", DTYPES.keys()) -@pytest.mark.parametrize("dim", range(4)) -def testHeader(dim, dtypeIdx): - from pySDC.helpers.fieldsIO import FieldsIO, Scalar, Rectilinear - - fileName = "testHeader.pysdc" - dtype = DTYPES[dtypeIdx] - - coords = [np.linspace(0, 1, num=256, endpoint=False) for n in [256, 64, 32]] - - if dim == 0: - Class = Scalar - args = {"nVar": 20} - else: - Class = Rectilinear - args = {"nVar": 10, "coords": coords[:dim]} - - 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 testScalar(nVar, nSteps, dtypeIdx): - from pySDC.helpers.fieldsIO import FieldsIO, Scalar - - fileName = "testScalar.pysdc" - dtype = DTYPES[dtypeIdx] - - f1 = Scalar(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})" + 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" - 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("nVar", [1, 2, 5]) -@pytest.mark.parametrize("dim", [1, 2, 3]) -def testRectilinear(dim, nVar, nSteps, dtypeIdx): - from pySDC.helpers.fieldsIO import FieldsIO, Rectilinear, DTYPES - - fileName = f"testRectilinear{dim}D.pysdc" - dtype = DTYPES[dtypeIdx] + 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 nX in itertools.product(*[[5, 10, 16]] * dim): + 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}" - coords = [np.linspace(0, 1, num=n, endpoint=False) for n in nX] + @pytest.mark.parametrize("dtypeIdx", DTYPES.keys()) + @pytest.mark.parametrize("nSteps", [1, 2, 10, 100]) + @pytest.mark.parametrize("nVar", [1, 2, 5]) + def testScalar(nVar, nSteps, dtypeIdx): + from pySDC.helpers.fieldsIO import FieldsIO, Scalar - f1 = Rectilinear(dtype, fileName) - f1.setHeader(nVar=nVar, coords=coords) + fileName = "testScalar.pysdc" + dtype = DTYPES[dtypeIdx] - assert f1.dim == dim, f"{f1} has incorrect dimension" - assert f1.nX == list(nX), f"{f1} has incorrect nX" - assert f1.nDoF == np.prod(nX), f"{f1} has incorrect nDOF" - assert f1.nItems == nVar * np.prod(nX), f"{f1} do not have nItems == nVar*nX**dim" + f1 = Scalar(dtype, fileName) + f1.setHeader(nVar=nVar) + assert f1.nItems == nVar, f"{f1} do not have nItems == nVar" f1.initialize() - u0 = np.random.rand(nVar, *nX).astype(f1.dtype) + + u0 = np.random.rand(nVar).astype(f1.dtype) times = np.arange(nSteps) / nSteps for t in times: @@ -147,98 +99,99 @@ def testRectilinear(dim, nVar, nSteps, dtypeIdx): 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("nVar", [1, 2, 5]) + @pytest.mark.parametrize("dim", [1, 2, 3]) + def testRectilinear(dim, nVar, nSteps, dtypeIdx): + from pySDC.helpers.fieldsIO import FieldsIO, Rectilinear, DTYPES -def initGrid(nVar, gridSizes): - dim = len(gridSizes) - coords = [np.linspace(0, 1, num=n, endpoint=False) for n in gridSizes] - s = [None] * dim - u0 = np.array(np.arange(nVar) + 1)[:, *s] - for x in np.meshgrid(*coords, indexing="ij"): - u0 = u0 * x - return coords, u0 + fileName = f"testRectilinear{dim}D.pysdc" + dtype = DTYPES[dtypeIdx] + for nX in itertools.product(*[[5, 10, 16]] * dim): -def writeFields_MPI(fileName, dtypeIdx, algo, nSteps, nVar, nX): - coords, u0 = initGrid(nVar, nX) + coords = [np.linspace(0, 1, num=n, endpoint=False) for n in nX] - from mpi4py import MPI - from pySDC.helpers.blocks import BlockDecomposition - from pySDC.helpers.fieldsIO import Rectilinear + f1 = Rectilinear(dtype, fileName) + f1.setHeader(nVar=nVar, coords=coords) - comm = MPI.COMM_WORLD - MPI_SIZE = comm.Get_size() - MPI_RANK = comm.Get_rank() + assert f1.dim == dim, f"{f1} has incorrect dimension" + assert f1.nX == list(nX), f"{f1} has incorrect nX" + assert f1.nDoF == np.prod(nX), f"{f1} has incorrect nDOF" + assert f1.nItems == nVar * np.prod(nX), f"{f1} do not have nItems == nVar*nX**dim" - blocks = BlockDecomposition(MPI_SIZE, nX, algo, MPI_RANK) + f1.initialize() + u0 = np.random.rand(nVar, *nX).astype(f1.dtype) + times = np.arange(nSteps) / nSteps - iLoc, nLoc = blocks.localBounds - Rectilinear.setupMPI(comm, iLoc, nLoc) - s = [slice(i, i + n) for i, n in zip(iLoc, nLoc)] - u0 = u0[:, *s] - print(MPI_RANK, u0.shape) + for t in times: + ut = (u0 * t).astype(f1.dtype) + f1.addField(t, ut) - f1 = Rectilinear(DTYPES[dtypeIdx], fileName) - f1.setHeader(nVar=nVar, coords=coords) + 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" - u0 = np.asarray(u0, dtype=f1.dtype) - f1.initialize() + f2 = FieldsIO.fromFile(fileName) - times = np.arange(nSteps) / nSteps - for t in times: - ut = (u0 * t).astype(f1.dtype) - f1.addField(t, ut) + 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})" - return u0 + 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, gridSizes): + dim = len(gridSizes) + coords = [np.linspace(0, 1, num=n, endpoint=False) for n in gridSizes] + s = [None] * dim + u0 = np.array(np.arange(nVar) + 1)[:, *s] + for x in np.meshgrid(*coords, indexing="ij"): + u0 = u0 * x + return coords, 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"fields[{idx}] in {f2} has incorrect time ({t2} instead of {t})" - 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" - + def writeFields_MPI(fileName, dtypeIdx, algo, nSteps, nVar, nX): + coords, u0 = initGrid(nVar, nX) -@pytest.mark.mpi4py -@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", [2, 4]) -@pytest.mark.parametrize("dim", [2, 3]) -def testRectilinear_MPI(dim, nProcs, dtypeIdx, algo, nSteps, nVar): + from mpi4py import MPI + from pySDC.helpers.blocks import BlockDecomposition + from pySDC.helpers.fieldsIO import Rectilinear - import subprocess + comm = MPI.COMM_WORLD + MPI_SIZE = comm.Get_size() + MPI_RANK = comm.Get_rank() - fileName = f"testRectilinear{dim}D_MPI.pysdc" + blocks = BlockDecomposition(MPI_SIZE, nX, algo, MPI_RANK) - for nX in itertools.product(*[[61, 16]] * dim): + iLoc, nLoc = blocks.localBounds + Rectilinear.setupMPI(comm, iLoc, nLoc) + s = [slice(i, i + n) for i, n in zip(iLoc, nLoc)] + u0 = u0[:, *s] + print(MPI_RANK, u0.shape) - cmd = f"mpirun -np {nProcs} python {__file__} --fileName {fileName}" - cmd += f" --dtypeIdx {dtypeIdx} --algo {algo} --nSteps {nSteps} --nVar {nVar} --nX {' '.join([str(n) for n in nX])}" + f1 = Rectilinear(DTYPES[dtypeIdx], fileName) + f1.setHeader(nVar=nVar, coords=coords) - p = subprocess.Popen(cmd.split(), cwd=".") - p.wait() - assert p.returncode == 0, f"MPI write with {nProcs} proc(s) did not return code 0, but {p.returncode}" + u0 = np.asarray(u0, dtype=f1.dtype) + f1.initialize() - from pySDC.helpers.fieldsIO import FieldsIO, Rectilinear + times = np.arange(nSteps) / nSteps + for t in times: + ut = (u0 * t).astype(f1.dtype) + f1.addField(t, ut) - f2: Rectilinear = FieldsIO.fromFile(fileName) + return u0 - assert type(f2) == Rectilinear, f"incorrect type in MPI written fields {f2}" - assert f2.nFields == nSteps, f"incorrect nFields in MPI written fields {f2} ({f2.nFields} instead of {nSteps})" - assert f2.nVar == nVar, f"incorrect nVar in MPI written fields {f2}" - assert f2.nX == list(nX), f"incorrect nX in MPI written fields {f2}" + def compareFields_MPI(fileName, u0, nSteps): + from pySDC.helpers.fieldsIO import FieldsIO - coords, u0 = initGrid(nVar, nX) - for i, (cFile, cRef) in enumerate(zip(f2.header['coords'], coords)): - assert np.allclose(cFile, cRef), f"incorrect coords[{i}] in MPI written fields {f2}" + f2 = FieldsIO.fromFile(fileName) times = np.arange(nSteps) / nSteps for idx, t in enumerate(times): @@ -248,18 +201,62 @@ def testRectilinear_MPI(dim, nProcs, dtypeIdx, algo, nSteps, nVar): 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('--dtypeIdx', type=int, help="dtype index", choices=DTYPES.keys()) - parser.add_argument('--algo', type=str, help="algorithm used for block decomposition") - parser.add_argument('--nSteps', type=int, help="number of time-steps") - parser.add_argument('--nVar', type=int, help="number of field variables") - parser.add_argument('--nX', type=int, nargs='+', help="number of grid points in each dimensions") - args = parser.parse_args() - - u0 = writeFields_MPI(**args.__dict__) - compareFields_MPI(args.fileName, u0, args.nSteps) + @pytest.mark.mpi4py + @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", [2, 4]) + @pytest.mark.parametrize("dim", [2, 3]) + def testRectilinear_MPI(dim, nProcs, dtypeIdx, algo, nSteps, nVar): + + import subprocess + + fileName = f"testRectilinear{dim}D_MPI.pysdc" + + for nX in itertools.product(*[[61, 16]] * dim): + + cmd = f"mpirun -np {nProcs} python {__file__} --fileName {fileName}" + cmd += f" --dtypeIdx {dtypeIdx} --algo {algo} --nSteps {nSteps} --nVar {nVar} --nX {' '.join([str(n) for n in nX])}" + + p = subprocess.Popen(cmd.split(), cwd=".") + p.wait() + assert p.returncode == 0, f"MPI write with {nProcs} proc(s) did not return code 0, but {p.returncode}" + + from pySDC.helpers.fieldsIO import FieldsIO, Rectilinear + + f2: Rectilinear = FieldsIO.fromFile(fileName) + + assert type(f2) == Rectilinear, f"incorrect type in MPI written fields {f2}" + assert ( + f2.nFields == nSteps + ), f"incorrect nFields in MPI written fields {f2} ({f2.nFields} instead of {nSteps})" + assert f2.nVar == nVar, f"incorrect nVar in MPI written fields {f2}" + assert f2.nX == list(nX), f"incorrect nX in MPI written fields {f2}" + + coords, u0 = initGrid(nVar, nX) + for i, (cFile, cRef) in enumerate(zip(f2.header['coords'], coords)): + assert np.allclose(cFile, cRef), f"incorrect coords[{i}] 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"fields[{idx}] in {f2} has incorrect time ({t2} instead of {t})" + 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('--dtypeIdx', type=int, help="dtype index", choices=DTYPES.keys()) + parser.add_argument('--algo', type=str, help="algorithm used for block decomposition") + parser.add_argument('--nSteps', type=int, help="number of time-steps") + parser.add_argument('--nVar', type=int, help="number of field variables") + parser.add_argument('--nX', type=int, nargs='+', help="number of grid points in each dimensions") + args = parser.parse_args() + + u0 = writeFields_MPI(**args.__dict__) + compareFields_MPI(args.fileName, u0, args.nSteps) From 830d2f12fdb9a45b164cc73de59c7b1d24e02dde Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 01:20:27 +0100 Subject: [PATCH 04/16] TL: another try --- pySDC/helpers/fieldsIO.py | 1173 ++++++++++----------- pySDC/tests/test_helpers/test_fieldsIO.py | 396 +++---- 2 files changed, 788 insertions(+), 781 deletions(-) diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py index 57a379907..dfe3c72eb 100644 --- a/pySDC/helpers/fieldsIO.py +++ b/pySDC/helpers/fieldsIO.py @@ -45,597 +45,596 @@ Also : this feature is only available with Python 3.11 or higher ... """ -import sys +import os +import numpy as np +from typing import Type, TypeVar +import logging +import itertools -if sys.version_info >= (3, 11): - import os - import numpy as np - from typing import Type, TypeVar - import logging - import itertools - - T = TypeVar("T") +T = TypeVar("T") +try: 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 + import dolfin as df # noqa: F841 (for some reason, dolfin always needs to be imported before mpi4py) except ImportError: - - class MPI: - COMM_WORLD = None - Intracomm = T - - # Supported data types - DTYPES = { - 0: np.float64, # double precision - 1: np.complex128, - } - try: - DTYPES.update( - { - 2: np.float128, # quadruple precision - 3: np.complex256, - } - ) - except AttributeError: - logging.getLogger('FieldsIO').debug('Warning: Quadruple precision not available on this machine') - try: - DTYPES.update( - { - 4: np.float32, # single precision - 5: np.complex64, - } - ) - except AttributeError: - logging.getLogger('FieldsIO').debug('Warning: Single precision not available on this machine') - - DTYPES_AVAIL = {val: key for key, val in DTYPES.items()} - - # Header dtype - H_DTYPE = np.int8 - T_DTYPE = np.float64 - - class FieldsIO: - """Abstract IO file handler""" - - STRUCTS = {} # supported structures, modified dynamically - sID = None # structure ID of the FieldsIO class, modified dynamically - - tSize = T_DTYPE().itemsize - - ALLOW_OVERWRITE = False - - def __init__(self, dtype, fileName): - """ - Parameters - ---------- - dtype : np.dtype - The data type of the fields values. - fileName : str - File. - """ + pass + from mpi4py import MPI +except ImportError: + + class MPI: + COMM_WORLD = None + Intracomm = T + + +# Supported data types +DTYPES = { + 0: np.float64, # double precision + 1: np.complex128, +} +try: + DTYPES.update( + { + 2: np.float128, # quadruple precision + 3: np.complex256, + } + ) +except AttributeError: + logging.getLogger('FieldsIO').debug('Warning: Quadruple precision not available on this machine') +try: + DTYPES.update( + { + 4: np.float32, # single precision + 5: np.complex64, + } + ) +except AttributeError: + logging.getLogger('FieldsIO').debug('Warning: Single precision not available on this machine') + +DTYPES_AVAIL = {val: key for key, val in DTYPES.items()} + +# Header dtype +H_DTYPE = np.int8 +T_DTYPE = np.float64 + + +class FieldsIO: + """Abstract IO file handler""" + + STRUCTS = {} # supported structures, modified dynamically + sID = None # structure ID of the FieldsIO class, modified dynamically + + tSize = T_DTYPE().itemsize + + ALLOW_OVERWRITE = False + + 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. Supported on this machine: {list(DTYPES_AVAIL.keys())}" + 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): + """ + 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) + fieldsIO: FieldsIO = cls.STRUCTS[STRUCT](DTYPES[DTYPE], fileName) + fieldsIO.readHeader(f) + fieldsIO.initialized = True + return fieldsIO + + @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 ( - dtype in DTYPES_AVAIL - ), f"{dtype=} not available. Supported on this machine: {list(DTYPES_AVAIL.keys())}" - 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): - """ - 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) - fieldsIO: FieldsIO = cls.STRUCTS[STRUCT](DTYPES[DTYPE], fileName) - fieldsIO.readHeader(f) - fieldsIO.initialized = True - return fieldsIO - - @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 - ), 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): - """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" - - if not self.ALLOW_OVERWRITE: - assert not os.path.isfile( - self.fileName - ), "file already exists, use FieldsIO.ALLOW_OVERWRITE = True to allow overwriting" - - with open(self.fileName, "w+b") as f: - self.hBase.tofile(f) - for array in self.hInfos: - array.tofile(f) + 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): + """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" + + if not self.ALLOW_OVERWRITE: + assert not os.path.isfile( + self.fileName + ), "file already exists, use FieldsIO.ALLOW_OVERWRITE = True to allow overwriting" + + 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): + """(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 + 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): + """Size of one field value (in bytes)""" + return self.dtype().itemsize + + @property + def fSize(self): + """Full size of a field (in bytes)""" + return self.nItems * self.itemSize + + @property + def fileSize(self): + """Current size of the file (in bytes)""" + return os.path.getsize(self.fileName) + + 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}" + 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): + """Number of fields currently stored in the binary file""" + return int((self.fileSize - self.hSize) // (self.tSize + self.fSize)) + + def formatIndex(self, idx): + """Utility method to format a fields index to a positional integer (negative starts from last field index, like python lists)""" + 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): + """Vector of all times stored in the binary file""" + 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): + """Time stored at a given field index""" + idx = self.formatIndex(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): + """ + 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.formatIndex(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): + """Eventually reshape the field to correspond to the grid structure""" + pass + + +@FieldsIO.register(sID=0) +class Scalar(FieldsIO): + """FieldsIO handler storing a given number of scalar""" + + # ------------------------------------------------------------------------- + # Overridden 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 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 Rectilinear(Scalar): + """FieldsIO handler storing a given number of scalar variables on a N-dimensional rectilinear grid""" + + @staticmethod + def setupCoords(*coords): + """Utility function to setup grids in multiple dimensions, given the keyword arguments""" + coords = [np.asarray(coord, dtype=np.float64) for coord in coords] + for axis, coord in enumerate(coords): + assert coord.ndim == 1, f"coord for {axis=} must be one dimensional" + return coords + + # ------------------------------------------------------------------------- + # Overridden methods + # ------------------------------------------------------------------------- + def setHeader(self, nVar, coords): + """ + Set the descriptive grid structure to be stored in the file header. + + Parameters + ---------- + nVar : int + Number of 1D variables stored. + coords : np.1darray or list[np.1darray] + The grid coordinates in each dimensions. + + Note + ---- + When used in MPI decomposition mode, all coordinate **must** be the global grid. + """ + if not isinstance(coords, (tuple, list)): + coords = [coords] + coords = self.setupCoords(*coords) + self.header = {"nVar": int(nVar), "coords": coords} + self.nItems = nVar * self.nDoF + + @property + def hInfos(self): + """Array representing the grid structure to be written in the binary file.""" + return [np.array([self.nVar, self.dim, *self.nX], dtype=np.int32)] + [ + np.array(coord, dtype=np.float64) for coord in self.header["coords"] + ] + + 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, dim = np.fromfile(f, dtype=np.int32, count=2) + nX = np.fromfile(f, dtype=np.int32, count=dim) + coords = [np.fromfile(f, dtype=np.float64, count=n) for n in nX] + self.setHeader(nVar, coords) + + def reshape(self, fields: np.ndarray): + """Reshape the fields to a N-d array (inplace operation)""" + fields.shape = (self.nVar, *self.nX) + + # ------------------------------------------------------------------------- + # Class specifics + # ------------------------------------------------------------------------- + @property + def nX(self): + """Number of points in y direction""" + return [coord.size for coord in self.header["coords"]] + + @property + def dim(self): + """Number of grid dimensions""" + return len(self.nX) + + @property + def nDoF(self): + """Number of degrees of freedom for one variable""" + return np.prod(self.nX) + + # ------------------------------------------------------------------------- + # MPI-parallel implementation + # ------------------------------------------------------------------------- + comm: MPI.Intracomm = None + + @classmethod + def setupMPI(cls, comm: MPI.Intracomm, iLoc, nLoc): + """ + 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. + iLoc : list[int] + Starting index of the local sub-domain in the global `coordX`. + nLoc : list[int] + Number of points in the local sub-domain. + """ + cls.comm = comm + cls.iLoc = iLoc + cls.nLoc = nLoc + cls.mpiFile = None + + @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, + }[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: + try: + super().initialize() + except AssertionError as e: + if self.MPI_ON: + print(f"{type(e)}: {e}") + self.comm.Abort() + else: + raise e + + if self.MPI_ON: + self.comm.Barrier() # Important, should not be removed ! self.initialized = True - def setHeader(self, **params): - """(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 - 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): - """Size of one field value (in bytes)""" - return self.dtype().itemsize - - @property - def fSize(self): - """Full size of a field (in bytes)""" - return self.nItems * self.itemSize - - @property - def fileSize(self): - """Current size of the file (in bytes)""" - return os.path.getsize(self.fileName) - - 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}" - 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): - """Number of fields currently stored in the binary file""" - return int((self.fileSize - self.hSize) // (self.tSize + self.fSize)) - - def formatIndex(self, idx): - """Utility method to format a fields index to a positional integer (negative starts from last field index, like python lists)""" - 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): - """Vector of all times stored in the binary file""" - 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): - """Time stored at a given field index""" - idx = self.formatIndex(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): - """ - 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.formatIndex(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): - """Eventually reshape the field to correspond to the grid structure""" - pass - - @FieldsIO.register(sID=0) - class Scalar(FieldsIO): - """FieldsIO handler storing a given number of scalar""" - - # ------------------------------------------------------------------------- - # Overridden 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 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 Rectilinear(Scalar): - """FieldsIO handler storing a given number of scalar variables on a N-dimensional rectilinear grid""" - - @staticmethod - def setupCoords(*coords): - """Utility function to setup grids in multiple dimensions, given the keyword arguments""" - coords = [np.asarray(coord, dtype=np.float64) for coord in coords] - for axis, coord in enumerate(coords): - assert coord.ndim == 1, f"coord for {axis=} must be one dimensional" - return coords - - # ------------------------------------------------------------------------- - # Overridden methods - # ------------------------------------------------------------------------- - def setHeader(self, nVar, coords): - """ - Set the descriptive grid structure to be stored in the file header. - - Parameters - ---------- - nVar : int - Number of 1D variables stored. - coords : np.1darray or list[np.1darray] - The grid coordinates in each dimensions. - - Note - ---- - When used in MPI decomposition mode, all coordinate **must** be the global grid. - """ - if not isinstance(coords, (tuple, list)): - coords = [coords] - coords = self.setupCoords(*coords) - self.header = {"nVar": int(nVar), "coords": coords} - self.nItems = nVar * self.nDoF - - @property - def hInfos(self): - """Array representing the grid structure to be written in the binary file.""" - return [np.array([self.nVar, self.dim, *self.nX], dtype=np.int32)] + [ - np.array(coord, dtype=np.float64) for coord in self.header["coords"] - ] - - 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, dim = np.fromfile(f, dtype=np.int32, count=2) - nX = np.fromfile(f, dtype=np.int32, count=dim) - coords = [np.fromfile(f, dtype=np.float64, count=n) for n in nX] - self.setHeader(nVar, coords) - - def reshape(self, fields: np.ndarray): - """Reshape the fields to a N-d array (inplace operation)""" - fields.shape = (self.nVar, *self.nX) - - # ------------------------------------------------------------------------- - # Class specifics - # ------------------------------------------------------------------------- - @property - def nX(self): - """Number of points in y direction""" - return [coord.size for coord in self.header["coords"]] - - @property - def dim(self): - """Number of grid dimensions""" - return len(self.nX) - - @property - def nDoF(self): - """Number of degrees of freedom for one variable""" - return np.prod(self.nX) - - # ------------------------------------------------------------------------- - # MPI-parallel implementation - # ------------------------------------------------------------------------- - comm: MPI.Intracomm = None - - @classmethod - def setupMPI(cls, comm: MPI.Intracomm, iLoc, nLoc): - """ - 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. - iLoc : list[int] - Starting index of the local sub-domain in the global `coordX`. - nLoc : list[int] - Number of points in the local sub-domain. - """ - cls.comm = comm - cls.iLoc = iLoc - cls.nLoc = nLoc - cls.mpiFile = None - - @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, - }[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: - try: - super().initialize() - except AssertionError as e: - if self.MPI_ON: - print(f"{type(e)}: {e}") - self.comm.Abort() - else: - raise e - - if self.MPI_ON: - 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, - possibly 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) - - 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.nLoc, - ), f"expected {(self.nVar, *self.nLoc)} 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, *iX) in itertools.product(range(self.nVar), *[range(nX) for nX in self.nLoc[:-1]]): - offset = offset0 + self.iPos(iVar, iX) * self.itemSize - self.MPI_WRITE_AT(offset, field[iVar, *iX]) - self.MPI_FILE_CLOSE() - - def iPos(self, iVar, iX): - iPos = iVar * self.nDoF - for axis in range(self.dim - 1): - iPos += (self.iLoc[axis] + iX[axis]) * np.prod(self.nX[axis + 1 :]) - iPos += self.iLoc[-1] - return iPos - - 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.formatIndex(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.nLoc), dtype=self.dtype) - - self.MPI_FILE_OPEN(mode="r") - for (iVar, *iX) in itertools.product(range(self.nVar), *[range(nX) for nX in self.nLoc[:-1]]): - offset = offset0 + self.iPos(iVar, iX) * self.itemSize - self.MPI_READ_AT(offset, field[iVar, *iX]) - self.MPI_FILE_CLOSE() - - return t, field + def addField(self, time, field): + """ + Append one field solution at the end of the file with one given time, + possibly 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) + + 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.nLoc, + ), f"expected {(self.nVar, *self.nLoc)} 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, *iX) in itertools.product(range(self.nVar), *[range(nX) for nX in self.nLoc[:-1]]): + offset = offset0 + self.iPos(iVar, iX) * self.itemSize + self.MPI_WRITE_AT(offset, field[iVar, *iX]) + self.MPI_FILE_CLOSE() + + def iPos(self, iVar, iX): + iPos = iVar * self.nDoF + for axis in range(self.dim - 1): + iPos += (self.iLoc[axis] + iX[axis]) * np.prod(self.nX[axis + 1 :]) + iPos += self.iLoc[-1] + return iPos + + 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.formatIndex(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.nLoc), dtype=self.dtype) + + self.MPI_FILE_OPEN(mode="r") + for (iVar, *iX) in itertools.product(range(self.nVar), *[range(nX) for nX in self.nLoc[:-1]]): + offset = offset0 + self.iPos(iVar, iX) * 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 index a59baf675..0cf826d41 100644 --- a/pySDC/tests/test_helpers/test_fieldsIO.py +++ b/pySDC/tests/test_helpers/test_fieldsIO.py @@ -1,80 +1,133 @@ import sys +import pytest if sys.version_info >= (3, 11): - import pytest - import itertools - import numpy as np - - from pySDC.helpers.fieldsIO import DTYPES, FieldsIO - - FieldsIO.ALLOW_OVERWRITE = True - - @pytest.mark.parametrize("dtypeIdx", DTYPES.keys()) - @pytest.mark.parametrize("dim", range(4)) - def testHeader(dim, dtypeIdx): - from pySDC.helpers.fieldsIO import FieldsIO, Scalar, Rectilinear - - fileName = "testHeader.pysdc" - dtype = DTYPES[dtypeIdx] - - coords = [np.linspace(0, 1, num=256, endpoint=False) for n in [256, 64, 32]] - - if dim == 0: - Class = Scalar - args = {"nVar": 20} - else: - Class = Rectilinear - args = {"nVar": 10, "coords": coords[:dim]} - - 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") + pytest.skip("skipping fieldsIO tests on python lower than 3.11", allow_module_level=True) - 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" +import itertools +import numpy as np - 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})" +from pySDC.helpers.fieldsIO import DTYPES, FieldsIO + +FieldsIO.ALLOW_OVERWRITE = True - 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 testScalar(nVar, nSteps, dtypeIdx): - from pySDC.helpers.fieldsIO import FieldsIO, Scalar +@pytest.mark.parametrize("dtypeIdx", DTYPES.keys()) +@pytest.mark.parametrize("dim", range(4)) +def testHeader(dim, dtypeIdx): + from pySDC.helpers.fieldsIO import FieldsIO, Scalar, Rectilinear - fileName = "testScalar.pysdc" - dtype = DTYPES[dtypeIdx] + fileName = "testHeader.pysdc" + dtype = DTYPES[dtypeIdx] - f1 = Scalar(dtype, fileName) - f1.setHeader(nVar=nVar) + coords = [np.linspace(0, 1, num=256, endpoint=False) for n in [256, 64, 32]] - assert f1.nItems == nVar, f"{f1} do not have nItems == nVar" + if dim == 0: + Class = Scalar + args = {"nVar": 20} + else: + Class = Rectilinear + args = {"nVar": 10, "coords": coords[:dim]} + + 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 testScalar(nVar, nSteps, dtypeIdx): + from pySDC.helpers.fieldsIO import FieldsIO, Scalar + + fileName = "testScalar.pysdc" + dtype = DTYPES[dtypeIdx] + + f1 = Scalar(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("nVar", [1, 2, 5]) +@pytest.mark.parametrize("dim", [1, 2, 3]) +def testRectilinear(dim, nVar, nSteps, dtypeIdx): + from pySDC.helpers.fieldsIO import FieldsIO, Rectilinear, DTYPES + + fileName = f"testRectilinear{dim}D.pysdc" + dtype = DTYPES[dtypeIdx] - u0 = np.random.rand(nVar).astype(f1.dtype) + for nX in itertools.product(*[[5, 10, 16]] * dim): + + coords = [np.linspace(0, 1, num=n, endpoint=False) for n in nX] + + f1 = Rectilinear(dtype, fileName) + f1.setHeader(nVar=nVar, coords=coords) + + assert f1.dim == dim, f"{f1} has incorrect dimension" + assert f1.nX == list(nX), f"{f1} has incorrect nX" + assert f1.nDoF == np.prod(nX), f"{f1} has incorrect nDOF" + assert f1.nItems == nVar * np.prod(nX), f"{f1} do not have nItems == nVar*nX**dim" + + f1.initialize() + u0 = np.random.rand(nVar, *nX).astype(f1.dtype) times = np.arange(nSteps) / nSteps for t in times: @@ -99,99 +152,98 @@ def testScalar(nVar, nSteps, dtypeIdx): 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("nVar", [1, 2, 5]) - @pytest.mark.parametrize("dim", [1, 2, 3]) - def testRectilinear(dim, nVar, nSteps, dtypeIdx): - from pySDC.helpers.fieldsIO import FieldsIO, Rectilinear, DTYPES - fileName = f"testRectilinear{dim}D.pysdc" - dtype = DTYPES[dtypeIdx] +def initGrid(nVar, gridSizes): + dim = len(gridSizes) + coords = [np.linspace(0, 1, num=n, endpoint=False) for n in gridSizes] + s = [None] * dim + u0 = np.array(np.arange(nVar) + 1)[:, *s] + for x in np.meshgrid(*coords, indexing="ij"): + u0 = u0 * x + return coords, u0 - for nX in itertools.product(*[[5, 10, 16]] * dim): - coords = [np.linspace(0, 1, num=n, endpoint=False) for n in nX] +def writeFields_MPI(fileName, dtypeIdx, algo, nSteps, nVar, nX): + coords, u0 = initGrid(nVar, nX) - f1 = Rectilinear(dtype, fileName) - f1.setHeader(nVar=nVar, coords=coords) + from mpi4py import MPI + from pySDC.helpers.blocks import BlockDecomposition + from pySDC.helpers.fieldsIO import Rectilinear - assert f1.dim == dim, f"{f1} has incorrect dimension" - assert f1.nX == list(nX), f"{f1} has incorrect nX" - assert f1.nDoF == np.prod(nX), f"{f1} has incorrect nDOF" - assert f1.nItems == nVar * np.prod(nX), f"{f1} do not have nItems == nVar*nX**dim" + comm = MPI.COMM_WORLD + MPI_SIZE = comm.Get_size() + MPI_RANK = comm.Get_rank() - f1.initialize() - u0 = np.random.rand(nVar, *nX).astype(f1.dtype) - times = np.arange(nSteps) / nSteps + blocks = BlockDecomposition(MPI_SIZE, nX, algo, MPI_RANK) - for t in times: - ut = (u0 * t).astype(f1.dtype) - f1.addField(t, ut) + iLoc, nLoc = blocks.localBounds + Rectilinear.setupMPI(comm, iLoc, nLoc) + s = [slice(i, i + n) for i, n in zip(iLoc, nLoc)] + u0 = u0[:, *s] + print(MPI_RANK, u0.shape) - 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" + f1 = Rectilinear(DTYPES[dtypeIdx], fileName) + f1.setHeader(nVar=nVar, coords=coords) - f2 = FieldsIO.fromFile(fileName) + u0 = np.asarray(u0, dtype=f1.dtype) + f1.initialize() - 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})" + times = np.arange(nSteps) / nSteps + for t in times: + ut = (u0 * t).astype(f1.dtype) + f1.addField(t, ut) - 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" + return u0 - def initGrid(nVar, gridSizes): - dim = len(gridSizes) - coords = [np.linspace(0, 1, num=n, endpoint=False) for n in gridSizes] - s = [None] * dim - u0 = np.array(np.arange(nVar) + 1)[:, *s] - for x in np.meshgrid(*coords, indexing="ij"): - u0 = u0 * x - return coords, u0 - def writeFields_MPI(fileName, dtypeIdx, algo, nSteps, nVar, nX): - coords, u0 = initGrid(nVar, nX) +def compareFields_MPI(fileName, u0, nSteps): + from pySDC.helpers.fieldsIO import FieldsIO - from mpi4py import MPI - from pySDC.helpers.blocks import BlockDecomposition - from pySDC.helpers.fieldsIO import Rectilinear + f2 = FieldsIO.fromFile(fileName) - comm = MPI.COMM_WORLD - MPI_SIZE = comm.Get_size() - MPI_RANK = comm.Get_rank() + times = np.arange(nSteps) / nSteps + for idx, t in enumerate(times): + u1 = u0 * t + t2, u2 = f2.readField(idx) + assert t2 == t, f"fields[{idx}] in {f2} has incorrect time ({t2} instead of {t})" + 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" - blocks = BlockDecomposition(MPI_SIZE, nX, algo, MPI_RANK) - iLoc, nLoc = blocks.localBounds - Rectilinear.setupMPI(comm, iLoc, nLoc) - s = [slice(i, i + n) for i, n in zip(iLoc, nLoc)] - u0 = u0[:, *s] - print(MPI_RANK, u0.shape) +@pytest.mark.mpi4py +@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", [2, 4]) +@pytest.mark.parametrize("dim", [2, 3]) +def testRectilinear_MPI(dim, nProcs, dtypeIdx, algo, nSteps, nVar): - f1 = Rectilinear(DTYPES[dtypeIdx], fileName) - f1.setHeader(nVar=nVar, coords=coords) + import subprocess - u0 = np.asarray(u0, dtype=f1.dtype) - f1.initialize() + fileName = f"testRectilinear{dim}D_MPI.pysdc" - times = np.arange(nSteps) / nSteps - for t in times: - ut = (u0 * t).astype(f1.dtype) - f1.addField(t, ut) + for nX in itertools.product(*[[61, 16]] * dim): - return u0 + cmd = f"mpirun -np {nProcs} python {__file__} --fileName {fileName}" + cmd += f" --dtypeIdx {dtypeIdx} --algo {algo} --nSteps {nSteps} --nVar {nVar} --nX {' '.join([str(n) for n in nX])}" - def compareFields_MPI(fileName, u0, nSteps): - from pySDC.helpers.fieldsIO import FieldsIO + p = subprocess.Popen(cmd.split(), cwd=".") + p.wait() + assert p.returncode == 0, f"MPI write with {nProcs} proc(s) did not return code 0, but {p.returncode}" - f2 = FieldsIO.fromFile(fileName) + from pySDC.helpers.fieldsIO import FieldsIO, Rectilinear + + f2: Rectilinear = FieldsIO.fromFile(fileName) + + assert type(f2) == Rectilinear, f"incorrect type in MPI written fields {f2}" + assert f2.nFields == nSteps, f"incorrect nFields in MPI written fields {f2} ({f2.nFields} instead of {nSteps})" + assert f2.nVar == nVar, f"incorrect nVar in MPI written fields {f2}" + assert f2.nX == list(nX), f"incorrect nX in MPI written fields {f2}" + + coords, u0 = initGrid(nVar, nX) + for i, (cFile, cRef) in enumerate(zip(f2.header['coords'], coords)): + assert np.allclose(cFile, cRef), f"incorrect coords[{i}] in MPI written fields {f2}" times = np.arange(nSteps) / nSteps for idx, t in enumerate(times): @@ -201,62 +253,18 @@ def compareFields_MPI(fileName, u0, nSteps): 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("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", [2, 4]) - @pytest.mark.parametrize("dim", [2, 3]) - def testRectilinear_MPI(dim, nProcs, dtypeIdx, algo, nSteps, nVar): - - import subprocess - - fileName = f"testRectilinear{dim}D_MPI.pysdc" - - for nX in itertools.product(*[[61, 16]] * dim): - - cmd = f"mpirun -np {nProcs} python {__file__} --fileName {fileName}" - cmd += f" --dtypeIdx {dtypeIdx} --algo {algo} --nSteps {nSteps} --nVar {nVar} --nX {' '.join([str(n) for n in nX])}" - - p = subprocess.Popen(cmd.split(), cwd=".") - p.wait() - assert p.returncode == 0, f"MPI write with {nProcs} proc(s) did not return code 0, but {p.returncode}" - - from pySDC.helpers.fieldsIO import FieldsIO, Rectilinear - - f2: Rectilinear = FieldsIO.fromFile(fileName) - - assert type(f2) == Rectilinear, f"incorrect type in MPI written fields {f2}" - assert ( - f2.nFields == nSteps - ), f"incorrect nFields in MPI written fields {f2} ({f2.nFields} instead of {nSteps})" - assert f2.nVar == nVar, f"incorrect nVar in MPI written fields {f2}" - assert f2.nX == list(nX), f"incorrect nX in MPI written fields {f2}" - - coords, u0 = initGrid(nVar, nX) - for i, (cFile, cRef) in enumerate(zip(f2.header['coords'], coords)): - assert np.allclose(cFile, cRef), f"incorrect coords[{i}] 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"fields[{idx}] in {f2} has incorrect time ({t2} instead of {t})" - 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('--dtypeIdx', type=int, help="dtype index", choices=DTYPES.keys()) - parser.add_argument('--algo', type=str, help="algorithm used for block decomposition") - parser.add_argument('--nSteps', type=int, help="number of time-steps") - parser.add_argument('--nVar', type=int, help="number of field variables") - parser.add_argument('--nX', type=int, nargs='+', help="number of grid points in each dimensions") - args = parser.parse_args() - - u0 = writeFields_MPI(**args.__dict__) - compareFields_MPI(args.fileName, u0, args.nSteps) + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('--fileName', type=str, help='fileName of the file') + parser.add_argument('--dtypeIdx', type=int, help="dtype index", choices=DTYPES.keys()) + parser.add_argument('--algo', type=str, help="algorithm used for block decomposition") + parser.add_argument('--nSteps', type=int, help="number of time-steps") + parser.add_argument('--nVar', type=int, help="number of field variables") + parser.add_argument('--nX', type=int, nargs='+', help="number of grid points in each dimensions") + args = parser.parse_args() + + u0 = writeFields_MPI(**args.__dict__) + compareFields_MPI(args.fileName, u0, args.nSteps) From 166a0e4ba9124115603756dc112afa01cccba195 Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 01:21:08 +0100 Subject: [PATCH 05/16] TL: I'm dumb --- pySDC/tests/test_helpers/test_fieldsIO.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pySDC/tests/test_helpers/test_fieldsIO.py b/pySDC/tests/test_helpers/test_fieldsIO.py index 0cf826d41..a7dec26e4 100644 --- a/pySDC/tests/test_helpers/test_fieldsIO.py +++ b/pySDC/tests/test_helpers/test_fieldsIO.py @@ -1,7 +1,7 @@ import sys import pytest -if sys.version_info >= (3, 11): +if sys.version_info < (3, 11): pytest.skip("skipping fieldsIO tests on python lower than 3.11", allow_module_level=True) import itertools From bf7e7d26fcc5f3bc81d26acf7b1f93db14dcaea0 Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 01:28:58 +0100 Subject: [PATCH 06/16] TL: yet another try ... --- pySDC/helpers/fieldsIO.py | 58 ++++++++++++++++++++ pySDC/tests/test_helpers/test_fieldsIO.py | 66 +++-------------------- 2 files changed, 64 insertions(+), 60 deletions(-) diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py index dfe3c72eb..4e6feaf8b 100644 --- a/pySDC/helpers/fieldsIO.py +++ b/pySDC/helpers/fieldsIO.py @@ -638,3 +638,61 @@ def readField(self, idx): self.MPI_FILE_CLOSE() return t, field + + +# Utility function used for testing +def initGrid(nVar, gridSizes): + dim = len(gridSizes) + coords = [np.linspace(0, 1, num=n, endpoint=False) for n in gridSizes] + s = [None] * dim + u0 = np.array(np.arange(nVar) + 1)[:, *s] + for x in np.meshgrid(*coords, indexing="ij"): + u0 = u0 * x + return coords, u0 + + +def writeFields_MPI(fileName, dtypeIdx, algo, nSteps, nVar, nX): + coords, u0 = initGrid(nVar, nX) + + from mpi4py import MPI + from pySDC.helpers.blocks import BlockDecomposition + from pySDC.helpers.fieldsIO import Rectilinear + + comm = MPI.COMM_WORLD + MPI_SIZE = comm.Get_size() + MPI_RANK = comm.Get_rank() + + blocks = BlockDecomposition(MPI_SIZE, nX, algo, MPI_RANK) + + iLoc, nLoc = blocks.localBounds + Rectilinear.setupMPI(comm, iLoc, nLoc) + s = [slice(i, i + n) for i, n in zip(iLoc, nLoc)] + u0 = u0[:, *s] + print(MPI_RANK, u0.shape) + + f1 = Rectilinear(DTYPES[dtypeIdx], fileName) + f1.setHeader(nVar=nVar, coords=coords) + + 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"fields[{idx}] in {f2} has incorrect time ({t2} instead of {t})" + 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" diff --git a/pySDC/tests/test_helpers/test_fieldsIO.py b/pySDC/tests/test_helpers/test_fieldsIO.py index a7dec26e4..5ee96fba3 100644 --- a/pySDC/tests/test_helpers/test_fieldsIO.py +++ b/pySDC/tests/test_helpers/test_fieldsIO.py @@ -153,63 +153,6 @@ def testRectilinear(dim, nVar, nSteps, dtypeIdx): assert np.allclose(u2, u1), f"{idx}'s fields in {f1} has incorrect values" -def initGrid(nVar, gridSizes): - dim = len(gridSizes) - coords = [np.linspace(0, 1, num=n, endpoint=False) for n in gridSizes] - s = [None] * dim - u0 = np.array(np.arange(nVar) + 1)[:, *s] - for x in np.meshgrid(*coords, indexing="ij"): - u0 = u0 * x - return coords, u0 - - -def writeFields_MPI(fileName, dtypeIdx, algo, nSteps, nVar, nX): - coords, u0 = initGrid(nVar, nX) - - from mpi4py import MPI - from pySDC.helpers.blocks import BlockDecomposition - from pySDC.helpers.fieldsIO import Rectilinear - - comm = MPI.COMM_WORLD - MPI_SIZE = comm.Get_size() - MPI_RANK = comm.Get_rank() - - blocks = BlockDecomposition(MPI_SIZE, nX, algo, MPI_RANK) - - iLoc, nLoc = blocks.localBounds - Rectilinear.setupMPI(comm, iLoc, nLoc) - s = [slice(i, i + n) for i, n in zip(iLoc, nLoc)] - u0 = u0[:, *s] - print(MPI_RANK, u0.shape) - - f1 = Rectilinear(DTYPES[dtypeIdx], fileName) - f1.setHeader(nVar=nVar, coords=coords) - - 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"fields[{idx}] in {f2} has incorrect time ({t2} instead of {t})" - 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("nVar", [1, 4]) @pytest.mark.parametrize("nSteps", [1, 10]) @@ -232,7 +175,7 @@ def testRectilinear_MPI(dim, nProcs, dtypeIdx, algo, nSteps, nVar): p.wait() assert p.returncode == 0, f"MPI write with {nProcs} proc(s) did not return code 0, but {p.returncode}" - from pySDC.helpers.fieldsIO import FieldsIO, Rectilinear + from pySDC.helpers.fieldsIO import FieldsIO, Rectilinear, initGrid f2: Rectilinear = FieldsIO.fromFile(fileName) @@ -266,5 +209,8 @@ def testRectilinear_MPI(dim, nProcs, dtypeIdx, algo, nSteps, nVar): parser.add_argument('--nX', type=int, nargs='+', help="number of grid points in each dimensions") args = parser.parse_args() - u0 = writeFields_MPI(**args.__dict__) - compareFields_MPI(args.fileName, u0, args.nSteps) + if sys.version_info >= (3, 11): + from pySDC.helpers.fieldsIO import writeFields_MPI, compareFields_MPI + + u0 = writeFields_MPI(**args.__dict__) + compareFields_MPI(args.fileName, u0, args.nSteps) From f0e77712357d1fdf8301c76829dc5e6750b98502 Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 14:01:48 +0100 Subject: [PATCH 07/16] TL: added vtk support for 3D VTR files --- etc/environment-base.yml | 1 + pySDC/helpers/fieldsIO.py | 6 +- pySDC/helpers/vtk.py | 103 +++++++++++++++++++++++++++ pySDC/tests/test_helpers/test_vtk.py | 23 ++++++ 4 files changed, 131 insertions(+), 2 deletions(-) create mode 100644 pySDC/helpers/vtk.py create mode 100644 pySDC/tests/test_helpers/test_vtk.py diff --git a/etc/environment-base.yml b/etc/environment-base.yml index 94b6af8ca..e5d3fb0aa 100644 --- a/etc/environment-base.yml +++ b/etc/environment-base.yml @@ -10,6 +10,7 @@ dependencies: - sympy>=1.0 - numba>=0.35 - dill>=0.2.6 + - vtk - pip - pip: - qmat>=0.1.8 diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py index 4e6feaf8b..80fd6787f 100644 --- a/pySDC/helpers/fieldsIO.py +++ b/pySDC/helpers/fieldsIO.py @@ -43,7 +43,7 @@ To use MPI collective writing, you need to call first the class methods :class:`Rectilinear.initMPI` (cf their docstring). Also, `Rectilinear.setHeader` **must be given the global grids coordinates**, wether the code is run in parallel or not. -Also : this feature is only available with Python 3.11 or higher ... +> ⚠️ Also : this module can only be imported with **Python 3.11 or higher** ! """ import os import numpy as np @@ -640,7 +640,9 @@ def readField(self, idx): return t, field -# Utility function used for testing +# ----------------------------------------------------------------------------------------------- +# Utility functions used for testing +# ----------------------------------------------------------------------------------------------- def initGrid(nVar, gridSizes): dim = len(gridSizes) coords = [np.linspace(0, 1, num=n, endpoint=False) for n in gridSizes] diff --git a/pySDC/helpers/vtk.py b/pySDC/helpers/vtk.py new file mode 100644 index 000000000..fad80e2e5 --- /dev/null +++ b/pySDC/helpers/vtk.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Helper functions to write and read fields from VTK files (to be used with Paraview or PyVista) +""" +import os +import vtk +from vtkmodules.util import numpy_support +import numpy as np + + +def writeToVTR(fileName: str, data, coords, varNames): + """ + Write a data array containing variables from a 3D rectilinear grid into a VTR file. + + Parameters + ---------- + fileName : str + Name of the VTR file, can be with or without the .vtr extension. + data : np.4darray + Array containing all the variables with [nVar, nX, nY, nZ] shape. + coords : list[np.1darray] + Coordinates in each direction. + varNames : list[str] + Variable names. + + Returns + ------- + fileName : str + Name of the VTR file. + """ + data = np.asarray(data) + nVar, *gridSizes = data.shape + + assert len(gridSizes) == 3, "function can be used only for 3D grid data" + assert nVar == len(varNames), "varNames must have as many variable as data" + assert [len(np.ravel(coord)) for coord in coords] == gridSizes, "coordinate size incompatible with data shape" + + nX, nY, nZ = gridSizes + vtr = vtk.vtkRectilinearGrid() + vtr.SetDimensions(nX, nY, nZ) + + vect = lambda x: numpy_support.numpy_to_vtk(num_array=x, deep=True, array_type=vtk.VTK_FLOAT) + x, y, z = coords + vtr.SetXCoordinates(vect(x)) + vtr.SetYCoordinates(vect(y)) + vtr.SetZCoordinates(vect(z)) + + field = lambda u: numpy_support.numpy_to_vtk(num_array=u.ravel(order='F'), deep=True, array_type=vtk.VTK_FLOAT) + pointData = vtr.GetPointData() + for name, u in zip(varNames, data): + uVTK = field(u) + uVTK.SetName(name) + pointData.AddArray(uVTK) + + writer = vtk.vtkXMLRectilinearGridWriter() + if not fileName.endswith(".vtr"): + fileName += ".vtr" + writer.SetFileName(fileName) + writer.SetInputData(vtr) + writer.Write() + + return fileName + + +def readFromVTR(fileName: str): + """ + Read a VTR file into a numpy 4darray + + Parameters + ---------- + fileName : str + Name of the VTR file, can be with or without the .vtr extension. + + Returns + ------- + data : np.4darray + Array containing all the variables with [nVar, nX, nY, nZ] shape. + coords : list[np.1darray] + Coordinates in each direction. + varNames : list[str] + Variable names. + """ + if not fileName.endswith(".vtr"): + fileName += ".vtr" + assert os.path.isfile(fileName), f"{fileName} does not exist" + + reader = vtk.vtkXMLRectilinearGridReader() + reader.SetFileName(fileName) + reader.Update() + + vtr = reader.GetOutput() + dims = vtr.GetDimensions() + assert len(dims) == 3, "can only read 3D data" + + vect = lambda x: numpy_support.vtk_to_numpy(x) + coords = [vect(vtr.GetXCoordinates()), vect(vtr.GetYCoordinates()), vect(vtr.GetZCoordinates())] + pointData = vtr.GetPointData() + varNames = [pointData.GetArrayName(i) for i in range(pointData.GetNumberOfArrays())] + data = [numpy_support.vtk_to_numpy(pointData.GetArray(name)).reshape(dims, order="F") for name in varNames] + data = np.array(data) + + return data, coords, varNames diff --git a/pySDC/tests/test_helpers/test_vtk.py b/pySDC/tests/test_helpers/test_vtk.py new file mode 100644 index 000000000..d69f9f046 --- /dev/null +++ b/pySDC/tests/test_helpers/test_vtk.py @@ -0,0 +1,23 @@ +import pytest +import numpy as np + + +@pytest.mark.parametrize("nZ", [1, 5, 16]) +@pytest.mark.parametrize("nY", [1, 5, 16]) +@pytest.mark.parametrize("nX", [1, 5, 16]) +@pytest.mark.parametrize("nVar", [1, 2, 3]) +def testVTR(nVar, nX, nY, nZ): + from pySDC.helpers.vtk import writeToVTR, readFromVTR + + data1 = np.random.rand(nVar, nX, nY, nZ) + coords1 = [np.sort(np.random.rand(n)) for n in [nX, nY, nZ]] + varNames1 = [f"var{i}" for i in range(nVar)] + + data2, coords2, varNames2 = readFromVTR(writeToVTR("testVTR", data1, coords1, varNames1)) + + for i, (x1, x2) in enumerate(zip(coords1, coords2)): + print(x1, x2) + assert np.allclose(x1, x2), f"coordinate mismatch in dir. {i}" + assert varNames1 == varNames2, f"varNames mismatch" + assert data1.shape == data2.shape, f"data shape mismatch" + assert np.allclose(data1, data2), f"data values mismatch" From e52d0e9dae2203489a23277464619e7cbd654416 Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 14:07:57 +0100 Subject: [PATCH 08/16] TL: avoided clashing name, gitignore update --- .gitignore | 1 + pySDC/helpers/{vtk.py => vtkIO.py} | 2 +- pySDC/tests/test_helpers/test_vtk.py | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) rename pySDC/helpers/{vtk.py => vtkIO.py} (97%) diff --git a/.gitignore b/.gitignore index 845782e78..d72c27e16 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,7 @@ step_*.png *_data.json !_dataRef.json *.pysdc +*.vtr # Created by https://www.gitignore.io diff --git a/pySDC/helpers/vtk.py b/pySDC/helpers/vtkIO.py similarity index 97% rename from pySDC/helpers/vtk.py rename to pySDC/helpers/vtkIO.py index fad80e2e5..544929291 100644 --- a/pySDC/helpers/vtk.py +++ b/pySDC/helpers/vtkIO.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- """ -Helper functions to write and read fields from VTK files (to be used with Paraview or PyVista) +Helper functions for VTK files IO (to be used with Paraview or PyVista) """ import os import vtk diff --git a/pySDC/tests/test_helpers/test_vtk.py b/pySDC/tests/test_helpers/test_vtk.py index d69f9f046..94613ac6f 100644 --- a/pySDC/tests/test_helpers/test_vtk.py +++ b/pySDC/tests/test_helpers/test_vtk.py @@ -7,7 +7,7 @@ @pytest.mark.parametrize("nX", [1, 5, 16]) @pytest.mark.parametrize("nVar", [1, 2, 3]) def testVTR(nVar, nX, nY, nZ): - from pySDC.helpers.vtk import writeToVTR, readFromVTR + from pySDC.helpers.vtkIO import writeToVTR, readFromVTR data1 = np.random.rand(nVar, nX, nY, nZ) coords1 = [np.sort(np.random.rand(n)) for n in [nX, nY, nZ]] From e8d5631d092f466152b04bad4722e194c5e33546 Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 14:12:03 +0100 Subject: [PATCH 09/16] TL: you gotta be kiding me ... --- pySDC/helpers/vtkIO.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/pySDC/helpers/vtkIO.py b/pySDC/helpers/vtkIO.py index 544929291..8fe2a06fb 100644 --- a/pySDC/helpers/vtkIO.py +++ b/pySDC/helpers/vtkIO.py @@ -40,13 +40,17 @@ def writeToVTR(fileName: str, data, coords, varNames): vtr = vtk.vtkRectilinearGrid() vtr.SetDimensions(nX, nY, nZ) - vect = lambda x: numpy_support.numpy_to_vtk(num_array=x, deep=True, array_type=vtk.VTK_FLOAT) + def vect(x): + return numpy_support.numpy_to_vtk(num_array=x, deep=True, array_type=vtk.VTK_FLOAT) + x, y, z = coords vtr.SetXCoordinates(vect(x)) vtr.SetYCoordinates(vect(y)) vtr.SetZCoordinates(vect(z)) - field = lambda u: numpy_support.numpy_to_vtk(num_array=u.ravel(order='F'), deep=True, array_type=vtk.VTK_FLOAT) + def field(u): + return numpy_support.numpy_to_vtk(num_array=u.ravel(order='F'), deep=True, array_type=vtk.VTK_FLOAT) + pointData = vtr.GetPointData() for name, u in zip(varNames, data): uVTK = field(u) @@ -93,7 +97,9 @@ def readFromVTR(fileName: str): dims = vtr.GetDimensions() assert len(dims) == 3, "can only read 3D data" - vect = lambda x: numpy_support.vtk_to_numpy(x) + def vect(x): + return numpy_support.vtk_to_numpy(x) + coords = [vect(vtr.GetXCoordinates()), vect(vtr.GetYCoordinates()), vect(vtr.GetZCoordinates())] pointData = vtr.GetPointData() varNames = [pointData.GetArrayName(i) for i in range(pointData.GetNumberOfArrays())] From 5569514314713f4a0eeedebbbc35b6dd5906d2a5 Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 14:50:43 +0100 Subject: [PATCH 10/16] TL: VTK output support for Rectilinear FieldsIO --- pySDC/helpers/fieldsIO.py | 9 +++++ pySDC/tests/test_helpers/test_fieldsIO.py | 44 ++++++++++++++++++++++- 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py index 80fd6787f..ce46ab0c5 100644 --- a/pySDC/helpers/fieldsIO.py +++ b/pySDC/helpers/fieldsIO.py @@ -453,6 +453,15 @@ def nDoF(self): """Number of degrees of freedom for one variable""" return np.prod(self.nX) + def toVTR(self, baseName, varNames, suffix="{:06d}_t={:1.2f}s"): + assert self.dim == 3, "can only be used with 3D fields" + from pySDC.helpers.vtkIO import writeToVTR + + template = f"{baseName}_{suffix}" + for i in range(self.nFields): + t, u = self.readField(i) + writeToVTR(template.format(i, t), u, self.header["coords"], varNames) + # ------------------------------------------------------------------------- # MPI-parallel implementation # ------------------------------------------------------------------------- diff --git a/pySDC/tests/test_helpers/test_fieldsIO.py b/pySDC/tests/test_helpers/test_fieldsIO.py index 5ee96fba3..e1dc7c6cd 100644 --- a/pySDC/tests/test_helpers/test_fieldsIO.py +++ b/pySDC/tests/test_helpers/test_fieldsIO.py @@ -1,4 +1,6 @@ +import os import sys +import glob import pytest if sys.version_info < (3, 11): @@ -153,6 +155,46 @@ def testRectilinear(dim, nVar, nSteps, dtypeIdx): assert np.allclose(u2, u1), f"{idx}'s fields in {f1} has incorrect values" +@pytest.mark.parametrize("nSteps", [1, 10]) +@pytest.mark.parametrize("nZ", [1, 5, 16]) +@pytest.mark.parametrize("nY", [1, 5, 16]) +@pytest.mark.parametrize("nX", [1, 5, 16]) +@pytest.mark.parametrize("nVar", [1, 2, 3]) +def testToVTR(nVar, nX, nY, nZ, nSteps): + + from pySDC.helpers.fieldsIO import Rectilinear + from pySDC.helpers.vtkIO import readFromVTR + + coords = [np.linspace(0, 1, num=n, endpoint=False) for n in [nX, nY, nZ]] + file = Rectilinear(np.float64, "testToVTR.pysdc") + file.setHeader(nVar=nVar, coords=coords) + file.initialize() + u0 = np.random.rand(nVar, nX, nY, nZ).astype(file.dtype) + times = np.arange(nSteps) / nSteps + for t in times: + ut = (u0 * t).astype(file.dtype) + file.addField(t, ut) + + # Cleaning after eventuall other tests ... + for f in glob.glob("testToVTR*.vtr"): + os.remove(f) + + file.toVTR("testToVTR", varNames=[f"var{i}" for i in range(nVar)]) + + vtrFiles = glob.glob("testToVTR*.vtr") + assert len(vtrFiles) == file.nFields + + vtrFiles.sort(key=lambda name: int(name.split("_")[1])) + for i, vFile in enumerate(vtrFiles): + uVTR, coords, _ = readFromVTR(vFile) + tVTR = float(vFile.split("_t=")[-1].split("s.vtr")[0]) + tFile, uFile = file.readField(i) + assert np.allclose(tFile, tVTR), "mismatch between field times" + assert np.allclose(uFile, uVTR), "mismatch between data" + for i, (xVTR, xFile) in enumerate(zip(coords, file.header["coords"])): + assert np.allclose(xVTR, xFile), f"coordinate mismatch in dir. {i}" + + @pytest.mark.mpi4py @pytest.mark.parametrize("nVar", [1, 4]) @pytest.mark.parametrize("nSteps", [1, 10]) @@ -175,7 +217,7 @@ def testRectilinear_MPI(dim, nProcs, dtypeIdx, algo, nSteps, nVar): p.wait() assert p.returncode == 0, f"MPI write with {nProcs} proc(s) did not return code 0, but {p.returncode}" - from pySDC.helpers.fieldsIO import FieldsIO, Rectilinear, initGrid + from pySDC.helpers.fieldsIO import Rectilinear, initGrid f2: Rectilinear = FieldsIO.fromFile(fileName) From f0136e68b36f908026c20ec7b46b2895229875d5 Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 15:01:39 +0100 Subject: [PATCH 11/16] TL: added documentation --- pySDC/helpers/fieldsIO.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py index ce46ab0c5..68d9666ae 100644 --- a/pySDC/helpers/fieldsIO.py +++ b/pySDC/helpers/fieldsIO.py @@ -454,6 +454,31 @@ def nDoF(self): return np.prod(self.nX) def toVTR(self, baseName, varNames, suffix="{:06d}_t={:1.2f}s"): + """ + Convert all 3D fields stored in binary format (FieldsIO) into a list + of VTR files, that can be read later with Paraview or equivalent to + make videos. + + Parameters + ---------- + baseName : str + Base name of the VTR file. + varNames : list[str] + Variable names of the fields. + suffix : str, optional + Formating string for the suffix of the VTR file, containing the + index in first position, and the time in second position. + The default is "{:06d}_t={:1.2f}s". + + Example + ------- + >>> # Suppose the FieldsIO object is already writen into outputs.pysdc + >>> import os + >>> from pySDC.utils.fieldsIO import Rectilinear + >>> os.makedirs("vtrFiles") # to store all VTR files into a subfolder + >>> Rectilinear.fromFile("outputs.pysdc").toVTR( + >>> baseName="field", varNames=["u", "v", "w", "T", "p"]) + """ assert self.dim == 3, "can only be used with 3D fields" from pySDC.helpers.vtkIO import writeToVTR From 56a74b5d20ded4f65b70712c6cbee08c5dfd206f Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 15:34:11 +0100 Subject: [PATCH 12/16] TL: polished documentation --- pySDC/helpers/fieldsIO.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py index 68d9666ae..7efd72046 100644 --- a/pySDC/helpers/fieldsIO.py +++ b/pySDC/helpers/fieldsIO.py @@ -14,19 +14,23 @@ Example ------- >>> import numpy as np ->>> from pySDC.helpers.fieldsIO import Rectilinear, FieldsIO +>>> from pySDC.helpers.fieldsIO import Rectilinear >>> >>> # Write some fields in files ->>> x = np.linspace(0, 1, 101) +>>> x = np.linspace(0, 1, 128) +>>> y = np.linspace(0, 1, 64) >>> fOut = Rectilinear(np.float64, "file.pysdc") ->>> fOut.setHeader(nVar=2, coords=x) +>>> fOut.setHeader(nVar=2, coords=[x, y]) >>> fOut.initialize() >>> times = [0, 1, 2] ->>> u0 = np.array([-1, 1])[:, None]*x[None, :] +>>> xGrid, yGrid = np.meshgrid(x, y, indexing="ij") +>>> u0 = np.array([-1, 1]).reshape((-1, 1, 1))*xGrid*yGrid +>>> # u0 has shape [2, nX, nY] >>> for t in times: >>> fOut.addField(t, t*u0) >>> ->>> # Read the file using a the generic interface +>>> # Read the file using the generic interface +>>> from pySDC.helpers.fieldsIO import FieldsIO >>> fIn = FieldsIO.fromFile("file.pysdc") >>> times = fIn.times >>> assert len(times) == fIn.nFields @@ -36,7 +40,7 @@ Notes ----- 🚀 :class:`Rectilinear` is compatible with a MPI-based cartesian decomposition. -See :class:`pySDC.tests.test_helpers.test_fieldsIO.writeFields_MPI` for an illustrative example. +See :class:`pySDC.helpers.fieldsIO.writeFields_MPI` for an illustrative example. Warning ------- From 7bac672b609e58c4b65926ed2537c9321356c8bd Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 15:57:07 +0100 Subject: [PATCH 13/16] TL: minor typo --- pySDC/helpers/fieldsIO.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py index 7efd72046..475bcac5f 100644 --- a/pySDC/helpers/fieldsIO.py +++ b/pySDC/helpers/fieldsIO.py @@ -481,7 +481,7 @@ def toVTR(self, baseName, varNames, suffix="{:06d}_t={:1.2f}s"): >>> from pySDC.utils.fieldsIO import Rectilinear >>> os.makedirs("vtrFiles") # to store all VTR files into a subfolder >>> Rectilinear.fromFile("outputs.pysdc").toVTR( - >>> baseName="field", varNames=["u", "v", "w", "T", "p"]) + >>> baseName="vtrFiles/field", varNames=["u", "v", "w", "T", "p"]) """ assert self.dim == 3, "can only be used with 3D fields" from pySDC.helpers.vtkIO import writeToVTR From 2795143477c8164c046fdc404589bfcdcb2dae97 Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 16:39:07 +0100 Subject: [PATCH 14/16] TL: refactoring following thomas's suggestions --- pySDC/helpers/fieldsIO.py | 36 +++++++++++------------ pySDC/helpers/vtkIO.py | 6 ++-- pySDC/tests/test_helpers/test_fieldsIO.py | 22 +++++++------- 3 files changed, 32 insertions(+), 32 deletions(-) diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py index 475bcac5f..0ae918916 100644 --- a/pySDC/helpers/fieldsIO.py +++ b/pySDC/helpers/fieldsIO.py @@ -417,7 +417,7 @@ def setHeader(self, nVar, coords): @property def hInfos(self): """Array representing the grid structure to be written in the binary file.""" - return [np.array([self.nVar, self.dim, *self.nX], dtype=np.int32)] + [ + return [np.array([self.nVar, self.dim, *self.gridSizes], dtype=np.int32)] + [ np.array(coord, dtype=np.float64) for coord in self.header["coords"] ] @@ -431,31 +431,31 @@ def readHeader(self, f): File to read the header from. """ nVar, dim = np.fromfile(f, dtype=np.int32, count=2) - nX = np.fromfile(f, dtype=np.int32, count=dim) - coords = [np.fromfile(f, dtype=np.float64, count=n) for n in nX] + gridSizes = np.fromfile(f, dtype=np.int32, count=dim) + coords = [np.fromfile(f, dtype=np.float64, count=n) for n in gridSizes] self.setHeader(nVar, coords) def reshape(self, fields: np.ndarray): """Reshape the fields to a N-d array (inplace operation)""" - fields.shape = (self.nVar, *self.nX) + fields.shape = (self.nVar, *self.gridSizes) # ------------------------------------------------------------------------- # Class specifics # ------------------------------------------------------------------------- @property - def nX(self): + def gridSizes(self): """Number of points in y direction""" return [coord.size for coord in self.header["coords"]] @property def dim(self): """Number of grid dimensions""" - return len(self.nX) + return len(self.gridSizes) @property def nDoF(self): """Number of degrees of freedom for one variable""" - return np.prod(self.nX) + return np.prod(self.gridSizes) def toVTR(self, baseName, varNames, suffix="{:06d}_t={:1.2f}s"): """ @@ -625,22 +625,22 @@ def addField(self, time, field): self.MPI_WRITE(np.array(time, dtype=T_DTYPE)) offset0 += self.tSize - for (iVar, *iX) in itertools.product(range(self.nVar), *[range(nX) for nX in self.nLoc[:-1]]): - offset = offset0 + self.iPos(iVar, iX) * self.itemSize - self.MPI_WRITE_AT(offset, field[iVar, *iX]) + for (iVar, *iBeg) in itertools.product(range(self.nVar), *[range(n) for n in self.nLoc[:-1]]): + offset = offset0 + self.iPos(iVar, iBeg) * self.itemSize + self.MPI_WRITE_AT(offset, field[iVar, *iBeg]) self.MPI_FILE_CLOSE() def iPos(self, iVar, iX): iPos = iVar * self.nDoF for axis in range(self.dim - 1): - iPos += (self.iLoc[axis] + iX[axis]) * np.prod(self.nX[axis + 1 :]) + iPos += (self.iLoc[axis] + iX[axis]) * np.prod(self.gridSizes[axis + 1 :]) iPos += self.iLoc[-1] return iPos def readField(self, idx): """ Read one field stored in the binary file, corresponding to the given - time index, eventually in MPI mode. + time index, using MPI in the eventuality of space parallel decomposition. Parameters ---------- @@ -670,9 +670,9 @@ def readField(self, idx): field = np.empty((self.nVar, *self.nLoc), dtype=self.dtype) self.MPI_FILE_OPEN(mode="r") - for (iVar, *iX) in itertools.product(range(self.nVar), *[range(nX) for nX in self.nLoc[:-1]]): - offset = offset0 + self.iPos(iVar, iX) * self.itemSize - self.MPI_READ_AT(offset, field[iVar, *iX]) + for (iVar, *iBeg) in itertools.product(range(self.nVar), *[range(n) for n in self.nLoc[:-1]]): + offset = offset0 + self.iPos(iVar, iBeg) * self.itemSize + self.MPI_READ_AT(offset, field[iVar, *iBeg]) self.MPI_FILE_CLOSE() return t, field @@ -691,8 +691,8 @@ def initGrid(nVar, gridSizes): return coords, u0 -def writeFields_MPI(fileName, dtypeIdx, algo, nSteps, nVar, nX): - coords, u0 = initGrid(nVar, nX) +def writeFields_MPI(fileName, dtypeIdx, algo, nSteps, nVar, gridSizes): + coords, u0 = initGrid(nVar, gridSizes) from mpi4py import MPI from pySDC.helpers.blocks import BlockDecomposition @@ -702,7 +702,7 @@ def writeFields_MPI(fileName, dtypeIdx, algo, nSteps, nVar, nX): MPI_SIZE = comm.Get_size() MPI_RANK = comm.Get_rank() - blocks = BlockDecomposition(MPI_SIZE, nX, algo, MPI_RANK) + blocks = BlockDecomposition(MPI_SIZE, gridSizes, algo, MPI_RANK) iLoc, nLoc = blocks.localBounds Rectilinear.setupMPI(comm, iLoc, nLoc) diff --git a/pySDC/helpers/vtkIO.py b/pySDC/helpers/vtkIO.py index 8fe2a06fb..9cf0c7f99 100644 --- a/pySDC/helpers/vtkIO.py +++ b/pySDC/helpers/vtkIO.py @@ -94,8 +94,8 @@ def readFromVTR(fileName: str): reader.Update() vtr = reader.GetOutput() - dims = vtr.GetDimensions() - assert len(dims) == 3, "can only read 3D data" + gridSizes = vtr.GetDimensions() + assert len(gridSizes) == 3, "can only read 3D data" def vect(x): return numpy_support.vtk_to_numpy(x) @@ -103,7 +103,7 @@ def vect(x): coords = [vect(vtr.GetXCoordinates()), vect(vtr.GetYCoordinates()), vect(vtr.GetZCoordinates())] pointData = vtr.GetPointData() varNames = [pointData.GetArrayName(i) for i in range(pointData.GetNumberOfArrays())] - data = [numpy_support.vtk_to_numpy(pointData.GetArray(name)).reshape(dims, order="F") for name in varNames] + data = [numpy_support.vtk_to_numpy(pointData.GetArray(name)).reshape(gridSizes, order="F") for name in varNames] data = np.array(data) return data, coords, varNames diff --git a/pySDC/tests/test_helpers/test_fieldsIO.py b/pySDC/tests/test_helpers/test_fieldsIO.py index e1dc7c6cd..d13fefd3e 100644 --- a/pySDC/tests/test_helpers/test_fieldsIO.py +++ b/pySDC/tests/test_helpers/test_fieldsIO.py @@ -116,20 +116,20 @@ def testRectilinear(dim, nVar, nSteps, dtypeIdx): fileName = f"testRectilinear{dim}D.pysdc" dtype = DTYPES[dtypeIdx] - for nX in itertools.product(*[[5, 10, 16]] * dim): + for gridSizes in itertools.product(*[[5, 10, 16]] * dim): - coords = [np.linspace(0, 1, num=n, endpoint=False) for n in nX] + coords = [np.linspace(0, 1, num=n, endpoint=False) for n in gridSizes] f1 = Rectilinear(dtype, fileName) f1.setHeader(nVar=nVar, coords=coords) assert f1.dim == dim, f"{f1} has incorrect dimension" - assert f1.nX == list(nX), f"{f1} has incorrect nX" - assert f1.nDoF == np.prod(nX), f"{f1} has incorrect nDOF" - assert f1.nItems == nVar * np.prod(nX), f"{f1} do not have nItems == nVar*nX**dim" + assert f1.gridSizes == list(gridSizes), f"{f1} has incorrect gridSizes" + assert f1.nDoF == np.prod(gridSizes), f"{f1} has incorrect nDOF" + assert f1.nItems == nVar * np.prod(gridSizes), f"{f1} do not have nItems == nVar*product(gridSizes)" f1.initialize() - u0 = np.random.rand(nVar, *nX).astype(f1.dtype) + u0 = np.random.rand(nVar, *gridSizes).astype(f1.dtype) times = np.arange(nSteps) / nSteps for t in times: @@ -208,10 +208,10 @@ def testRectilinear_MPI(dim, nProcs, dtypeIdx, algo, nSteps, nVar): fileName = f"testRectilinear{dim}D_MPI.pysdc" - for nX in itertools.product(*[[61, 16]] * dim): + for gridSizes in itertools.product(*[[61, 16]] * dim): cmd = f"mpirun -np {nProcs} python {__file__} --fileName {fileName}" - cmd += f" --dtypeIdx {dtypeIdx} --algo {algo} --nSteps {nSteps} --nVar {nVar} --nX {' '.join([str(n) for n in nX])}" + cmd += f" --dtypeIdx {dtypeIdx} --algo {algo} --nSteps {nSteps} --nVar {nVar} --gridSizes {' '.join([str(n) for n in gridSizes])}" p = subprocess.Popen(cmd.split(), cwd=".") p.wait() @@ -224,9 +224,9 @@ def testRectilinear_MPI(dim, nProcs, dtypeIdx, algo, nSteps, nVar): assert type(f2) == Rectilinear, f"incorrect type in MPI written fields {f2}" assert f2.nFields == nSteps, f"incorrect nFields in MPI written fields {f2} ({f2.nFields} instead of {nSteps})" assert f2.nVar == nVar, f"incorrect nVar in MPI written fields {f2}" - assert f2.nX == list(nX), f"incorrect nX in MPI written fields {f2}" + assert f2.gridSizes == list(gridSizes), f"incorrect gridSizes in MPI written fields {f2}" - coords, u0 = initGrid(nVar, nX) + coords, u0 = initGrid(nVar, gridSizes) for i, (cFile, cRef) in enumerate(zip(f2.header['coords'], coords)): assert np.allclose(cFile, cRef), f"incorrect coords[{i}] in MPI written fields {f2}" @@ -248,7 +248,7 @@ def testRectilinear_MPI(dim, nProcs, dtypeIdx, algo, nSteps, nVar): parser.add_argument('--algo', type=str, help="algorithm used for block decomposition") parser.add_argument('--nSteps', type=int, help="number of time-steps") parser.add_argument('--nVar', type=int, help="number of field variables") - parser.add_argument('--nX', type=int, nargs='+', help="number of grid points in each dimensions") + parser.add_argument('--gridSizes', type=int, nargs='+', help="number of grid points in each dimensions") args = parser.parse_args() if sys.version_info >= (3, 11): From 1874244aef1388db2306247d777a5dda825520bd Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 16:40:59 +0100 Subject: [PATCH 15/16] TL: minor typo --- pySDC/helpers/fieldsIO.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py index 0ae918916..2a515f812 100644 --- a/pySDC/helpers/fieldsIO.py +++ b/pySDC/helpers/fieldsIO.py @@ -507,7 +507,7 @@ def setupMPI(cls, comm: MPI.Intracomm, iLoc, nLoc): comm : MPI.Intracomm The space decomposition communicator. iLoc : list[int] - Starting index of the local sub-domain in the global `coordX`. + Starting index of the local sub-domain in the global coordinates. nLoc : list[int] Number of points in the local sub-domain. """ From a132a92f5c55ce5444ca0ad8e8d89bee70ef219c Mon Sep 17 00:00:00 2001 From: Thibaut Lunet Date: Wed, 19 Feb 2025 17:41:17 +0100 Subject: [PATCH 16/16] TL: corrected Rectilinear.toVTR for paraview groups --- pySDC/helpers/fieldsIO.py | 15 +++++++-------- pySDC/tests/test_helpers/test_fieldsIO.py | 6 ++---- 2 files changed, 9 insertions(+), 12 deletions(-) diff --git a/pySDC/helpers/fieldsIO.py b/pySDC/helpers/fieldsIO.py index 2a515f812..3b124dd10 100644 --- a/pySDC/helpers/fieldsIO.py +++ b/pySDC/helpers/fieldsIO.py @@ -457,7 +457,7 @@ def nDoF(self): """Number of degrees of freedom for one variable""" return np.prod(self.gridSizes) - def toVTR(self, baseName, varNames, suffix="{:06d}_t={:1.2f}s"): + def toVTR(self, baseName, varNames, idxFormat="{:06d}"): """ Convert all 3D fields stored in binary format (FieldsIO) into a list of VTR files, that can be read later with Paraview or equivalent to @@ -469,10 +469,9 @@ def toVTR(self, baseName, varNames, suffix="{:06d}_t={:1.2f}s"): Base name of the VTR file. varNames : list[str] Variable names of the fields. - suffix : str, optional - Formating string for the suffix of the VTR file, containing the - index in first position, and the time in second position. - The default is "{:06d}_t={:1.2f}s". + idxFormat : str, optional + Formatting string for the index of the VTR file. + The default is "{:06d}". Example ------- @@ -486,10 +485,10 @@ def toVTR(self, baseName, varNames, suffix="{:06d}_t={:1.2f}s"): assert self.dim == 3, "can only be used with 3D fields" from pySDC.helpers.vtkIO import writeToVTR - template = f"{baseName}_{suffix}" + template = f"{baseName}_{idxFormat}" for i in range(self.nFields): - t, u = self.readField(i) - writeToVTR(template.format(i, t), u, self.header["coords"], varNames) + _, u = self.readField(i) + writeToVTR(template.format(i), u, self.header["coords"], varNames) # ------------------------------------------------------------------------- # MPI-parallel implementation diff --git a/pySDC/tests/test_helpers/test_fieldsIO.py b/pySDC/tests/test_helpers/test_fieldsIO.py index d13fefd3e..1e75505d0 100644 --- a/pySDC/tests/test_helpers/test_fieldsIO.py +++ b/pySDC/tests/test_helpers/test_fieldsIO.py @@ -184,12 +184,10 @@ def testToVTR(nVar, nX, nY, nZ, nSteps): vtrFiles = glob.glob("testToVTR*.vtr") assert len(vtrFiles) == file.nFields - vtrFiles.sort(key=lambda name: int(name.split("_")[1])) + vtrFiles.sort(key=lambda name: int(name.split("_")[-1][:-4])) for i, vFile in enumerate(vtrFiles): uVTR, coords, _ = readFromVTR(vFile) - tVTR = float(vFile.split("_t=")[-1].split("s.vtr")[0]) - tFile, uFile = file.readField(i) - assert np.allclose(tFile, tVTR), "mismatch between field times" + _, uFile = file.readField(i) assert np.allclose(uFile, uVTR), "mismatch between data" for i, (xVTR, xFile) in enumerate(zip(coords, file.header["coords"])): assert np.allclose(xVTR, xFile), f"coordinate mismatch in dir. {i}"