import taichi as ti import numpy as np from utils import Sinusoid # TODO: fix oob memory access for first and last equations """ NOTE: This uses globals 'passes' and 'n' to allow static for-loop unrolling! """ @ti.kernel def solve_tridiagonal( now: ti.template(), dt: float, ld: ti.template(), dd: ti.template(), ud: ti.template(), rhs: ti.template(), soln: ti.template(), ): for _ in ti.static(range(UPDATES_PER_KERNEL)): # TODO: decide if this is actually faster or not """ Copy current solution to rhs """ for i in range(soln.shape[0]): rhs[0, i] = soln[i] soln[i] = 0 ti.sync() update_boundaries(rhs, now[None]) """ Forward reduction """ for pass_idx in ti.static(range(passes-1)): eqns = int(0.5**(pass_idx+1)*(n+1))-1 # print(f"Beginning to update solutions for dataset {pass_idx+1}, which will contain {eqns} eqns once complete") for i in range(eqns): a0, b0, c0 = ld[pass_idx, i*2], dd[pass_idx, i*2], ud[pass_idx, i*2] a1, b1, c1 = ld[pass_idx, i*2 + 1], dd[pass_idx, i*2 + 1], ud[pass_idx, i*2 + 1] a2, b2, c2 = ld[pass_idx, i*2 + 2], dd[pass_idx, i*2 + 2], ud[pass_idx, i*2 + 2] beta = 1 alpha = -beta*a1/b0 gamma = -beta*c1/b2 rhs0 = alpha * rhs[pass_idx, i*2] rhs1 = beta * rhs[pass_idx, i*2+1] rhs2 = gamma * rhs[pass_idx, i*2+2] rhs_new = rhs0+rhs1+rhs2 a = alpha*a0 b = alpha*c0 + beta*b1 + gamma*a2 c = gamma*c2 ld[pass_idx+1, i] = a dd[pass_idx+1, i] = b ud[pass_idx+1, i] = c rhs[pass_idx+1, i] = rhs_new ti.sync() """ Backwards substitution """ center_idx = int((n-1)/2) soln[center_idx] = rhs[passes-1,0] / dd[passes-1,0] for pass_counter in ti.static(range(passes-1)): pass_idx = passes-2-pass_counter n_eqs_in_pass = 2**(pass_counter+2)-1 n_unknowns_in_pass = (n_eqs_in_pass+1)/2 hop_size = 2**pass_idx for unknown_counter in range(int(n_unknowns_in_pass)): reduced_mat_row_idx = unknown_counter*2 soln_idx = hop_size-1 + reduced_mat_row_idx*hop_size # TODO: Make sure that the zero padding on ld/ud successfully allows the following statement which skips conditional format for edges soln[soln_idx] = ( rhs[pass_idx, reduced_mat_row_idx] - ld[pass_idx, reduced_mat_row_idx]*soln[soln_idx-hop_size] - ud[pass_idx, reduced_mat_row_idx]*soln[soln_idx+hop_size] ) / dd[pass_idx,reduced_mat_row_idx] ti.sync() now[None] += dt @ti.kernel def assemble_tridiagonal(ld: ti.template(), dd: ti.template(), ud: ti.template(), s: float): for i in range(n): if i == 0: ld[0, i] = 0 # zero pad the ld which doesn't exist dd[0, i] = 1 # BC is given in rhs ud[0, i] = 0 elif i == n-1: ld[0, i] = 0 dd[0, i] = 1 # BC is given in rhs ud[0, i] = 0 # zero pad the ud which doesn't exist else: ld[0, i] = -s dd[0, i] = 1 + 2*s ud[0, i] = -s ti.sync() """ Initial Conditions """ impulse_types = { "null": 0, "gaussian": 1, "splitting_gaussian": 2, "impulse": 3, "pulse_width": 4, "tri": 5, "fundamental": 6, "sine": 7, "harmonic": 8, } IMPULSE_MODE = impulse_types["null"] @ti.func def gaussian(x_norm: float, mu: float, sigma: float) -> float: return 0.35/ti.sqrt(2*np.pi*(sigma**2)) * ti.exp((-(x_norm- mu)**2)/(2*(sigma**2))) @ti.func def pt_impulse(soln: ti.template(), loc: int, a: float): soln[loc] = a @ti.kernel def apply_initial_conditions(soln: ti.template(), impulse_type: int): mu = 0.5 sigma = 0.10 for i in range(n): if impulse_type != impulse_types["impulse"]: x_norm = i/(n-1) if impulse_type == impulse_types["gaussian"]: soln[i] = gaussian(x_norm, mu, sigma) ti.sync() if impulse_type == impulse_types["impulse"]: loc = int((n+1)/2) # center height = 0.35 pt_impulse(soln, loc, height) @ti.func def update_boundaries(rhs: ti.template(), t: float): rhs[0, 0] = daily_sine.z(t) + annual_sine.z(t) rhs[0, n-1] = 0 @ti.kernel def update_geo(soln: ti.template(), circ_pts: ti.template(), line_pts: ti.template()): for i in range(n): p_x = i / (n-1) p_y = soln[i] * Y_VIEW_SCALE + Y_VIEW_OFFSET circ_pts[i].x = p_x circ_pts[i].y = p_y if (i != n-1): line_pts[2*i].x = p_x line_pts[2*i+1].x = (i+1) / (n-1) line_pts[2*i].y = p_y line_pts[2*i+1].y = soln[i+1] * Y_VIEW_SCALE + Y_VIEW_OFFSET if __name__ == '__main__': import time """ Taichi Init """ # TODO: fix oob error on edges for final pass ti.init(arch=ti.cuda, default_fp=ti.f64) # ti.init(arch=ti.cuda, default_fp=ti.f64, kernel_profiler=True) """ Finite Difference Grid Setup """ n = 2**9 - 1 #2**9 = 512, Needs to be 2^k-1 for tridiag solver, TODO: automatically pad out numbers dx = 1 D = 100 L = (n-1)*dx # courant condition: 2D*Delta T/Delta x^2 <= 1 dt = 1 s = D*dt/(dx**2) IMPULSE_TYPE = impulse_types["gaussian"] now = ti.field(dtype=float, shape=()) now[None] = 0 UPDATES_PER_KERNEL = 1 """ Tridiag setup """ passes = int(np.log2(n+1)) print(f"Solving system will require {passes} forward passes and {passes} backward passes.") """ Set up the fields """ # TODO: use a better data structure so that the higher number of passes only need 2**(passes-pass_idx)-1 entries lower_diagonal = ti.field(dtype=float, shape=(passes, n)) upper_diagonal = ti.field(dtype=float, shape=(passes, n)) diagonal = ti.field(dtype=float, shape=(passes, n)) right_hand_side = ti.field(dtype=float, shape=(passes, n)) solution = ti.field(dtype=float, shape=(n,)) """ Geometry fields """ line_points = ti.Vector.field(2,dtype=ti.f32, shape=(2*(n-1),)) circ_points = ti.Vector.field(2,dtype=ti.f32, shape=(n,)) Y_VIEW_SCALE = 0.4 Y_VIEW_OFFSET = 0.5 """ Setup """ assemble_tridiagonal(lower_diagonal, diagonal, upper_diagonal, s) apply_initial_conditions(solution, IMPULSE_TYPE) update_geo(solution, circ_points, line_points) daily_sine = Sinusoid( amp=0.15, freq=0.005, bias=0 ) annual_sine = Sinusoid( amp=0.6, freq=0.00005, bias=0 ) """ Benchmark """ # it = 0 # solve_tridiagonal(now, dt, lower_diagonal, diagonal, upper_diagonal, right_hand_side, solution) # print("Starting benchmark!") # start = time.time() # while it<1000: # solve_tridiagonal(now, dt, lower_diagonal, diagonal, upper_diagonal, right_hand_side, solution) # it += 1 # end = time.time() # print(f"{1000*(end-start)/(it*UPDATES_PER_KERNEL):0.3f}ms") """ Render setup """ window = ti.ui.Window('Diffusion 1D Implicit (Parallel)', res = (1000,1000), pos = (150, 150)) canvas = window.get_canvas() while window.running: canvas.set_background_color(color=(0,0,0)) canvas.lines(line_points, 0.002, indices=None, color=(1,1,1)) update_geo(solution, circ_points, line_points) solve_tridiagonal(now, dt, lower_diagonal, diagonal, upper_diagonal, right_hand_side, solution) # run(now, dt, lower_diagonal, diagonal, upper_diagonal, right_hand_side, solution, circ_points, line_points) window.show() # ti.profiler.print_scoped_profiler_info() # ti.profiler.print_kernel_profiler_info() # The default mode: 'count'