Skip to content

Commit

Permalink
gpt_io can now use zip files! sparse_propagator official support
Browse files Browse the repository at this point in the history
  • Loading branch information
lehner committed Jan 1, 2025
1 parent 88d563c commit 94a7f17
Showing 7 changed files with 123 additions and 8 deletions.
1 change: 1 addition & 0 deletions lib/gpt/core/__init__.py
Original file line number Diff line number Diff line change
@@ -73,6 +73,7 @@
format,
mview,
FILE,
FILE_exists,
LoadError,
gpt_io,
corr_io,
46 changes: 43 additions & 3 deletions lib/gpt/core/io/FILE.py
Original file line number Diff line number Diff line change
@@ -16,7 +16,7 @@
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
import cgpt, gpt, os, shutil
import cgpt, gpt, os, shutil, zipfile


def cache_file(root, src, md):
@@ -36,7 +36,30 @@ def cache_file(root, src, md):
cache_root = gpt.default.get("--cache-root", None)


class FILE:
def zip_split(fn):
# check if any of the directories is a .zip archive
fn_split = fn.split("/")
fn_partial = fn_split[0]
for i, fn_part in enumerate(fn_split[1:]):
zip_candidate = fn_partial + ".zip"
if os.path.exists(zip_candidate):
return zip_candidate, "/".join(fn_split[i:])
fn_partial = f"{fn_partial}/{fn_part}"
return None, None


def FILE_exists(fn):
if os.path.exists(fn):
return True

fn_zip, fn_element = zip_split(fn)
if fn_zip is not None:
return True

return False


class FILE_base:
def __init__(self, fn, md):
if cache_root is not None:
fn = cache_file(cache_root, fn, md)
@@ -64,7 +87,14 @@ def seek(self, offset, whence):
r = cgpt.fseek(self.f, offset, whence)
return r

def read(self, sz):
def read(self, sz=None):
if sz is None:
pos = self.tell()
self.seek(0, 2)
size = self.tell()
self.seek(pos, 0)
return self.read(size - pos)

assert self.f is not None
t = bytes(sz)
if sz > 0:
@@ -81,3 +111,13 @@ def write(self, d):
def flush(self):
assert self.f is not None
cgpt.fflush(self.f)


def FILE(fn, mode):
if mode[0] == "w" or os.path.exists(fn):
return FILE_base(fn, mode)

# if file does not exists but should be read, try zip route
fn_zip, fn_element = zip_split(fn)
zip_file = zipfile.ZipFile(fn_zip, "r")
return zip_file.open(fn_element, "r")
2 changes: 1 addition & 1 deletion lib/gpt/core/io/__init__.py
Original file line number Diff line number Diff line change
@@ -21,7 +21,7 @@
import gpt.core.io.cevec_io
import gpt.core.io.qlat_io
import gpt.core.io.nersc_io
from gpt.core.io.FILE import FILE
from gpt.core.io.FILE import FILE, FILE_exists
from gpt.core.io.error import LoadError
from gpt.core.io.util import mview, crc32
from gpt.core.io.load import load
8 changes: 4 additions & 4 deletions lib/gpt/core/io/gpt_io.py
Original file line number Diff line number Diff line change
@@ -53,8 +53,8 @@ def __init__(self, root, write, params):
self.params["paths"] = [p.translate(replace) for p in self.params["paths"]]

if gpt.rank() == 0:
os.makedirs(self.root, exist_ok=True)
if write:
os.makedirs(self.root, exist_ok=True)
self.glb = gpt.FILE(root + "/global", "wb")
for f in glob.glob("%s/??/*.field" % self.root):
os.unlink(f)
@@ -539,7 +539,7 @@ def save(filename, objs, params):

def load(filename, params):
# first check if this is right file format
if not (os.path.exists(filename + "/index.crc32") and os.path.exists(filename + "/global")):
if not (gpt.FILE_exists(filename + "/index.crc32") and gpt.FILE_exists(filename + "/global")):
raise NotImplementedError()

# timing
@@ -551,8 +551,8 @@ def load(filename, params):
gpt.message("Reading %s" % filename)

# read index
idx = open(filename + "/index", "rb").read()
crc_expected = int(open(filename + "/index.crc32", "rt").read(), 16)
idx = gpt.FILE(filename + "/index", "rb").read()
crc_expected = int(gpt.FILE(filename + "/index.crc32", "rt").read(), 16)
crc_computed = gpt.crc32(memoryview(idx))
assert crc_expected == crc_computed

1 change: 1 addition & 0 deletions lib/gpt/qcd/__init__.py
Original file line number Diff line number Diff line change
@@ -22,5 +22,6 @@
import gpt.qcd.scalar
import gpt.qcd.pseudofermion
import gpt.qcd.baryon
import gpt.qcd.sparse_propagator
from gpt.qcd.utils import ferm_to_prop, prop_to_ferm
from gpt.qcd.wick import wick
72 changes: 72 additions & 0 deletions tests/qcd/sparse_propagator.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
#!/usr/bin/env python3
#
# Authors: Christoph Lehner 2025
#
import gpt as g
import numpy as np
import os


# workdir
if "WORK_DIR" in os.environ:
work_dir = os.environ["WORK_DIR"]
else:
work_dir = "."

# request test files
files = ["30_combined.zip"]
for f in files:
g.repository.load(f"{work_dir}/{f}", f"gpt://tests/qcd/propagators/{f}")

cache_params = (16, 4)
quark = g.qcd.sparse_propagator.flavor("30_combined/light.e", *cache_params)

coordinates = quark.sink_domain.coordinates
nsrc = quark.source_domain.sampled_sites

# test cache optimized sampler
rng = g.random("test")
to_sample = [[rng.uniform_int(min=0, max=nsrc - 1) for i in range(3)] for j in range(20)]
sampled = []
for s in g.qcd.sparse_propagator.cache_optimized_sampler([quark], to_sample):
q0 = quark[s[0]]
q1 = quark[s[1]]
q2 = quark[s[2]]
sampled.append(s)

# check that sampled is only a permutation of to_sample
to_sample = [s[0] + 1000 * s[1] + 1000000 * s[2] for s in to_sample]
sample = [s[0] + 1000 * s[1] + 1000000 * s[2] for s in sampled]
for s in sample:
assert s in to_sample
to_sample.pop(to_sample.index(s))
assert to_sample == []

# now perform self-consistency checks of quark
for i_src in range(nsrc):

g.message(f"Self-consistency check for src={i_src}")
# get src-to-src propagator as numpy matrices
src_src = quark(np.arange(nsrc), i_src)

# get embedding in full lattice to test
full = quark.sink_domain.sdomain.promote(quark[i_src])

# test that src_src is consistent with quark[i_src] which embeds all into the sparse domain
for i_snk in range(nsrc):
eps2 = g.norm2(full[coordinates[i_snk]] - g.mspincolor(src_src[i_snk]))
assert eps2 == 0.0

# now remove elements
mask = quark.source_mask()
remove = [3, 8, 11]
sparsened = g(mask * quark[i_src, remove])

slice = quark.sink_domain.sdomain.slice(sparsened, 3)
for i_snk in range(nsrc):
if i_snk in remove:
continue
slice[coordinates[i_snk, 3]] -= g.mspincolor(src_src[i_snk])
for s in slice:
eps2 = g.norm2(s)
assert eps2 < 1e-13
1 change: 1 addition & 0 deletions tests/run
Original file line number Diff line number Diff line change
@@ -44,6 +44,7 @@ for f in \
${test_path}/random/simple.py \
${test_path}/create/create.py \
${test_path}/create/smear.py \
${test_path}/qcd/sparse_propagator.py \
${test_path}/qcd/split_chiral.py \
${test_path}/qcd/coarsen.py \
${test_path}/qcd/coarse_even_odd.py \

0 comments on commit 94a7f17

Please sign in to comment.