first impl. handle reps with list of lists

fjosw · Jan 17, 2024 · f60e623 · f60e623
1 parent df1873b
commit f60e623
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 3 deletions.
diff --git a/pyerrors/input/sfcf.py b/pyerrors/input/sfcf.py
@@ -121,7 +121,7 @@ def read_sfcf_multi(path, prefix, name_list, quarks_list=['.*'], corr_type_list=
         String that separates the ensemble identifier from the configuration number (default 'n').
     replica: list
         list of replica to be read, default is all
-    files: list
+    files: list[list[int]]
         list of files to be read per replica, default is all.
         for non-compact output format, hand the folders to be read here.
     check_configs: list[list[int]]
@@ -236,6 +236,10 @@ def read_sfcf_multi(path, prefix, name_list, quarks_list=['.*'], corr_type_list=
             rep_path = path + '/' + item
             if "files" in kwargs:
                 files = kwargs.get("files")
+                if isinstance(files, list):
+                    if isinstance(files[0], list):
+                        files = files[i]
+
             else:
                 files = []
             sub_ls = _find_files(rep_path, prefix, compact, files)
@@ -248,7 +252,7 @@ def read_sfcf_multi(path, prefix, name_list, quarks_list=['.*'], corr_type_list=
                     else:
                         rep_idl.append(int(cfg[3:]))
                 except Exception:
-                    raise Exception("Couldn't parse idl from directroy, problem with file " + cfg)
+                    raise Exception("Couldn't parse idl from directory, problem with file " + cfg)
             rep_idl.sort()
             # maybe there is a better way to print the idls
             if not silent:
@@ -309,7 +313,7 @@ def read_sfcf_multi(path, prefix, name_list, quarks_list=['.*'], corr_type_list=
             w = specs[3]
             w2 = specs[4]
             if "files" in kwargs:
-                ls = kwargs.get("files")
+                name_ls = kwargs.get("files")
             else:
                 name_ls = ls
                 for exc in name_ls:

diff --git a/tests/sfcf_in_test.py b/tests/sfcf_in_test.py
@@ -50,6 +50,18 @@ def test_o_bi(tmp_path):
     assert f_A[2].value == -41.025094911185185
 
 
+def test_o_bi_files(tmp_path):
+    build_test_environment(str(tmp_path), "o", 10, 3)
+    f_A = sfin.read_sfcf(str(tmp_path) + "/data_o", "test", "f_A", quarks="lquark lquark", wf=0, version="2.0",
+                         files=[["cfg" + str(i) for i in range(1, 11, 2)], ["cfg" + str(i) for i in range(2, 11, 2)], ["cfg" + str(i) for i in range(1, 11, 2)]])
+    print(f_A)
+    assert len(f_A) == 3
+    assert list(f_A[0].shape.keys()) == ["test_|r0", "test_|r1", "test_|r2"]
+    assert f_A[0].value == 65.4711887279723
+    assert f_A[1].value == 1.0447210336915187
+    assert f_A[2].value == -41.025094911185185
+
+
 def test_o_bib(tmp_path):
     build_test_environment(str(tmp_path), "o", 5, 3)
     f_V0 = sfin.read_sfcf(str(tmp_path) + "/data_o", "test", "F_V0", quarks="lquark lquark", wf=0, wf2=0, version="2.0", corr_type="bib")
@@ -120,6 +132,18 @@ def test_c_bi(tmp_path):
     assert f_A[2].value == -41.025094911185185
 
 
+def test_c_bi_files(tmp_path):
+    build_test_environment(str(tmp_path), "c", 10, 3)
+    f_A = sfin.read_sfcf(str(tmp_path) + "/data_c", "data_c", "f_A", quarks="lquark lquark", wf=0, version="2.0c",
+                         files=[["data_c_r0_n" + str(i) for i in range(1, 11, 2)], ["data_c_r1_n" + str(i) for i in range(2, 11, 2)], ["data_c_r2_n" + str(i) for i in range(1, 11, 2)]])
+    print(f_A)
+    assert len(f_A) == 3
+    assert list(f_A[0].shape.keys()) == ["data_c_|r0", "data_c_|r1", "data_c_|r2"]
+    assert f_A[0].value == 65.4711887279723
+    assert f_A[1].value == 1.0447210336915187
+    assert f_A[2].value == -41.025094911185185
+
+
 def test_c_bib(tmp_path):
     build_test_environment(str(tmp_path), "c", 5, 3)
     f_V0 = sfin.read_sfcf(str(tmp_path) + "/data_c", "data_c", "F_V0", quarks="lquark lquark", wf=0, wf2=0, version="2.0c", corr_type="bib")
@@ -256,6 +280,17 @@ def test_a_bi(tmp_path):
     assert f_A[2].value == -41.025094911185185
 
 
+def test_a_bi_files(tmp_path):
+    build_test_environment(str(tmp_path), "a", 5, 3)
+    f_A = sfin.read_sfcf(str(tmp_path) + "/data_a", "data_a", "f_A", quarks="lquark lquark", wf=0, version="2.0a", files=["data_a_r0.f_A", "data_a_r1.f_A", "data_a_r2.f_A"])
+    print(f_A)
+    assert len(f_A) == 3
+    assert list(f_A[0].shape.keys()) == ["data_a_|r0", "data_a_|r1", "data_a_|r2"]
+    assert f_A[0].value == 65.4711887279723
+    assert f_A[1].value == 1.0447210336915187
+    assert f_A[2].value == -41.025094911185185
+
+
 def test_a_bib(tmp_path):
     build_test_environment(str(tmp_path), "a", 5, 3)
     f_V0 = sfin.read_sfcf(str(tmp_path) + "/data_a", "data_a", "F_V0", quarks="lquark lquark", wf=0, wf2=0, version="2.0a", corr_type="bib")