diff --git a/pyquda/utils/io/__init__.py b/pyquda/utils/io/__init__.py index 7b59127..524ea28 100644 --- a/pyquda/utils/io/__init__.py +++ b/pyquda/utils/io/__init__.py @@ -10,7 +10,9 @@ readChromaQIO as readQIOPropagator, readChromaQIO as readChromaQIOPropagator, readMILCQIO as readMILCQIOPropagator, - readKYU as readKYUPropagator, - writeKYU as writeKYUPropagator, +) +from .fermion import ( + readKYU as readKYUFermion, + writeKYU as writeKYUFermion, ) from .eigen import readTimeSlice as readTimeSliceEivenvector diff --git a/pyquda/utils/io/fermion.py b/pyquda/utils/io/fermion.py new file mode 100644 index 0000000..d556ba5 --- /dev/null +++ b/pyquda/utils/io/fermion.py @@ -0,0 +1,97 @@ +from os import path + +import numpy + +from ...field import Ns, Nc, LatticeInfo, LatticeFermion, cb2 + + +def gatherFermionRaw(fermion_send: numpy.ndarray, latt_info: LatticeInfo): + from ... import getMPIComm, getCoordFromRank + + Gx, Gy, Gz, Gt = latt_info.grid_size + Lt, Lz, Ly, Lx = latt_info.size + dtype = fermion_send.dtype + + if latt_info.mpi_rank == 0: + fermion_recv = numpy.zeros((Gt * Gz * Gy * Gx, Lt, Lz, Ly, Lx, Ns, Nc), dtype) + getMPIComm().Gatherv(fermion_send, fermion_recv) + + fermion_raw = numpy.zeros((Gt * Lt, Gz * Lz, Gy * Ly, Gx * Lx, Ns, Nc), dtype) + for rank in range(latt_info.mpi_size): + gx, gy, gz, gt = getCoordFromRank(rank, [Gx, Gy, Gz, Gt]) + fermion_raw[ + gt * Lt : (gt + 1) * Lt, + gz * Lz : (gz + 1) * Lz, + gy * Ly : (gy + 1) * Ly, + gx * Lx : (gx + 1) * Lx, + ] = fermion_recv[rank] + else: + fermion_recv = None + getMPIComm().Gatherv(fermion_send, fermion_recv) + + fermion_raw = None + + return fermion_raw + + +def fromKYUBuffer(buffer: bytes, dtype: str, latt_info: LatticeInfo): + Gx, Gy, Gz, Gt = latt_info.grid_size + gx, gy, gz, gt = latt_info.grid_coord + Lx, Ly, Lz, Lt = latt_info.size + + fermion_raw = ( + numpy.frombuffer(buffer, dtype) + .reshape(2, Ns, Nc, Gt * Lt, Gz * Lz, Gy * Ly, Gx * Lx)[ + :, + :, + :, + gt * Lt : (gt + 1) * Lt, + gz * Lz : (gz + 1) * Lz, + gy * Ly : (gy + 1) * Ly, + gx * Lx : (gx + 1) * Lx, + ] + .astype("f8") + .tobytes() + ) + else: + buffer = None + + return buffer + + +def readKYU(filename: str, latt_info: LatticeInfo): + filename = path.expanduser(path.expandvars(filename)) + with open(filename, "rb") as f: + # kyu_binary_data = f.read(2 * Ns * Nc * Lt * Lz * Ly * Lx * 8) + kyu_binary_data = f.read() + fermion_raw = fromKYUBuffer(kyu_binary_data, ">f8", latt_info) + + return LatticeFermion(latt_info, cb2(fermion_raw, [0, 1, 2, 3])) + + +def writeKYU(filename: str, fermion: LatticeFermion): + filename = path.expanduser(path.expandvars(filename)) + latt_info = fermion.latt_info + kyu_binary_data = toKYUBuffer(fermion.lexico(), latt_info) + if latt_info.mpi_rank == 0: + with open(filename, "wb") as f: + f.write(kyu_binary_data) diff --git a/pyquda/utils/io/propagator.py b/pyquda/utils/io/propagator.py index fd5fccc..85b46d2 100644 --- a/pyquda/utils/io/propagator.py +++ b/pyquda/utils/io/propagator.py @@ -11,35 +11,6 @@ _precision_map = {"D": 8, "S": 4} -def gatherPropagatorRaw(propagator_send: numpy.ndarray, latt_info: LatticeInfo): - from ... import getMPIComm, getCoordFromRank - - Gx, Gy, Gz, Gt = latt_info.grid_size - Lt, Lz, Ly, Lx = latt_info.size - dtype = propagator_send.dtype - - if latt_info.mpi_rank == 0: - propagator_recv = numpy.zeros((Gt * Gz * Gy * Gx, Lt, Lz, Ly, Lx, Ns, Nc), dtype) - getMPIComm().Gatherv(propagator_send, propagator_recv) - - propagator_raw = numpy.zeros((Gt * Lt, Gz * Lz, Gy * Ly, Gx * Lx, Ns, Nc), dtype) - for rank in range(latt_info.mpi_size): - gx, gy, gz, gt = getCoordFromRank(rank, [Gx, Gy, Gz, Gt]) - propagator_raw[ - gt * Lt : (gt + 1) * Lt, - gz * Lz : (gz + 1) * Lz, - gy * Ly : (gy + 1) * Ly, - gx * Lx : (gx + 1) * Lx, - ] = propagator_recv[rank] - else: - propagator_recv = None - getMPIComm().Gatherv(propagator_send, propagator_recv) - - propagator_raw = None - - return propagator_raw - - def fromSCIDACBuffer(buffer: bytes, dtype: str, latt_info: LatticeInfo, staggered: bool): Gx, Gy, Gz, Gt = latt_info.grid_size gx, gy, gz, gt = latt_info.grid_coord @@ -110,50 +81,6 @@ def fromMultiSCIDACBuffer(buffer: bytes, dtype: str, latt_info: LatticeInfo, sta return propagator_raw -def fromKYUBuffer(buffer: bytes, dtype: str, latt_info: LatticeInfo): - Gx, Gy, Gz, Gt = latt_info.grid_size - gx, gy, gz, gt = latt_info.grid_coord - Lx, Ly, Lz, Lt = latt_info.size - - propagator_raw = ( - numpy.frombuffer(buffer, dtype) - .reshape(2, Ns, Nc, Gt * Lt, Gz * Lz, Gy * Ly, Gx * Lx)[ - :, - :, - :, - gt * Lt : (gt + 1) * Lt, - gz * Lz : (gz + 1) * Lz, - gy * Ly : (gy + 1) * Ly, - gx * Lx : (gx + 1) * Lx, - ] - .astype("f8") - .tobytes() - ) - else: - buffer = None - - return buffer - - def readChromaQIO(filename: str): filename = path.expanduser(path.expandvars(filename)) with open(filename, "rb") as f: @@ -250,22 +177,3 @@ def readMILCQIO(filename: str): return LatticePropagator(latt_info, cb2(propagator_raw, [0, 1, 2, 3])) else: return LatticeStaggeredPropagator(latt_info, cb2(propagator_raw, [0, 1, 2, 3])) - - -def readKYU(filename: str, latt_info: LatticeInfo): - filename = path.expanduser(path.expandvars(filename)) - with open(filename, "rb") as f: - # kyu_binary_data = f.read(2 * Ns * Nc * Lt * Lz * Ly * Lx * 8) - kyu_binary_data = f.read() - propagator_raw = fromKYUBuffer(kyu_binary_data, ">f8", latt_info) - - return LatticePropagator(latt_info, cb2(propagator_raw, [0, 1, 2, 3])) - - -def writeKYU(filename: str, propagator: LatticePropagator): - filename = path.expanduser(path.expandvars(filename)) - latt_info = propagator.latt_info - kyu_binary_data = toKYUBuffer(propagator.lexico(), latt_info) - if latt_info.mpi_rank == 0: - with open(filename, "wb") as f: - f.write(kyu_binary_data) diff --git a/pyquda/version.py b/pyquda/version.py index f6104e0..afd45a4 100644 --- a/pyquda/version.py +++ b/pyquda/version.py @@ -1 +1 @@ -__version__ = "0.6.7" +__version__ = "0.6.8"