Skip to content

Commit

Permalink
feat: added support for addition and multiplication of complex numbers (
Browse files Browse the repository at this point in the history
#209)

to Corr objects.
  • Loading branch information
fjosw authored Jul 21, 2023
1 parent 01ef97f commit 7f8c2ce
Show file tree
Hide file tree
Showing 4 changed files with 21 additions and 2 deletions.
4 changes: 2 additions & 2 deletions pyerrors/correlators.py
Original file line number Diff line number Diff line change
Expand Up @@ -1075,7 +1075,7 @@ def __add__(self, y):
newcontent.append(self.content[t] + y.content[t])
return Corr(newcontent)

elif isinstance(y, (Obs, int, float, CObs)):
elif isinstance(y, (Obs, int, float, CObs, complex)):
newcontent = []
for t in range(self.T):
if _check_for_none(self, self.content[t]):
Expand Down Expand Up @@ -1103,7 +1103,7 @@ def __mul__(self, y):
newcontent.append(self.content[t] * y.content[t])
return Corr(newcontent)

elif isinstance(y, (Obs, int, float, CObs)):
elif isinstance(y, (Obs, int, float, CObs, complex)):
newcontent = []
for t in range(self.T):
if _check_for_none(self, self.content[t]):
Expand Down
2 changes: 2 additions & 0 deletions pyerrors/obs.py
Original file line number Diff line number Diff line change
Expand Up @@ -784,6 +784,8 @@ def __add__(self, y):
else:
if isinstance(y, np.ndarray):
return np.array([self + o for o in y])
elif isinstance(y, complex):
return CObs(self, 0) + y
elif y.__class__.__name__ in ['Corr', 'CObs']:
return NotImplemented
else:
Expand Down
10 changes: 10 additions & 0 deletions tests/correlators_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -749,3 +749,13 @@ def test_corr_item():
corr_mat = pe.Corr(np.array([[corr_aa, corr_ab], [corr_ab, corr_aa]]))
corr_mat.item(0, 0)
assert corr_mat[0].item(0, 1) == corr_mat.item(0, 1)[0]


def test_complex_add_and_mul():
o = pe.pseudo_Obs(1.0, 0.3, "my_r345sfg16£$%&$%^%$^$", samples=47)
co = pe.CObs(o, 0.341 * o)
for obs in [o, co]:
cc = pe.Corr([obs for _ in range(4)])
cc += 2j
cc = cc * 4j
cc.real + cc.imag
7 changes: 7 additions & 0 deletions tests/obs_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -1333,3 +1333,10 @@ def test_vec_gm():
cc = pe.Corr(obs)
pe.gm(cc, S=4.12)
assert np.all(np.vectorize(lambda x: x.S["qq"])(cc.content) == 4.12)

def test_complex_addition():
o = pe.pseudo_Obs(34.12, 1e-4, "testens")
r = o + 2j
assert r.real == o
r = r * 1j
assert r.imag == o

0 comments on commit 7f8c2ce

Please sign in to comment.