Coverage for pySDC/helpers/fieldsIO.py: 56%

1#!/usr/bin/env python3

2# -*- coding: utf-8 -*-

3"""

4Generic utility class to write and read cartesian grid field solutions into binary files.

5It implements the base file handler class :class:`FieldsIO`, that is specialized into :

7- :class:`Scalar` : for 0D fields (scalar) with a given number of variables

8- :class:`Rectilinear` : for fields on N-dimensional rectilinear grids

10While each file handler need to be setup with specific parameters (grid, ...),

11each written file can be read using the same interface implemented in the

12base abstract class.

14Example

15-------

16>>> import numpy as np

17>>> from pySDC.helpers.fieldsIO import Rectilinear

18>>>

19>>> # Write some fields in files

20>>> x = np.linspace(0, 1, 128)

21>>> y = np.linspace(0, 1, 64)

22>>> fOut = Rectilinear(np.float64, "file.pysdc")

23>>> fOut.setHeader(nVar=2, coords=[x, y])

24>>> fOut.initialize()

25>>> times = [0, 1, 2]

26>>> xGrid, yGrid = np.meshgrid(x, y, indexing="ij")

27>>> u0 = np.array([-1, 1]).reshape((-1, 1, 1))*xGrid*yGrid

28>>> # u0 has shape [2, nX, nY]

29>>> for t in times:

30>>> fOut.addField(t, t*u0)

31>>>

32>>> # Read the file using the generic interface

33>>> from pySDC.helpers.fieldsIO import FieldsIO

34>>> fIn = FieldsIO.fromFile("file.pysdc")

35>>> times = fIn.times

36>>> assert len(times) == fIn.nFields

37>>> tEnd, uEnd = fIn.readField(-1)

38>>> assert tEnd == times[-1]

40Notes

41-----

42🚀 :class:`Rectilinear` is compatible with a MPI-based cartesian decomposition.

43See :class:`pySDC.helpers.fieldsIO.writeFields_MPI` for an illustrative example.

45Warning

46-------

47To use MPI collective writing, you need to call first the class methods :class:`Rectilinear.initMPI` (cf their docstring).

48Also, `Rectilinear.setHeader` **must be given the global grids coordinates**, wether the code is run in parallel or not.

50> ⚠️ Also : this module can only be imported with **Python 3.11 or higher** !

51"""

52import os

53import numpy as np

54from typing import Type, TypeVar

55import logging

56import itertools

58T = TypeVar("T")

60try:

61 try:

62 import dolfin as df # noqa: F841 (for some reason, dolfin always needs to be imported before mpi4py)

63 except ImportError:

64 pass

65 from mpi4py import MPI

66except ImportError:

68 class MPI:

69 COMM_WORLD = None

70 Intracomm = T

73# Supported data types

74DTYPES = {

75 0: np.float64, # double precision

76 1: np.complex128,

77}

78try:

79 DTYPES.update(

80 {

81 2: np.float128, # quadruple precision

82 3: np.complex256,

83 }

84 )

85except AttributeError:

86 logging.getLogger('FieldsIO').debug('Warning: Quadruple precision not available on this machine')

87try:

88 DTYPES.update(

89 {

90 4: np.float32, # single precision

91 5: np.complex64,

92 }

93 )

94except AttributeError:

95 logging.getLogger('FieldsIO').debug('Warning: Single precision not available on this machine')

97DTYPES_AVAIL = {val: key for key, val in DTYPES.items()}

99# Header dtype

100H_DTYPE = np.int8

101T_DTYPE = np.float64

102

103

104class FieldsIO:

105 """Abstract IO file handler"""

106

107 STRUCTS = {} # supported structures, modified dynamically

108 sID = None # structure ID of the FieldsIO class, modified dynamically

109

110 tSize = T_DTYPE().itemsize

111

112 ALLOW_OVERWRITE = False

113

114 def __init__(self, dtype, fileName):

115 """

116 Parameters

117 ----------

118 dtype : np.dtype

119 The data type of the fields values.

120 fileName : str

121 File.

122 """

123 assert dtype in DTYPES_AVAIL, f"{dtype=} not available. Supported on this machine: {list(DTYPES_AVAIL.keys())}"

124 self.dtype = dtype

125 self.fileName = fileName

126 self.initialized = False

127

128 # Initialized by the setHeader abstract method

129 self.header = None

130 self.nItems = None # number of values (dof) stored into one field

131

132 def __str__(self):

133 return f"FieldsIO[{self.__class__.__name__}|{self.dtype.__name__}|file:{self.fileName}]<{hex(id(self))}>"

134

135 def __repr__(self):

136 return self.__str__()

137

138 @classmethod

139 def fromFile(cls, fileName):

140 """

141 Read a file storing fields, and return the `FieldsIO` of the appropriate

142 field type (structure).

143

144 Parameters

145 ----------

146 fileName : str

147 Name of the binary file.

148

149 Returns

150 -------

151 fieldsIO : :class:`FieldsIO`

152 The specialized `FieldsIO` adapted to the file.

153 """

154 assert os.path.isfile(fileName), f"not a file ({fileName})"

155 with open(fileName, "rb") as f:

156 STRUCT, DTYPE = np.fromfile(f, dtype=H_DTYPE, count=2)

157 fieldsIO: FieldsIO = cls.STRUCTS[STRUCT](DTYPES[DTYPE], fileName)

158 fieldsIO.readHeader(f)

159 fieldsIO.initialized = True

160 return fieldsIO

161

162 @property

163 def hBase(self) -> np.ndarray:

164 """Base header into numpy array format"""

165 return np.array([self.sID, DTYPES_AVAIL[self.dtype]], dtype=H_DTYPE)

166

167 @classmethod

168 def register(cls, sID):

169 """

170 Decorator used to register a new class FieldsIO specialized class

171

172 Parameters

173 ----------

174 sID : int

175 Unique identifyer for the file, used in the binary file.

176 Since it's encoded on a 8-bytes signed integer,

177 it must be between -128 and 127

178

179 Example

180 -------

181 >>> # New specialized FieldsIO class

182 >>> @FieldsIO.register(sID=31)

183 >>> class HexaMesh2D(FieldsIO):

184 >>> pass # ... implementation

185 """

186

187 def wrapper(registered: Type[T]) -> Type[T]:

188 assert (

189 sID not in cls.STRUCTS

190 ), f"struct ID already taken by {cls.STRUCTS[sID]}, cannot use it for {registered}"

191 cls.STRUCTS[sID] = registered

192 registered.sID = sID

193 return registered

194

195 return wrapper

196

197 def initialize(self):

198 """Initialize the file handler : create the file with header, removing any existing file with the same name"""

199 assert self.header is not None, "header must be set before initializing FieldsIO"

200 assert not self.initialized, "FieldsIO already initialized"

201

202 if not self.ALLOW_OVERWRITE:

203 assert not os.path.isfile(

204 self.fileName

205 ), "file already exists, use FieldsIO.ALLOW_OVERWRITE = True to allow overwriting"

206

207 with open(self.fileName, "w+b") as f:

208 self.hBase.tofile(f)

209 for array in self.hInfos:

210 array.tofile(f)

211 self.initialized = True

212

213 def setHeader(self, **params):

214 """(Abstract) Set the header before creating a new file to store the fields"""

215 raise NotImplementedError()

216

217 @property

218 def hInfos(self) -> list[np.ndarray]:

219 """(Abstract) Array representing the grid structure to be written in the binary file."""

220 raise NotImplementedError()

221

222 def readHeader(self, f):

223 """

224 (Abstract) Read the header from the file storing the fields.

225

226 Parameters

227 ----------

228 f : `_io.TextIOWrapper`

229 File to read the header from.

230 """

231 raise NotImplementedError()

232

233 @property

234 def hSize(self):

235 """Size of the full header (in bytes)"""

236 return self.hBase.nbytes + sum(hInfo.nbytes for hInfo in self.hInfos)

237

238 @property

239 def itemSize(self):

240 """Size of one field value (in bytes)"""

241 return self.dtype().itemsize

242

243 @property

244 def fSize(self):

245 """Full size of a field (in bytes)"""

246 return self.nItems * self.itemSize

247

248 @property

249 def fileSize(self):

250 """Current size of the file (in bytes)"""

251 return os.path.getsize(self.fileName)

252

253 def addField(self, time, field):

254 """

255 Append one field solution at the end of the file with one given time.

256

257 Parameters

258 ----------

259 time : float-like

260 The associated time of the field solution.

261 field : np.ndarray

262 The field values.

263 """

264 assert self.initialized, "cannot add field to a non initialized FieldsIO"

265 field = np.asarray(field)

266 assert field.dtype == self.dtype, f"expected {self.dtype} dtype, got {field.dtype}"

267 assert field.size == self.nItems, f"expected {self.nItems} values, got {field.size}"

268 with open(self.fileName, "ab") as f:

269 np.array(time, dtype=T_DTYPE).tofile(f)

270 field.tofile(f)

271

272 @property

273 def nFields(self):

274 """Number of fields currently stored in the binary file"""

275 return int((self.fileSize - self.hSize) // (self.tSize + self.fSize))

276

277 def formatIndex(self, idx):

278 """Utility method to format a fields index to a positional integer (negative starts from last field index, like python lists)"""

279 nFields = self.nFields

280 if idx < 0:

281 idx = nFields + idx

282 assert idx < nFields, f"cannot read index {idx} from {nFields} fields"

283 assert idx >= 0, f"cannot read index {idx-nFields} from {nFields} fields"

284 return idx

285

286 @property

287 def times(self):

288 """Vector of all times stored in the binary file"""

289 times = []

290 with open(self.fileName, "rb") as f:

291 f.seek(self.hSize)

292 for i in range(self.nFields):

293 t = np.fromfile(f, dtype=T_DTYPE, count=1, offset=0 if i == 0 else self.fSize)[0]

294 times.append(float(t))

295 return times

296

297 def time(self, idx):

298 """Time stored at a given field index"""

299 idx = self.formatIndex(idx)

300 offset = self.hSize + idx * (self.tSize + self.fSize)

301 with open(self.fileName, "rb") as f:

302 t = np.fromfile(f, dtype=T_DTYPE, count=1, offset=offset)[0]

303 return float(t)

304

305 def readField(self, idx):

306 """

307 Read one field stored in the binary file, corresponding to the given

308 time index.

309

310 Parameters

311 ----------

312 idx : int

313 Positional index of the field.

314

315 Returns

316 -------

317 t : float

318 Stored time for this field.

319 field : np.ndarray

320 Read fields in a numpy array.

321 """

322 idx = self.formatIndex(idx)

323 offset = self.hSize + idx * (self.tSize + self.fSize)

324 with open(self.fileName, "rb") as f:

325 f.seek(offset)

326 t = float(np.fromfile(f, dtype=T_DTYPE, count=1)[0])

327 field = np.fromfile(f, dtype=self.dtype, count=self.nItems)

328 self.reshape(field)

329 return t, field

330

331 def reshape(self, field):

332 """Eventually reshape the field to correspond to the grid structure"""

333 pass

334

335

336@FieldsIO.register(sID=0)

337class Scalar(FieldsIO):

338 """FieldsIO handler storing a given number of scalar"""

339

340 # -------------------------------------------------------------------------

341 # Overridden methods

342 # -------------------------------------------------------------------------

343 def setHeader(self, nVar):

344 """

345 Set the descriptive grid structure to be stored in the file header.

346

347 Parameters

348 ----------

349 nVar : int

350 Number of scalar variable stored.

351 """

352 self.header = {"nVar": int(nVar)}

353 self.nItems = self.nVar

354

355 @property

356 def hInfos(self):

357 """Array representing the grid structure to be written in the binary file."""

358 return [np.array([self.nVar], dtype=np.int64)]

359

360 def readHeader(self, f):

361 """

362 Read the header from the binary file storing the fields.

363

364 Parameters

365 ----------

366 f : `_io.TextIOWrapper`

367 File to read the header from.

368 """

369 (nVar,) = np.fromfile(f, dtype=np.int64, count=1)

370 self.setHeader(nVar)

371

372 # -------------------------------------------------------------------------

373 # Class specifics

374 # -------------------------------------------------------------------------

375 @property

376 def nVar(self):

377 """Number of variables in a fields, as described in the header"""

378 return self.header["nVar"]

379

380

381@FieldsIO.register(sID=1)

382class Rectilinear(Scalar):

383 """FieldsIO handler storing a given number of scalar variables on a N-dimensional rectilinear grid"""

384

385 @staticmethod

386 def setupCoords(*coords):

387 """Utility function to setup grids in multiple dimensions, given the keyword arguments"""

388 coords = [np.asarray(coord, dtype=np.float64) for coord in coords]

389 for axis, coord in enumerate(coords):

390 assert coord.ndim == 1, f"coord for {axis=} must be one dimensional"

391 return coords

392

393 # -------------------------------------------------------------------------

394 # Overridden methods

395 # -------------------------------------------------------------------------

396 def setHeader(self, nVar, coords):

397 """

398 Set the descriptive grid structure to be stored in the file header.

399

400 Parameters

401 ----------

402 nVar : int

403 Number of 1D variables stored.

404 coords : np.1darray or list[np.1darray]

405 The grid coordinates in each dimensions.

406

407 Note

408 ----

409 When used in MPI decomposition mode, all coordinate **must** be the global grid.

410 """

411 if not isinstance(coords, (tuple, list)):

412 coords = [coords]

413 coords = self.setupCoords(*coords)

414 self.header = {"nVar": int(nVar), "coords": coords}

415 self.nItems = nVar * self.nDoF

416

417 @property

418 def hInfos(self):

419 """Array representing the grid structure to be written in the binary file."""

420 return [np.array([self.nVar, self.dim, *self.gridSizes], dtype=np.int32)] + [

421 np.array(coord, dtype=np.float64) for coord in self.header["coords"]

422 ]

423

424 def readHeader(self, f):

425 """

426 Read the header from the binary file storing the fields.

427

428 Parameters

429 ----------

430 f : `_io.TextIOWrapper`

431 File to read the header from.

432 """

433 nVar, dim = np.fromfile(f, dtype=np.int32, count=2)

434 gridSizes = np.fromfile(f, dtype=np.int32, count=dim)

435 coords = [np.fromfile(f, dtype=np.float64, count=n) for n in gridSizes]

436 self.setHeader(nVar, coords)

437

438 def reshape(self, fields: np.ndarray):

439 """Reshape the fields to a N-d array (inplace operation)"""

440 fields.shape = (self.nVar, *self.gridSizes)

441

442 # -------------------------------------------------------------------------

443 # Class specifics

444 # -------------------------------------------------------------------------

445 @property

446 def gridSizes(self):

447 """Number of points in y direction"""

448 return [coord.size for coord in self.header["coords"]]

449

450 @property

451 def dim(self):

452 """Number of grid dimensions"""

453 return len(self.gridSizes)

454

455 @property

456 def nDoF(self):

457 """Number of degrees of freedom for one variable"""

458 return np.prod(self.gridSizes)

459

460 def toVTR(self, baseName, varNames, idxFormat="{:06d}"):

461 """

462 Convert all 3D fields stored in binary format (FieldsIO) into a list

463 of VTR files, that can be read later with Paraview or equivalent to

464 make videos.

465

466 Parameters

467 ----------

468 baseName : str

469 Base name of the VTR file.

470 varNames : list[str]

471 Variable names of the fields.

472 idxFormat : str, optional

473 Formatting string for the index of the VTR file.

474 The default is "{:06d}".

475

476 Example

477 -------

478 >>> # Suppose the FieldsIO object is already writen into outputs.pysdc

479 >>> import os

480 >>> from pySDC.utils.fieldsIO import Rectilinear

481 >>> os.makedirs("vtrFiles") # to store all VTR files into a subfolder

482 >>> Rectilinear.fromFile("outputs.pysdc").toVTR(

483 >>> baseName="vtrFiles/field", varNames=["u", "v", "w", "T", "p"])

484 """

485 assert self.dim == 3, "can only be used with 3D fields"

486 from pySDC.helpers.vtkIO import writeToVTR

487

488 template = f"{baseName}_{idxFormat}"

489 for i in range(self.nFields):

490 _, u = self.readField(i)

491 writeToVTR(template.format(i), u, self.header["coords"], varNames)

492

493 # -------------------------------------------------------------------------

494 # MPI-parallel implementation

495 # -------------------------------------------------------------------------

496 comm: MPI.Intracomm = None

497

498 @classmethod

499 def setupMPI(cls, comm: MPI.Intracomm, iLoc, nLoc):

500 """

501 Setup the MPI mode for the files IO, considering a decomposition

502 of the 1D grid into contiuous subintervals.

503

504 Parameters

505 ----------

506 comm : MPI.Intracomm

507 The space decomposition communicator.

508 iLoc : list[int]

509 Starting index of the local sub-domain in the global coordinates.

510 nLoc : list[int]

511 Number of points in the local sub-domain.

512 """

513 cls.comm = comm

514 cls.iLoc = iLoc

515 cls.nLoc = nLoc

516 cls.mpiFile = None

517

518 @property

519 def MPI_ON(self):

520 """Wether or not MPI is activated"""

521 if self.comm is None:

522 return False

523 return self.comm.Get_size() > 1

524

525 @property

526 def MPI_ROOT(self):

527 """Wether or not the process is MPI Root"""

528 if self.comm is None:

529 return True

530 return self.comm.Get_rank() == 0

531

532 def MPI_FILE_OPEN(self, mode):

533 """Open the binary file in MPI mode"""

534 amode = {

535 "r": MPI.MODE_RDONLY,

536 "a": MPI.MODE_WRONLY | MPI.MODE_APPEND,

537 }[mode]

538 self.mpiFile = MPI.File.Open(self.comm, self.fileName, amode)

539

540 def MPI_WRITE(self, data):

541 """Write data (np.ndarray) in the binary file in MPI mode, at the current file cursor position."""

542 self.mpiFile.Write(data)

543

544 def MPI_WRITE_AT(self, offset, data: np.ndarray):

545 """

546 Write data in the binary file in MPI mode, with a given offset

547 **relative to the beginning of the file**.

548

549 Parameters

550 ----------

551 offset : int

552 Offset to write at, relative to the beginning of the file, in bytes.

553 data : np.ndarray

554 Data to be written in the binary file.

555 """

556 self.mpiFile.Write_at(offset, data)

557

558 def MPI_READ_AT(self, offset, data):

559 """

560 Read data from the binary file in MPI mode, with a given offset

561 **relative to the beginning of the file**.

562

563 Parameters

564 ----------

565 offset : int

566 Offset to read at, relative to the beginning of the file, in bytes.

567 data : np.ndarray

568 Array on which to read the data from the binary file.

569 """

570 self.mpiFile.Read_at(offset, data)

571

572 def MPI_FILE_CLOSE(self):

573 """Close the binary file in MPI mode"""

574 self.mpiFile.Close()

575 self.mpiFile = None

576

577 def initialize(self):

578 """Initialize the binary file (write header) in MPI mode"""

579 if self.MPI_ROOT:

580 try:

581 super().initialize()

582 except AssertionError as e:

583 if self.MPI_ON:

584 print(f"{type(e)}: {e}")

585 self.comm.Abort()

586 else:

587 raise e

588

589 if self.MPI_ON:

590 self.comm.Barrier() # Important, should not be removed !

591 self.initialized = True

592

593 def addField(self, time, field):

594 """

595 Append one field solution at the end of the file with one given time,

596 possibly using MPI.

597

598 Parameters

599 ----------

600 time : float-like

601 The associated time of the field solution.

602 field : np.ndarray

603 The (local) field values.

604

605 Note

606 ----

607 If a MPI decomposition is used, field **must be** the local field values.

608 """

609 if not self.MPI_ON:

610 return super().addField(time, field)

611

612 assert self.initialized, "cannot add field to a non initialized FieldsIO"

613

614 field = np.asarray(field)

615 assert field.dtype == self.dtype, f"expected {self.dtype} dtype, got {field.dtype}"

616 assert field.shape == (

617 self.nVar,

618 *self.nLoc,

619 ), f"expected {(self.nVar, *self.nLoc)} shape, got {field.shape}"

620

621 offset0 = self.fileSize

622 self.MPI_FILE_OPEN(mode="a")

623 if self.MPI_ROOT:

624 self.MPI_WRITE(np.array(time, dtype=T_DTYPE))

625 offset0 += self.tSize

626

627 for (iVar, *iBeg) in itertools.product(range(self.nVar), *[range(n) for n in self.nLoc[:-1]]):

628 offset = offset0 + self.iPos(iVar, iBeg) * self.itemSize

629 self.MPI_WRITE_AT(offset, field[iVar, *iBeg])

630 self.MPI_FILE_CLOSE()

631

632 def iPos(self, iVar, iX):

633 iPos = iVar * self.nDoF

634 for axis in range(self.dim - 1):

635 iPos += (self.iLoc[axis] + iX[axis]) * np.prod(self.gridSizes[axis + 1 :])

636 iPos += self.iLoc[-1]

637 return iPos

638

639 def readField(self, idx):

640 """

641 Read one field stored in the binary file, corresponding to the given

642 time index, using MPI in the eventuality of space parallel decomposition.

643

644 Parameters

645 ----------

646 idx : int

647 Positional index of the field.

648

649 Returns

650 -------

651 t : float

652 Stored time for this field.

653 field : np.ndarray

654 Read (local) fields in a numpy array.

655

656 Note

657 ----

658 If a MPI decomposition is used, it reads and returns the local fields values only.

659 """

660 if not self.MPI_ON:

661 return super().readField(idx)

662

663 idx = self.formatIndex(idx)

664 offset0 = self.hSize + idx * (self.tSize + self.fSize)

665 with open(self.fileName, "rb") as f:

666 t = float(np.fromfile(f, dtype=T_DTYPE, count=1, offset=offset0)[0])

667 offset0 += self.tSize

668

669 field = np.empty((self.nVar, *self.nLoc), dtype=self.dtype)

670

671 self.MPI_FILE_OPEN(mode="r")

672 for (iVar, *iBeg) in itertools.product(range(self.nVar), *[range(n) for n in self.nLoc[:-1]]):

673 offset = offset0 + self.iPos(iVar, iBeg) * self.itemSize

674 self.MPI_READ_AT(offset, field[iVar, *iBeg])

675 self.MPI_FILE_CLOSE()

676

677 return t, field

678

679

680# -----------------------------------------------------------------------------------------------

681# Utility functions used for testing

682# -----------------------------------------------------------------------------------------------

683def initGrid(nVar, gridSizes):

684 dim = len(gridSizes)

685 coords = [np.linspace(0, 1, num=n, endpoint=False) for n in gridSizes]

686 s = [None] * dim

687 u0 = np.array(np.arange(nVar) + 1)[:, *s]

688 for x in np.meshgrid(*coords, indexing="ij"):

689 u0 = u0 * x

690 return coords, u0

691

692

693def writeFields_MPI(fileName, dtypeIdx, algo, nSteps, nVar, gridSizes):

694 coords, u0 = initGrid(nVar, gridSizes)

695

696 from mpi4py import MPI

697 from pySDC.helpers.blocks import BlockDecomposition

698 from pySDC.helpers.fieldsIO import Rectilinear

699

700 comm = MPI.COMM_WORLD

701 MPI_SIZE = comm.Get_size()

702 MPI_RANK = comm.Get_rank()

703

704 blocks = BlockDecomposition(MPI_SIZE, gridSizes, algo, MPI_RANK)

705

706 iLoc, nLoc = blocks.localBounds

707 Rectilinear.setupMPI(comm, iLoc, nLoc)

708 s = [slice(i, i + n) for i, n in zip(iLoc, nLoc)]

709 u0 = u0[:, *s]

710 print(MPI_RANK, u0.shape)

711

712 f1 = Rectilinear(DTYPES[dtypeIdx], fileName)

713 f1.setHeader(nVar=nVar, coords=coords)

714

715 u0 = np.asarray(u0, dtype=f1.dtype)

716 f1.initialize()

717

718 times = np.arange(nSteps) / nSteps

719 for t in times:

720 ut = (u0 * t).astype(f1.dtype)

721 f1.addField(t, ut)

722

723 return u0

724

725

726def compareFields_MPI(fileName, u0, nSteps):

727 from pySDC.helpers.fieldsIO import FieldsIO

728

729 f2 = FieldsIO.fromFile(fileName)

730

731 times = np.arange(nSteps) / nSteps

732 for idx, t in enumerate(times):

733 u1 = u0 * t

734 t2, u2 = f2.readField(idx)

735 assert t2 == t, f"fields[{idx}] in {f2} has incorrect time ({t2} instead of {t})"

736 assert u2.shape == u1.shape, f"{idx}'s fields in {f2} has incorrect shape"

737 assert np.allclose(u2, u1), f"{idx}'s fields in {f2} has incorrect values"