diff --git a/src/gvec_to_python/geometry/domain.py b/src/gvec_to_python/geometry/domain.py
index 5034e2d25e772a6683a66e3b6b79387369739dc3..af8e258815962495026d3bdcde65970ec588e19c 100644
--- a/src/gvec_to_python/geometry/domain.py
+++ b/src/gvec_to_python/geometry/domain.py
@@ -1025,9 +1025,9 @@ class GVEC_domain:
for i in range(s.shape[0]):
for j in range(th2.shape[1]):
for k in range(ze.shape[2]):
- func = lambda th: th + self.LA_base.eval_stz(s[i, 0, 0], th, ze[0, 0, k], self.LA_coef) - th2[0, j, 0]
+ func = lambda th: th + self.LA_base.eval_stz(s[i, 0, 0], th, ze[0, 0, k], self.LA_coef) - th2[i, j, k]
fprime = lambda th: 1 + self.LA_base.eval_stz(s[i, 0, 0], th, ze[0, 0, k], self.LA_coef, der='th')
- th[i, j, k] = newton(func, th2[0, j, 0], fprime)
+ th[i, j, k] = newton(func, th2[i, j, k], fprime)
return s + 0*th, th, ze + 0*th
diff --git a/tests/test_gvec_domain.py b/tests/test_gvec_domain.py
index 6da13fc93b7e73c1c1c30c1db170c77b148fc802..be23ad088787219bf713ab41b21b3bc170590ef3 100644
--- a/tests/test_gvec_domain.py
+++ b/tests/test_gvec_domain.py
@@ -50,24 +50,25 @@ def test_output_formats(mapping, use_pyccel,gvec_file):
assert gvec.g(s, a1, a2).shape == (1, 1, 1, 3, 3)
assert gvec.g_inv(s, a1, a2).shape == (1, 1, 1, 3, 3)
- # test 1d array evaluation
- s = np.linspace(.01, 1, 3) # mappings are singular at s=0
- if 'unit' not in mapping:
- a1 = np.linspace(0, 2*np.pi, 4)
- a2 = np.linspace(0, 2*np.pi, 5)
- else:
- a1 = np.linspace(0, 1, 4)
- a2 = np.linspace(0, 1, 5)
-
- shp3 = (s.size, a1.size, a2.size)
- shp5 = (s.size, a1.size, a2.size, 3, 3)
-
- assert gvec.f(s, a1, a2)[0].shape == shp3
- assert gvec.df(s, a1, a2).shape == shp5
- assert gvec.det_df(s, a1, a2).shape == shp3
- assert gvec.df_inv(s, a1, a2).shape == shp5
- assert gvec.g(s, a1, a2).shape == shp5
- assert gvec.g_inv(s, a1, a2).shape == shp5
+ # test 1d array evaluation (all combinations)
+ for n0,n1,n2 in [(3,4,5),(3,4,1),(3,1,5),(1,4,5),(3,1,1),(1,4,1),(1,1,5),(1,1,1)]:
+ s = np.linspace(.01, 1, n0) # mappings are singular at s=0
+ if 'unit' not in mapping:
+ a1 = np.linspace(0, 2*np.pi, n1,endpoint=False)
+ a2 = np.linspace(0, 2*np.pi, n2,endpoint=False)
+ else:
+ a1 = np.linspace(0, 1, n1,endpoint=False)
+ a2 = np.linspace(0, 1, n2,endpoint=False)
+
+ shp3 = (s.size, a1.size, a2.size)
+ shp5 = (s.size, a1.size, a2.size, 3, 3)
+
+ assert gvec.f(s, a1, a2)[0].shape == shp3
+ assert gvec.df(s, a1, a2).shape == shp5
+ assert gvec.det_df(s, a1, a2).shape == shp3
+ assert gvec.df_inv(s, a1, a2).shape == shp5
+ assert gvec.g(s, a1, a2).shape == shp5
+ assert gvec.g_inv(s, a1, a2).shape == shp5
# test 3d array evaluation
s, a1, a2 = np.meshgrid(s, a1, a2, indexing='ij')
@@ -101,31 +102,48 @@ def test_values(use_pyccel,gvec_file,unit_tor_domain):
gvec = GVEC(json_file_out, use_pyccel=use_pyccel,unit_tor_domain=unit_tor_domain)
# compare f and f_pest
- s = np.array([.5])
- th = np.random.rand(10)*2*np.pi
- ze = np.array([np.pi])
- s, th, ze = gvec.domain.prepare_args(s, th, ze)
- th2 = th + gvec._lambda(s, th, ze)
-
- # check Pmap
- assert np.allclose(s, gvec.domain.Pmap(s, th2, ze)[0])
- assert np.allclose(th, gvec.domain.Pmap(s, th2, ze)[1])
- assert np.allclose(ze, gvec.domain.Pmap(s, th2, ze)[2])
-
- # compare f and f_pest
- assert np.allclose(gvec.domain.f_gvec(s, th, ze)[0], gvec.domain.f_pest(s, th2, ze)[0])
- assert np.allclose(gvec.domain.f_gvec(s, th, ze)[1], gvec.domain.f_pest(s, th2, ze)[1])
- assert np.allclose(gvec.domain.f_gvec(s, th, ze)[2], gvec.domain.f_pest(s, th2, ze)[2])
-
- # compare f and f_unit
- assert np.allclose(gvec.domain.f_gvec(s, th, ze)[0], gvec.domain.f_unit(s, th/(2*np.pi), ze/(2*np.pi*gvec.domain.tor_fraction))[0])
- assert np.allclose(gvec.domain.f_gvec(s, th, ze)[1], gvec.domain.f_unit(s, th/(2*np.pi), ze/(2*np.pi*gvec.domain.tor_fraction))[1])
- assert np.allclose(gvec.domain.f_gvec(s, th, ze)[2], gvec.domain.f_unit(s, th/(2*np.pi), ze/(2*np.pi*gvec.domain.tor_fraction))[2])
+ s0 = np.array([.52,.63])
+ th0 = np.random.rand(5)*2*np.pi
+ ze0 = np.random.rand(3)*np.pi
+ # all combinations of 1D arrays
+ for n0,n1,n2 in [(2,5,3),(2,5,1),(2,1,3),(1,5,3),(2,1,1),(1,5,1),(1,1,3),(1,1,1)]:
+ s = s0[0:n0]
+ th = th0[0:n1]
+ ze = ze0[0:n2]
+ s, th, ze = gvec.domain.prepare_args(s, th, ze)
+ th2 = th + gvec._lambda(s, th, ze)
+
+
+ # check Pmap
+ assert np.allclose(s, gvec.domain.Pmap(s, th2, ze)[0])
+ assert np.allclose(th, gvec.domain.Pmap(s, th2, ze)[1])
+ assert np.allclose(ze, gvec.domain.Pmap(s, th2, ze)[2])
+
+ # compare f and f_pest
+ assert np.allclose(gvec.domain.f_gvec(s, th, ze)[0], gvec.domain.f_pest(s, th2, ze)[0])
+ assert np.allclose(gvec.domain.f_gvec(s, th, ze)[1], gvec.domain.f_pest(s, th2, ze)[1])
+ assert np.allclose(gvec.domain.f_gvec(s, th, ze)[2], gvec.domain.f_pest(s, th2, ze)[2])
+
+ # compare f and f_unit
+ assert np.allclose(gvec.domain.f_gvec(s, th, ze)[0], gvec.domain.f_unit(s, th/(2*np.pi), ze/(2*np.pi*gvec.domain.tor_fraction))[0])
+ assert np.allclose(gvec.domain.f_gvec(s, th, ze)[1], gvec.domain.f_unit(s, th/(2*np.pi), ze/(2*np.pi*gvec.domain.tor_fraction))[1])
+ assert np.allclose(gvec.domain.f_gvec(s, th, ze)[2], gvec.domain.f_unit(s, th/(2*np.pi), ze/(2*np.pi*gvec.domain.tor_fraction))[2])
+
+
+ # compute theta with equidistant in PEST, compare f and f_pest again
+ s2=s0
+ ze2=ze0
+ th2=np.linspace(0.,1.,6)*2*np.pi
+ s,th,ze=gvec.domain.Pmap(s2, th2, ze2)
+
+ assert np.allclose(gvec.domain.f_gvec(s, th, ze)[0], gvec.domain.f_pest(s2, th2, ze2)[0])
+ assert np.allclose(gvec.domain.f_gvec(s, th, ze)[1], gvec.domain.f_pest(s2, th2, ze2)[1])
+ assert np.allclose(gvec.domain.f_gvec(s, th, ze)[2], gvec.domain.f_pest(s2, th2, ze2)[2])
# compare f_pest and f_unit_pest
- assert np.allclose(gvec.domain.f_pest(s, th2, ze)[0], gvec.domain.f_unit_pest(s, th2/(2*np.pi), ze/(2*np.pi*gvec.domain.tor_fraction))[0])
- assert np.allclose(gvec.domain.f_pest(s, th2, ze)[1], gvec.domain.f_unit_pest(s, th2/(2*np.pi), ze/(2*np.pi*gvec.domain.tor_fraction))[1])
- assert np.allclose(gvec.domain.f_pest(s, th2, ze)[2], gvec.domain.f_unit_pest(s, th2/(2*np.pi), ze/(2*np.pi*gvec.domain.tor_fraction))[2])
+ assert np.allclose(gvec.domain.f_pest(s2, th2, ze2)[0], gvec.domain.f_unit_pest(s2, th2/(2*np.pi), ze2/(2*np.pi*gvec.domain.tor_fraction))[0])
+ assert np.allclose(gvec.domain.f_pest(s2, th2, ze2)[1], gvec.domain.f_unit_pest(s2, th2/(2*np.pi), ze2/(2*np.pi*gvec.domain.tor_fraction))[1])
+ assert np.allclose(gvec.domain.f_pest(s2, th2, ze2)[2], gvec.domain.f_unit_pest(s2, th2/(2*np.pi), ze2/(2*np.pi*gvec.domain.tor_fraction))[2])
if __name__ == '__main__':