From 9ff34c27d74a3312deda1219455568a7c065fa15 Mon Sep 17 00:00:00 2001
From: Justus Kuhlmann <82444481+jkuhl-uni@users.noreply.github.com>
Date: Mon, 6 Jan 2025 10:46:49 +0100
Subject: [PATCH] Fix/sfcf ensname (#253)

* correct strings in _get_rep_names, add option for rep_sep

* doc

* add test for rep name getters
---
 pyerrors/input/sfcf.py | 19 ++++++++++---------
 tests/sfcf_in_test.py  | 30 ++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/pyerrors/input/sfcf.py b/pyerrors/input/sfcf.py
index e9f2837e..0431788a 100644
--- a/pyerrors/input/sfcf.py
+++ b/pyerrors/input/sfcf.py
@@ -127,7 +127,8 @@ def read_sfcf_multi(path, prefix, name_list, quarks_list=['.*'], corr_type_list=
     check_configs: list[list[int]]
         list of list of supposed configs, eg. [range(1,1000)]
         for one replicum with 1000 configs
-
+    rep_string: str
+        Separator of ensemble name and replicum. Example: In "ensAr0", "r" would be the separator string.
     Returns
     -------
     result: dict[list[Obs]]
@@ -199,9 +200,9 @@ def read_sfcf_multi(path, prefix, name_list, quarks_list=['.*'], corr_type_list=
     else:
         ens_name = kwargs.get("ens_name")
         if not appended:
-            new_names = _get_rep_names(ls, ens_name)
+            new_names = _get_rep_names(ls, ens_name, rep_sep=(kwargs.get('rep_string', 'r')))
         else:
-            new_names = _get_appended_rep_names(ls, prefix, name_list[0], ens_name)
+            new_names = _get_appended_rep_names(ls, prefix, name_list[0], ens_name, rep_sep=(kwargs.get('rep_string', 'r')))
         new_names = sort_names(new_names)
 
     idl = []
@@ -646,22 +647,22 @@ def _read_append_rep(filename, pattern, b2b, cfg_separator, im, single):
         return T, rep_idl, data
 
 
-def _get_rep_names(ls, ens_name=None):
+def _get_rep_names(ls, ens_name=None, rep_sep='r'):
     new_names = []
     for entry in ls:
         try:
-            idx = entry.index('r')
+            idx = entry.index(rep_sep)
         except Exception:
             raise Exception("Automatic recognition of replicum failed, please enter the key word 'names'.")
 
         if ens_name:
-            new_names.append('ens_name' + '|' + entry[idx:])
+            new_names.append(ens_name + '|' + entry[idx:])
         else:
             new_names.append(entry[:idx] + '|' + entry[idx:])
     return new_names
 
 
-def _get_appended_rep_names(ls, prefix, name, ens_name=None):
+def _get_appended_rep_names(ls, prefix, name, ens_name=None, rep_sep='r'):
     new_names = []
     for exc in ls:
         if not fnmatch.fnmatch(exc, prefix + '*.' + name):
@@ -670,12 +671,12 @@ def _get_appended_rep_names(ls, prefix, name, ens_name=None):
     for entry in ls:
         myentry = entry[:-len(name) - 1]
         try:
-            idx = myentry.index('r')
+            idx = myentry.index(rep_sep)
         except Exception:
             raise Exception("Automatic recognition of replicum failed, please enter the key word 'names'.")
 
         if ens_name:
-            new_names.append('ens_name' + '|' + entry[idx:])
+            new_names.append(ens_name + '|' + entry[idx:])
         else:
             new_names.append(myentry[:idx] + '|' + myentry[idx:])
     return new_names
diff --git a/tests/sfcf_in_test.py b/tests/sfcf_in_test.py
index f92126f9..60a71433 100644
--- a/tests/sfcf_in_test.py
+++ b/tests/sfcf_in_test.py
@@ -387,3 +387,33 @@ def test_find_correlator():
     found_start, found_T = sfin._find_correlator(file, "2.0", "name      f_A\nquarks    lquark lquark\noffset    0\nwf        0", False, False)
     assert found_start == 21
     assert found_T == 3
+
+
+def test_get_rep_name():
+    names = ['data_r0', 'data_r1', 'data_r2']
+    new_names = sfin._get_rep_names(names)
+    assert len(new_names) == 3
+    assert new_names[0] == 'data_|r0'
+    assert new_names[1] == 'data_|r1'
+    assert new_names[2] == 'data_|r2'
+    names = ['data_q0', 'data_q1', 'data_q2']
+    new_names = sfin._get_rep_names(names, rep_sep='q')
+    assert len(new_names) == 3
+    assert new_names[0] == 'data_|q0'
+    assert new_names[1] == 'data_|q1'
+    assert new_names[2] == 'data_|q2'
+
+
+def test_get_appended_rep_name():
+    names = ['data_r0.f_1', 'data_r1.f_1', 'data_r2.f_1']
+    new_names = sfin._get_appended_rep_names(names, 'data', 'f_1')
+    assert len(new_names) == 3
+    assert new_names[0] == 'data_|r0'
+    assert new_names[1] == 'data_|r1'
+    assert new_names[2] == 'data_|r2'
+    names = ['data_q0.f_1', 'data_q1.f_1', 'data_q2.f_1']
+    new_names = sfin._get_appended_rep_names(names, 'data', 'f_1', rep_sep='q')
+    assert len(new_names) == 3
+    assert new_names[0] == 'data_|q0'
+    assert new_names[1] == 'data_|q1'
+    assert new_names[2] == 'data_|q2'