bugfix in Pmap (solve th+lamdba-th2=0),th2 was not correctly set if multiple s...

bugfix in Pmap (solve th+lamdba-th2=0),th2 was not correctly set if multiple s or zeta positions were evaluated. Extended the domain tests to check all combinations of 1d arrays of s,th,ze with multiple / single element

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
 

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__':