Divergence-free reconstruction of B field

scripts/divergenceless_reconstruction.py

@@ -0,0 +1,342 @@
+import pytools as pt
+import numpy as np
+from scipy.special import legendre
+
+P = tuple(map(lambda n: legendre(n), (0, 1, 2, 3, 4)))
+
+# y and z are simply cyclic rotations of this
+def interpolate_x(a, x, y, z):
+    B_x = 0
+    for i in range(5):
+        for j in range(4):
+            for k in range(4):
+                B_x += a[i][j][k] * P[i](x) * P[j](y) * P[k](z)
+    return B_x
+
+def interpolate_y(b, x, y, z):
+    return interpolate_x(b, y, z, x)
+
+def interpolate_z(c, x, y, z):
+    return interpolate_x(c, z, x, y)
+
+# Returns moments as [ijk, order, x, y, z]
+def solve_moments_from_B(fg_b):
+    x_moments = solve_x_moments(fg_b[:, :, :, 0])
+    y_moments = solve_y_moments(fg_b[:, :, :, 1])
+    z_moments = solve_z_moments(fg_b[:, :, :, 2])
+    return (x_moments, y_moments, z_moments)
+
+# Returns moments for x as [order, x, y, z]
+# With order (0, y, z, yy, yz, zz, yyy, yyz, yzz, zzz)
+# y and z as cyclic rotations
+def solve_x_moments(B_x):
+    x_moments = np.zeros((10,) + np.shape(B_x))
+    x_moments[0] = B_x
+    x_moments[1:3] = np.gradient(B_x)[1:3]
+    start = 3
+    for i in range(2):
+        j = i + 1
+        x_moments[start:start+3-j] = np.gradient(x_moments[1+i])[j:3]
+        start += (3-j)
+    for i in range(3):
+        j = i + 1 if i < 2 else 2
+        x_moments[start:start+3-j] = np.gradient(x_moments[3+i])[j:3]
+        start += (3-j)
+    return x_moments
+
+def solve_y_moments(B_y):
+    return np.transpose(solve_x_moments(np.transpose(B_y, (1, 2, 0))), (0, 3, 1, 2))
+
+def solve_z_moments(B_z):
+    return np.transpose(solve_x_moments(np.transpose(B_z, (2, 0, 1))), (0, 2, 3, 1))
+
+# Solves a, b, c components with a[x, y, z], b[y, z, x] and c[z, x, y] up to given order
+# Input B should be the output of solve_moments_from_B
+# Might be deprecated. A lot faster than calculating all coefficients but that's only a few seconds anyway for a 100^3 array
+def solve_coefficients(B_moments, xyz, order = 4):
+    abc = np.zeros((3, 5, 4, 4))
+    x = xyz[0]
+    y = xyz[1]
+    z = xyz[2]
+
+    # 4th order
+    if (order > 3):
+        for i in range(3):
+            j = (i+1) % 3
+            k = (i+2) % 3
+            coords = (x if i else x + 1, y if k else y + 1, z if j else z + 1)
+            a = abc[i]
+            Bx = B_moments[i]
+
+            a[0][3][0] = 1/2 * (Bx[6][xyz] + Bx[6][coords])
+            a[0][2][1] = 1/2 * (Bx[7][xyz] + Bx[7][coords])
+            a[0][1][2] = 1/2 * (Bx[8][xyz] + Bx[8][coords])
+            a[0][0][3] = 1/2 * (Bx[9][xyz] + Bx[9][coords])
+
+            a[1][3][0] = Bx[6][xyz] - Bx[6][coords]
+            a[1][2][1] = Bx[7][xyz] - Bx[7][coords]
+            a[1][1][2] = Bx[8][xyz] - Bx[8][coords]
+            a[1][0][3] = Bx[9][xyz] - Bx[9][coords]
+
+        for i in range(3):
+            j = (i+1) % 3
+            k = (i+2) % 3
+            a = abc[i]
+            b = abc[j]
+            c = abc[k]
+
+            # Should be correct, check later
+            a[4][0][0] = -1/4 * (b[1][0][3] + c[1][3][0])
+            a[3][1][0] = -7/30 * c[1][2][1]
+            a[3][0][1] = -7/30 * b[1][1][2]
+            a[2][2][0] = -3/20 * c[1][1][2]
+            a[2][0][2] = -3/20 * b[1][2][1]
+
+    # 3rd order
+    if (order > 2):
+        for i in range(3):
+            j = (i+1) % 3
+            k = (i+2) % 3
+            coords = (x if i else x + 1, y if k else y + 1, z if j else z + 1)
+            a = abc[i]
+            Bx = B_moments[i]
+
+            a[0][2][0] = 1/2 * (Bx[3][xyz] + Bx[3][coords]) - 1/6 * a[2][2][0]
+            a[0][1][1] = 1/2 * (Bx[4][xyz] + Bx[4][coords])
+            a[0][0][2] = 1/2 * (Bx[5][xyz] + Bx[5][coords]) - 1/6 * a[2][0][2]
+
+            a[1][2][0] = Bx[3][xyz] - Bx[3][coords]
+            a[1][1][1] = Bx[4][xyz] - Bx[4][coords]
+            a[1][0][2] = Bx[5][xyz] - Bx[5][coords]
+
+        for i in range(3):
+            j = (i+1) % 3
+            k = (i+2) % 3
+            a = abc[i]
+            b = abc[j]
+            c = abc[k]
+
+            # Should be correct, check later
+            a[3][0][0] = -1/3 * (b[1][0][2] + c[1][2][0])
+            a[2][1][0] = -1/4 * c[1][1][1]
+            a[2][0][1] = -1/4 * b[1][1][1]
+
+    # 2nd order
+    if (order > 1):
+        for i in range(3):
+            j = (i+1) % 3
+            k = (i+2) % 3
+            coords = (x if i else x + 1, y if k else y + 1, z if j else z + 1)
+            a = abc[i]
+            Bx = B_moments[i]
+
+            a[0][1][0] = 1/2 * (Bx[1][xyz] + Bx[1][coords]) - 1/6 * a[2][1][0]
+            a[0][0][1] = 1/2 * (Bx[2][xyz] + Bx[2][coords]) - 1/6 * a[2][0][1]
+
+            a[1][1][0] = (Bx[1][xyz] - Bx[1][coords]) - 1/10 * a[3][1][0]
+            a[1][0][1] = (Bx[2][xyz] - Bx[2][coords]) - 1/10 * a[3][0][1]
+
+        for i in range(3):
+            j = (i+1) % 3
+            k = (i+2) % 3
+            a = abc[i]
+            b = abc[j]
+            c = abc[k]
+
+            # Should be correct, check later
+            a[2][0][0] = -1/2 * (b[1][0][1] + c[1][1][0]) - 3/35 * a[4][0][0] - 1/20 * (b[3][0][1] + c[3][1][0])
+
+    # 1st order
+    if (order > 0):
+        for i in range(3):
+            j = (i+1) % 3
+            k = (i+2) % 3
+            coords = (x if i else x + 1, y if k else y + 1, z if j else z + 1)
+            a = abc[i]
+            Bx = B_moments[i]
+
+            a[1][0][0] = (Bx[0][xyz] - Bx[0][coords]) - 1/10 * a[3][0][0]
+
+    # 0th order
+    for i in range(3):
+        j = (i+1) % 3
+        k = (i+2) % 3
+        coords = (x if i else x + 1, y if k else y + 1, z if j else z + 1)
+        a = abc[i]
+        Bx = B_moments[i]
+
+        a[0][0][0] = 1/2 * (Bx[0, x, y, z] + Bx[0][coords]) - 1/6 * a[2][0][0] - 1/70 * a[4][0][0]
+
+    # Check constraint:
+    test = abc[0][1][0][0] + abc[1][1][0][0] + abc[2][1][0][0] + 1/10 * (abc[0][3][0][0] + abc[1][3][0][0] + abc[2][3][0][0])
+    #if abs(test) > 0:
+    #    print("Something went wrong, sum (17) is " + str(test))
+
+    return abc
+
+def neighboursum(a, idx):
+    second = a[:-1, :-1, :-1]
+    if idx == 0:
+        first = a[1:, :-1, :-1]
+    elif idx == 1:
+        first = a[:-1, 1:, :-1]
+    elif idx == 2:
+        first = a[:-1, :-1, 1:]
+    return first + second
+
+def neighbourdiff(a, idx):
+    second = a[:-1, :-1, :-1]
+    if idx == 0:
+        first = a[1:, :-1, :-1]
+    elif idx == 1:
+        first = a[:-1, 1:, :-1]
+    elif idx == 2:
+        first = a[:-1, :-1, 1:]
+    return first - second
+
+# Solves a, b, c components with a[x, y, z], b[y, z, x] and c[z, x, y] up to given order
+# Input B should be the output of solve_moments_from_B
+def solve_all_coefficients(B_moments, order = 4):
+    shp = np.shape(B_moments[0][0])
+    abc = np.zeros((3, 5, 4, 4, shp[0] - 1, shp[1] - 1, shp[2] - 1))
+
+    # 4th order
+    if (order > 3):
+        for i in range(3):
+            a = abc[i]
+            Bx = B_moments[i]
+
+            a[0][3][0] = 1/2 * neighboursum(Bx[6], i)
+            a[0][2][1] = 1/2 * neighboursum(Bx[7], i)
+            a[0][1][2] = 1/2 * neighboursum(Bx[8], i)
+            a[0][0][3] = 1/2 * neighboursum(Bx[9], i)
+
+            a[1][3][0] = neighbourdiff(Bx[6], i)
+            a[1][2][1] = neighbourdiff(Bx[7], i)
+            a[1][1][2] = neighbourdiff(Bx[8], i)
+            a[1][0][3] = neighbourdiff(Bx[9], i)
+
+        for i in range(3):
+            j = (i+1) % 3
+            k = (i+2) % 3
+            a = abc[i]
+            b = abc[j]
+            c = abc[k]
+
+            # Should be correct, check later
+            a[4][0][0] = -1/4 * (b[1][0][3] + c[1][3][0])
+            a[3][1][0] = -7/30 * c[1][2][1]
+            a[3][0][1] = -7/30 * b[1][1][2]
+            a[2][2][0] = -3/20 * c[1][1][2]
+            a[2][0][2] = -3/20 * b[1][2][1]
+
+    # 3rd order
+    if (order > 2):
+        for i in range(3):
+            a = abc[i]
+            Bx = B_moments[i]
+
+            a[0][2][0] = 1/2 * neighboursum(Bx[3], i) - 1/6 * a[2][2][0]
+            a[0][1][1] = 1/2 * neighboursum(Bx[4], i)
+            a[0][0][2] = 1/2 * neighboursum(Bx[5], i) - 1/6 * a[2][0][2]
+
+            a[1][2][0] = neighbourdiff(Bx[3], i)
+            a[1][1][1] = neighbourdiff(Bx[4], i)
+            a[1][0][2] = neighbourdiff(Bx[5], i)
+
+        for i in range(3):
+            j = (i+1) % 3
+            k = (i+2) % 3
+            a = abc[i]
+            b = abc[j]
+            c = abc[k]
+
+            # Should be correct, check later
+            a[3][0][0] = -1/3 * (b[1][0][2] + c[1][2][0])
+            a[2][1][0] = -1/4 * c[1][1][1]
+            a[2][0][1] = -1/4 * b[1][1][1]
+
+    # 2nd order
+
+    if (order > 1):
+        for i in range(3):
+            a = abc[i]
+            Bx = B_moments[i]
+
+            a[0][1][0] = 1/2 * neighboursum(Bx[1], i) - 1/6 * a[2][1][0]
+            a[0][0][1] = 1/2 * neighboursum(Bx[2], i) - 1/6 * a[2][0][1]
+
+            a[1][1][0] = neighbourdiff(Bx[1], i) - 1/10 * a[3][1][0]
+            a[1][0][1] = neighbourdiff(Bx[2], i) - 1/10 * a[3][0][1]
+
+        for i in range(3):
+            j = (i+1) % 3
+            k = (i+2) % 3
+            a = abc[i]
+            b = abc[j]
+            c = abc[k]
+
+            # Should be correct, check later
+            a[2][0][0] = -1/2 * (b[1][0][1] + c[1][1][0]) - 3/35 * a[4][0][0] - 1/20 * (b[3][0][1] + c[3][1][0])
+
+    # 1st order
+    if (order > 0):
+        for i in range(3):
+            a = abc[i]
+            Bx = B_moments[i]
+
+            a[1][0][0] = neighbourdiff(Bx[0], i) - 1/10 * a[3][0][0]
+
+    # 0th order
+    for i in range(3):
+        a = abc[i]
+        Bx = B_moments[i]
+
+        a[0][0][0] = 1/2 * neighboursum(Bx[0], i) - 1/6 * a[2][0][0] - 1/70 * a[4][0][0]
+
+    # Check constraint:
+    #test = abc[0][1][0][0] + abc[1][1][0][0] + abc[2][1][0][0] + 1/10 * (abc[0][3][0][0] + abc[1][3][0][0] + abc[2][3][0][0])
+    #print(np.amax(test))
+    #if abs(test) > 0:
+    #    print("Something went wrong, sum (17) is " + str(test))
+
+    # return np.pad(np.transpose(abc, (4, 5, 6, 0, 1, 2, 3)), [(0, 1), (0, 1), (0, 1), (0, 0), (0, 0), (0, 0), (0, 0)])
+    return np.pad(abc, [(0, 0), (0, 0), (0, 0), (0, 0), (0, 1), (0, 1), (0, 1)])
+
+def center_value(B_moments, xyz, order=4):
+    abc = solve_coefficients(B_moments, xyz, order)
+    return [interpolate_x(abc[0], 0, 0, 0), interpolate_y(abc[1], 0, 0, 0), interpolate_z(abc[2], 0, 0, 0)]
+
+def ave_B(B_moments, xyz, order=4):
+    abc = solve_coefficients(B_moments, xyz, order)
+    a = abc[0]
+    b = abc[1]
+    c = abc[2]
+    return (
+        a[0][0][0] - 3/8 * (a[2][0][0] + a[0][2][0] + a[0][0][2]) + 9/64 * (a[2][2][0] + a[2][0][2]) + 15/128 * a[4][0][0],
+        b[0][0][0] - 3/8 * (b[2][0][0] + b[0][2][0] + b[0][0][2]) + 9/64 * (b[2][2][0] + b[2][0][2]) + 15/128 * b[4][0][0],
+        c[0][0][0] - 3/8 * (c[2][0][0] + c[0][2][0] + c[0][0][2]) + 9/64 * (c[2][2][0] + c[2][0][2]) + 15/128 * c[4][0][0]
+    )
+
+def all_ave_Bs(B_moments, order=4):
+    abc = solve_all_coefficients(B_moments, order)
+    a = abc[0]
+    b = abc[1]
+    c = abc[2]
+    return np.transpose(np.array(
+        (
+            a[0][0][0] - 3/8 * (a[2][0][0] + a[0][2][0] + a[0][0][2]) + 9/64 * (a[2][2][0] + a[2][0][2]) + 15/128 * a[4][0][0],
+            b[0][0][0] - 3/8 * (b[2][0][0] + b[0][2][0] + b[0][0][2]) + 9/64 * (b[2][2][0] + b[2][0][2]) + 15/128 * b[4][0][0],
+            c[0][0][0] - 3/8 * (c[2][0][0] + c[0][2][0] + c[0][0][2]) + 9/64 * (c[2][2][0] + c[2][0][2]) + 15/128 * c[4][0][0]
+        )
+    ), [1, 2, 3, 0])
+
+def center_values(B_moments, coords, order=4):
+    # Looks scuffed but is faster
+    abc = np.transpose(np.transpose(solve_all_coefficients(B_moments, order), (4, 5, 6, 0, 1, 2, 3))[coords[:, 0], coords[:, 1], coords[:, 2]], (1, 2, 3, 4, 0))
+    return np.transpose(np.array([interpolate_x(abc[0], 0, 0, 0), interpolate_y(abc[1], 0, 0, 0), interpolate_z(abc[2], 0, 0, 0)]))
+    #return all_center_values(B_moments, order)[coords[:, 0], coords[:, 1], coords[:, 2], :]
+
+def all_center_values(B_moments, order=4):
+    abc = solve_all_coefficients(B_moments, order)
+    return np.transpose(np.array([interpolate_x(abc[0], 0, 0, 0), interpolate_y(abc[1], 0, 0, 0), interpolate_z(abc[2], 0, 0, 0)]), (1, 2, 3, 0))

scripts/divergenceless_reconstruction_test.py

@@ -0,0 +1,35 @@
+from divergenceless_reconstruction import solve_moments_from_B, center_value, center_values, all_center_values, all_ave_Bs
+import pytools as pt
+import numpy as np
+from time import perf_counter
+# Code for testing divergenceless_resconstruction.py
+
+f = pt.vlsvfile.VlsvReader('/wrk/users/lkotipal/TEST/bulk.0000057.vlsv')
+fg_b = f.read_fsgrid_variable('fg_b')
+print(np.shape(fg_b))
+fg_b_vol = f.read_fsgrid_variable('fg_b_vol')
+print('File read!')
+t = perf_counter()
+B_moments = solve_moments_from_B(fg_b[0:100, 0:100, 0:100])
+print(f'B_moments solved in {perf_counter() - t} seconds!')
+print()
+print(f'Volumetric b in (50, 50, 50): {fg_b_vol[50, 50, 50]}')
+print()
+for i in range(5):
+    t = perf_counter()
+    print(center_value(B_moments, (50, 50, 50), i))
+    print(f'Single coefficients up to order {i} solved in {perf_counter() - t} seconds!')
+
+    t = perf_counter()
+    print(np.shape(center_values(B_moments, np.array([(50, 50, 50), (51, 51, 51), (54, 54, 64), (43, 65, 42), (43, 65, 32)]), i)))
+    print(f'Multiple coefficients up to order {i} solved in {perf_counter() - t} seconds!')
+
+    t = perf_counter()
+    print(all_center_values(B_moments, i)[50, 50, 50])
+    print(f'All coefficients up to order {i} solved in {perf_counter() - t} seconds!')
+
+    print(f'Average scaled difference between Analysator and Vlasiator in order {i}:')
+    print(np.average(((all_ave_Bs(B_moments, i))[:-1, :-1, :-1] - fg_b_vol[:99, :99, :99])/fg_b_vol[:99,:99,:99]))
+    print(f'Max scaled difference between Analysator and Vlasiator in order {i}:')
+    print(np.max(abs(((all_ave_Bs(B_moments, i))[:-1, :-1, :-1] - fg_b_vol[:99, :99, :99])/fg_b_vol[:99,:99,:99])))
+    print()