Skip to content

Commit

Permalink
Merge branch 'fieldsIO' into neuralpint
Browse files Browse the repository at this point in the history
  • Loading branch information
@tlunet
tlunet committed Jan 17, 2025
2 parents f493d63 + 8ee8815 commit 450f6be
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 43 deletions.
46 changes: 23 additions & 23 deletions pySDC/helpers/fieldsIO.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
>>> # Write some fields in files
>>> x = np.linspace(0, 1, 101)
>>> fOut = Cart2D(np.float64, "file.pysdc")
>>> fOut.setHeader(nVar=2, gridX=x)
>>> fOut.setHeader(nVar=2, coordX=x)
>>> fOut.initialize()
>>> times = [0, 1, 2]
>>> u0 = np.array([-1, 1])[:, None]*x[None, :]
Expand Down Expand Up @@ -365,29 +365,29 @@ class Cart1D(Scal0D):
# -------------------------------------------------------------------------
# Overriden methods
# -------------------------------------------------------------------------
def setHeader(self, nVar, gridX):
def setHeader(self, nVar, coordX):
"""
Set the descriptive grid structure to be stored in the file header.
Parameters
----------
nVar : int
Number of 1D variables stored.
gridX : np.1darray
coordX : np.1darray
The grid coordinates in X direction.
Note
----
When used in MPI decomposition mode, `gridX` **must** be the global grid.
When used in MPI decomposition mode, `coordX` **must** be the global grid.
"""
grids = self.setupGrids(gridX=gridX)
grids = self.setupCoords(coordX=coordX)
self.header = {"nVar": int(nVar), **grids}
self.nItems = nVar * self.nX

@property
def hInfos(self):
"""Array representing the grid structure to be written in the binary file."""
return [np.array([self.nVar, self.nX], dtype=np.int64), np.array(self.header["gridX"], dtype=np.float64)]
return [np.array([self.nVar, self.nX], dtype=np.int64), np.array(self.header["coordX"], dtype=np.float64)]

def readHeader(self, f):
"""
Expand All @@ -399,8 +399,8 @@ def readHeader(self, f):
File to read the header from.
"""
nVar, nX = np.fromfile(f, dtype=np.int64, count=2)
gridX = np.fromfile(f, dtype=np.float64, count=nX)
self.setHeader(nVar, gridX)
coordX = np.fromfile(f, dtype=np.float64, count=nX)
self.setHeader(nVar, coordX)

def reshape(self, fields: np.ndarray):
fields.shape = (self.nVar, self.nX)
Expand All @@ -411,10 +411,10 @@ def reshape(self, fields: np.ndarray):
@property
def nX(self):
"""Number of points in x direction"""
return self.header["gridX"].size
return self.header["coordX"].size

@staticmethod
def setupGrids(**grids):
def setupCoords(**grids):
"""Utility function to setup grids in multuple dimensions, given the keyword arguments"""
grids = {name: np.asarray(grid, dtype=np.float64) for name, grid in grids.items() if name.startswith("grid")}
for name, grid in grids.items():
Expand All @@ -437,7 +437,7 @@ def setupMPI(cls, comm: MPI.Intracomm, iLocX, nLocX):
comm : MPI.Intracomm
The space decomposition communicator.
iLocX : int
Starting index of the local sub-domain in the global `gridX`.
Starting index of the local sub-domain in the global `coordX`.
nLocX : int
Number of points in the local sub-domain.
"""
Expand Down Expand Up @@ -602,34 +602,34 @@ class Cart2D(Cart1D):
# -------------------------------------------------------------------------
# Overriden methods
# -------------------------------------------------------------------------
def setHeader(self, nVar, gridX, gridY):
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.
gridX : np.1darray
coordX : np.1darray
The grid coordinates in x direction.
gridY : np.1darray
coordY : np.1darray
The grid coordinates in y direction.
Note
----
When used in MPI decomposition mode, `gridX` and `gridY` **must** be the global grid.
When used in MPI decomposition mode, `coordX` and `coordX` **must** be the global grid.
"""
grids = self.setupGrids(gridX=gridX, gridY=gridY)
self.header = {"nVar": int(nVar), **grids}
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["gridX"], dtype=np.float64),
np.array(self.header["gridY"], dtype=np.float64),
np.array(self.header["coordX"], dtype=np.float64),
np.array(self.header["coordY"], dtype=np.float64),
]

def readHeader(self, f):
Expand All @@ -642,9 +642,9 @@ def readHeader(self, f):
File to read the header from.
"""
nVar, nX, nY = np.fromfile(f, dtype=np.int64, count=3)
gridX = np.fromfile(f, dtype=np.float64, count=nX)
gridY = np.fromfile(f, dtype=np.float64, count=nY)
self.setHeader(nVar, gridX, gridY)
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)"""
Expand All @@ -656,7 +656,7 @@ def reshape(self, fields: np.ndarray):
@property
def nY(self):
"""Number of points in y direction"""
return self.header["gridY"].size
return self.header["coordY"].size

# -------------------------------------------------------------------------
# MPI-parallel implementation
Expand Down
40 changes: 20 additions & 20 deletions pySDC/tests/test_helpers/test_fieldsIO.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,18 +14,18 @@ def testHeader(nDim, dtypeIdx):
fileName = "testHeader.pysdc"
dtype = DTYPES[dtypeIdx]

gridX = np.linspace(0, 1, num=256, endpoint=False)
gridY = np.linspace(0, 1, num=64, endpoint=False)
coordX = np.linspace(0, 1, num=256, endpoint=False)
coordY = np.linspace(0, 1, num=64, endpoint=False)

if nDim == 0:
Class = Scal0D
args = {"nVar": 20}
elif nDim == 1:
Class = Cart1D
args = {"nVar": 10, "gridX": gridX}
args = {"nVar": 10, "coordX": coordX}
elif nDim == 2:
Class = Cart2D
args = {"nVar": 10, "gridX": gridX, "gridY": gridY}
args = {"nVar": 10, "coordX": coordX, "coordY": coordY}

f1 = Class(dtype, fileName)
assert f1.__str__() == f1.__repr__(), "__repr__ and __str__ do not return the same result"
Expand Down Expand Up @@ -112,11 +112,11 @@ def testCart1D(nVar, nX, nSteps, dtypeIdx):
fileName = "testCart1D.pysdc"
dtype = DTYPES[dtypeIdx]

gridX = np.linspace(0, 1, num=nX, endpoint=False)
nX = gridX.size
coordX = np.linspace(0, 1, num=nX, endpoint=False)
nX = coordX.size

f1 = Cart1D(dtype, fileName)
f1.setHeader(nVar=nVar, gridX=gridX)
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"
Expand Down Expand Up @@ -159,11 +159,11 @@ def testCart2D(nVar, nX, nY, nSteps, dtypeIdx):
fileName = "testCart2D.pysdc"
dtype = DTYPES[dtypeIdx]

gridX = np.linspace(0, 1, num=nX, endpoint=False)
gridY = np.linspace(0, 1, num=nY, endpoint=False)
coordX = np.linspace(0, 1, num=nX, endpoint=False)
coordY = np.linspace(0, 1, num=nY, endpoint=False)

f1 = Cart2D(dtype, fileName)
f1.setHeader(nVar=nVar, gridX=gridX, gridY=gridY)
f1.setHeader(nVar=nVar, coordX=coordX, coordY=coordY)

assert f1.nItems == nVar * nX * nY, f"{f1} do not have nItems == nVar*nX"
assert f1.nX == nX, f"{f1} has incorrect nX"
Expand Down Expand Up @@ -201,21 +201,21 @@ def initGrid(nVar, nX, nY=None):
if nY is not None:
nDim += 1
x = np.linspace(0, 1, num=nX, endpoint=False)
grids = (x,)
coords = (x,)
gridSizes = (nX,)
u0 = np.array(np.arange(nVar) + 1)[:, None] * x[None, :]

if nDim > 1:
y = np.linspace(0, 1, num=nY, endpoint=False)
grids += (y,)
coords += (y,)
gridSizes += (nY,)
u0 = u0[:, :, None] * y[None, None, :]

return grids, gridSizes, u0
return coords, gridSizes, u0


def writeFields_MPI(fileName, nDim, dtypeIdx, algo, nSteps, nVar, nX, nY=None):
grids, gridSizes, u0 = initGrid(nVar, nX, nY)
coords, gridSizes, u0 = initGrid(nVar, nX, nY)

from mpi4py import MPI
from pySDC.helpers.blocks import BlockDecomposition
Expand All @@ -234,15 +234,15 @@ def writeFields_MPI(fileName, nDim, dtypeIdx, algo, nSteps, nVar, nX, nY=None):
u0 = u0[:, iLocX : iLocX + nLocX]

f1 = Cart1D(DTYPES[dtypeIdx], fileName)
f1.setHeader(nVar=nVar, gridX=grids[0])
f1.setHeader(nVar=nVar, coordX=coords[0])

if nDim == 2:
(iLocX, iLocY), (nLocX, nLocY) = blocks.localBounds
Cart2D.setupMPI(comm, iLocX, nLocX, iLocY, nLocY)
u0 = u0[:, iLocX : iLocX + nLocX, iLocY : iLocY + nLocY]

f1 = Cart2D(DTYPES[dtypeIdx], fileName)
f1.setHeader(nVar=nVar, gridX=grids[0], gridY=grids[1])
f1.setHeader(nVar=nVar, coordX=coords[0], coordY=coords[1])

u0 = np.asarray(u0, dtype=f1.dtype)
f1.initialize()
Expand Down Expand Up @@ -298,8 +298,8 @@ def testCart1D_MPI(nProcs, dtypeIdx, algo, nSteps, nVar, nX):
assert f2.nVar == nVar, f"incorrect nVar in MPI written fields {f2}"
assert f2.nX == nX, f"incorrect nX in MPI written fields {f2}"

grids, _, u0 = initGrid(nVar, nX)
assert np.allclose(f2.header['gridX'], grids[0]), f"incorrect gridX in MPI written fields {f2}"
coords, _, u0 = initGrid(nVar, nX)
assert np.allclose(f2.header['coordX'], coords[0]), f"incorrect coordX in MPI written fields {f2}"

times = np.arange(nSteps) / nSteps
for idx, t in enumerate(times):
Expand Down Expand Up @@ -342,8 +342,8 @@ def testCart2D_MPI(nProcs, dtypeIdx, algo, nSteps, nVar, nX, nY):
assert f2.nY == nY, f"incorrect nY in MPI written fields {f2}"

grids, _, u0 = initGrid(nVar, nX, nY)
assert np.allclose(f2.header['gridX'], grids[0]), f"incorrect gridX in MPI written fields {f2}"
assert np.allclose(f2.header['gridY'], grids[1]), f"incorrect gridY in MPI written fields {f2}"
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):
Expand Down

0 comments on commit 450f6be

Please sign in to comment.