import taichi as ti import numpy as np import matplotlib.pyplot as plt """ Init Taichi Backend, f64 precision """ ti.init(arch=ti.cuda, default_fp=ti.f64) """ Enable and Disable Plots """ PLOT_EULER_VS_RK = True PLOT_ADAPTIVE_RK = True DEBUG_INDEX = 0 """ Consts: Algo Selection """ EULER_INDEX = 0 RK_INDEX = 1 RK_ADAPTIVE_INDEX = 2 N_ALGORITHMS = 2 """ Consts: Design of Experiments """ CYCLE_COUNT = 100 STOPPING_POINT = np.pi*CYCLE_COUNT N_STEP_SIZE_TRIALS = 200 STEP_SIZE_MIN = 0.01 STEP_SIZE_MAX = np.pi/12 #0.5 STEP_SIZE_TRIAL_INCREMENT = (STEP_SIZE_MAX - STEP_SIZE_MIN)/N_STEP_SIZE_TRIALS N_ADAPTIVE_TOLERANCE_TRIALS = 200 ADAPTIVE_TOLERANCE_MIN = 0.00000000001 ADAPTIVE_TOLERANCE_MAX = 0.01 ADAPTIVE_TOLERANCE_INCREMENT = (ADAPTIVE_TOLERANCE_MAX - ADAPTIVE_TOLERANCE_MIN)/N_ADAPTIVE_TOLERANCE_TRIALS ADAPTIVE_START_STEP = 0.5 USE_O5_ESTIMATE = True MAX_ITERATIONS_FIXED = int(STOPPING_POINT/STEP_SIZE_MIN)+1 # Mem Limited by this MAX_ITERATIONS_ADAPTIVE = 30000 MAX_ADAPTIVE_COVERGENCE_ITER = 100 """ Const: how many derivatives will be used """ N_DERIVATIVES = 2 """ Taichi Data Structures """ model_def = { "Y": ti.types.matrix(N_DERIVATIVES, 1, dtype=float), "step_size": float, "t": float, } step_def = { "step": int, "t": float, } config_def = { "tol": float, "step_size": float } output_def = { "step": int, "step_size": float, "err": float, "err_total": float, "err_slope": float, } StepValues = ti.Struct.field(step_def , shape=(N_STEP_SIZE_TRIALS, MAX_ITERATIONS_FIXED)) StepSizeTrialConfig = ti.Struct.field(config_def, shape=(N_ALGORITHMS, N_STEP_SIZE_TRIALS)) StepSizeTrialHistory = ti.Struct.field(model_def, shape=(N_ALGORITHMS, N_STEP_SIZE_TRIALS, MAX_ITERATIONS_FIXED)) StepSizeTrialOutputs = ti.Struct.field(output_def, shape=(N_ALGORITHMS, N_STEP_SIZE_TRIALS)) AdaptiveTrialConfig = ti.Struct.field(config_def, shape=(1, N_ADAPTIVE_TOLERANCE_TRIALS)) AdaptiveTrialHistory = ti.Struct.field(model_def, shape=(1, N_ADAPTIVE_TOLERANCE_TRIALS, MAX_ITERATIONS_ADAPTIVE)) AdaptiveTrialOutputs = ti.Struct.field(output_def, shape=(1, N_ADAPTIVE_TOLERANCE_TRIALS)) """ Initial Conditions """ y0 = 1.0 yprime_0 = 0.0 Y0 = ti.Matrix.cols([ [ y0, yprime_0 ] ]) """ System of ODEs """ M = ti.Matrix([ [0, 1], [-1, 0] ]) """ Field Initialization Kernels """ @ti.kernel def fill_state(state_history: ti.template()): for model_mode, trial_index, step in state_history: addr = (model_mode, trial_index, step) state_history[addr].step_size = 0 if step == 0: state_history[addr].Y = Y0 else: state_history[addr].Y = ti.Matrix.cols([ [ 0.0, 0.0 ] ]) @ti.kernel def fill_config(config: ti.template()): for model_mode, trial_index in config: addr = (model_mode, trial_index) config[addr].tol = float(trial_index)*ADAPTIVE_TOLERANCE_INCREMENT + ADAPTIVE_TOLERANCE_MIN config[addr].step_size = float(trial_index)*STEP_SIZE_TRIAL_INCREMENT + STEP_SIZE_MIN @ti.kernel def fill_steps(): for trial_index, step_index in StepValues: addr = (trial_index, step_index) step_size = float(trial_index)*STEP_SIZE_TRIAL_INCREMENT + STEP_SIZE_MIN StepValues[addr].t = float(step_index)*step_size StepValues[addr].step = step_index """ Parallelized Experiments """ @ti.kernel def compute(): for algorithm_index, trial_index in ti.ndrange(StepSizeTrialHistory.shape[0], StepSizeTrialHistory.shape[1]): assert algorithm_index < N_ALGORITHMS, "Step size comparator only currently supports EULER ad RUNGE_KUTTE approximation. The input matrix is too long in the first dimension." compute_approximation_with_caching(algorithm_index, M, StepSizeTrialConfig, StepSizeTrialHistory, StepSizeTrialOutputs, trial_index) @ti.kernel def compute_adaptive_trials(): for trial_index in range(N_ADAPTIVE_TOLERANCE_TRIALS): compute_approximation_with_caching(RK_ADAPTIVE_INDEX, M, AdaptiveTrialConfig, AdaptiveTrialHistory, AdaptiveTrialOutputs, trial_index) @ti.func def compute_dy_rk(dt: float, M: ti.template(), Y: ti.template()) -> ti.template(): # Compute K values k1 = dt * M @ Y k2 = dt * M @ (Y + 0.5 * k1) k3 = dt * M @ (Y + 0.5 * k2) k4 = dt * M @ (Y + k3) # Compute dY dY = (1/6)*k1 + (1/3)*k2 + (1/3)*k3 + (1/6)*k4 return dY @ti.func def compute_dy_euler(dt: float, M: ti.template(), Y: ti.template()) -> ti.template(): return dt * (M @ Y) @ti.func def compute_dy_halfstep_err(M: ti.template(), Y: ti.template(), dY: ti.template(), step_size) -> ti.template(): dY_halfstep = compute_dy_rk(step_size/2, M, Y) Y_half = Y + dY_halfstep dY_halfstep = compute_dy_rk(step_size/2, M, Y_half) Y_test = Y_half + dY_halfstep Y_predicted = Y + compute_dy_rk(step_size, M, Y) # or + dY diff = ti.abs(Y_test - Y_predicted)[0, 0] return diff, Y_test - Y_predicted @ti.func def compute_approximation_with_caching(model_mode: int, M: ti.template(), config: ti.template(), state_history: ti.template(), outputs: ti.template(), trial_index: int): root_addr = (model_mode, trial_index) if model_mode == RK_ADAPTIVE_INDEX: root_addr[0] = 0 # Time Locations step_index = 0 t = 0.0 step_size = config[root_addr].step_size if model_mode == RK_ADAPTIVE_INDEX: step_size = ADAPTIVE_START_STEP tol = config[root_addr].tol error = 0.0 while t <= STOPPING_POINT and step_index < state_history.shape[2] - 1: # Get the previous state known_state = state_history[root_addr, step_index] Y = known_state.Y # Update the step size to the previous step's if in adaptive mode if model_mode == RK_ADAPTIVE_INDEX: if step_index != 0: step_size = state_history[root_addr, step_index - 1].step_size # Compute dY (Vector of dy,dy', dy'' etc) dY = ti.select(model_mode == EULER_INDEX, compute_dy_euler(step_size, M, Y), compute_dy_rk(step_size, M, Y)) if model_mode == RK_ADAPTIVE_INDEX: convergence_iter = 0 # TODO: implement bias and mean/width controls for adaptive stepper random gen errHalfstep, errHalfstepY = compute_dy_halfstep_err(M, Y, dY, step_size) if errHalfstep > tol: while errHalfstep > tol: step_size *= ((0.5) + ti.random(float)* 0.1 ) dY = compute_dy_rk(step_size, M, Y) errHalfstep, errHalfstepY = compute_dy_halfstep_err(M, Y, dY, step_size) convergence_iter += 1 assert convergence_iter < MAX_ADAPTIVE_COVERGENCE_ITER, f"Failed to converge adaptive step sizer to a tolerance of {tol} after {convergence_iter} iterations" state_history[root_addr, step_index].step_size = step_size dY = dY # + errHalfstepY else: new_test_step_size = step_size new_test_step_size *= ((1.02) + ti.random(float)*0.2) dY_test_step = compute_dy_rk(new_test_step_size, M, Y) errHalfstep_test, errHalfstepY_test = compute_dy_halfstep_err(M, Y, dY, new_test_step_size) if (errHalfstep_test < tol): step_size = new_test_step_size errHalfstep = errHalfstep_test errHalfstepY = errHalfstepY_test dY = dY_test_step # + errHalfstepY state_history[root_addr, step_index].step_size = step_size if USE_O5_ESTIMATE: NextY = known_state.Y + dY Y2dt = errHalfstepY + NextY CorrectedY = Y2dt + (errHalfstepY)/15 dY = CorrectedY - known_state.Y # Store Resulting data step_index += 1 state_history[root_addr, step_index].Y += known_state.Y + dY # Update the time counter Step = StepValues[trial_index, step_index] if model_mode == RK_ADAPTIVE_INDEX: t += step_size state_history[root_addr, step_index].t = t outputs[root_addr].step_size += step_size StepValues[trial_index, step_index].t = t else: t = Step.t # Compute and store the error actual_value = ti.cos(t) error = actual_value - state_history[root_addr, step_index].Y[0, 0] outputs[root_addr].err_total += ti.abs(error) # Update the outputs outputs[root_addr].step = step_index outputs[root_addr].err = error outputs[root_addr].err_total /= (step_index+ 1) if model_mode == RK_ADAPTIVE_INDEX: outputs[root_addr].step_size /= (step_index+ 1) else: outputs[root_addr].step_size = state_history[root_addr, step_index].step_size actual_slope = -ti.sin(t) err_slope = actual_slope - state_history[root_addr, step_index].Y[1, 0] outputs[root_addr].err_slope = err_slope if __name__ == '__main__': """ Execute the model """ fill_config(StepSizeTrialConfig) fill_state(StepSizeTrialHistory) fill_steps() compute() """Perform Data Analysis""" if PLOT_EULER_VS_RK: ConfigData = StepSizeTrialConfig.to_numpy() OutputData = StepSizeTrialOutputs.to_numpy() StateData = StepSizeTrialHistory.to_numpy() StepData = StepValues.to_numpy() for model_mode in [EULER_INDEX, RK_INDEX]: step_size_series = ConfigData["step_size"][model_mode,:] fig_final_error = plt.figure() err_data = OutputData["err"][model_mode,:] plt.scatter(step_size_series, err_data, s=2) plt.plot(step_size_series, err_data, alpha=0.2) plt.title(f"{'Euler' if model_mode == EULER_INDEX else 'RK'} Final Error vs Step Size") plt.xlabel("Step Size (rads)") plt.ylabel("Final Error") fig_final_error_slope = plt.figure() err_slope_data = OutputData["err_slope"][model_mode,:] plt.scatter(step_size_series, err_slope_data, s=2) plt.plot(step_size_series, err_slope_data, alpha=0.2) plt.title(f"{'Euler' if model_mode == EULER_INDEX else 'RK'} Final Slope Error vs Step Size") plt.xlabel("Step Size (rads)") plt.ylabel("Final Slope Error") avg_err_data = OutputData["err_total"][model_mode, :] fig_avg_local_err = plt.figure() plt.scatter(step_size_series, avg_err_data, s=2) plt.plot(step_size_series, avg_err_data, alpha=0.2) plt.title(f"{'Euler' if model_mode == EULER_INDEX else 'RK'} Average Abs. Local Error vs Step Size") plt.xlabel("Step Size (rads)") plt.ylabel("Avg Abs. Local Error") if model_mode == EULER_INDEX: plt.yscale('log') t_plt = np.linspace(0, STOPPING_POINT, 2000) trial_index = 0 fig_approx = plt.figure() final_step = OutputData['step'][model_mode, trial_index] approximation_data = StateData["Y"][model_mode, trial_index, 0:final_step+1, 0] step_data = StepData["t"][trial_index, :][0:final_step+1] plt.plot(step_data / np.pi, approximation_data, linewidth=1.25, alpha=0.5) plt.scatter(step_data / np.pi, approximation_data, s=1.5, alpha=0.3) plt.plot(t_plt / np.pi, np.cos(t_plt), linewidth = 0.5, alpha=0.5) plt.title(f"{'Euler' if model_mode == EULER_INDEX else 'RK'} Approximation for dt = {ConfigData['step_size'][model_mode, trial_index]}") plt.xlabel("t/PI") plt.ylabel("y(t)") trial_index = int(N_STEP_SIZE_TRIALS / 2) fig_approx = plt.figure() final_step = OutputData['step'][model_mode, trial_index] approximation_data = StateData["Y"][model_mode, trial_index, 0:final_step+1, 0] step_data = StepData["t"][trial_index, :][0:final_step+1] plt.plot(step_data / np.pi, approximation_data, linewidth=1.25, alpha=0.5) plt.scatter(step_data / np.pi, approximation_data, s=1.5, alpha=0.3) plt.plot(t_plt / np.pi, np.cos(t_plt), linewidth = 0.5, alpha=0.5) plt.title(f"{'Euler' if model_mode == EULER_INDEX else 'RK'} Approximation for dt = {ConfigData['step_size'][model_mode, trial_index]}") plt.xlabel("t/PI") plt.ylabel("y(t)") trial_index =N_STEP_SIZE_TRIALS-1 fig_approx = plt.figure() final_step = OutputData['step'][model_mode, trial_index] approximation_data = StateData["Y"][model_mode, trial_index, 0:final_step+1, 0] step_data = StepData["t"][trial_index, :][0:final_step+1] plt.plot(step_data / np.pi, approximation_data, linewidth=1.25, alpha=0.5) plt.scatter(step_data / np.pi, approximation_data, s=1.5, alpha=0.3) plt.plot(t_plt / np.pi, np.cos(t_plt), linewidth = 0.5, alpha=0.5) plt.title(f"{'Euler' if model_mode == EULER_INDEX else 'RK'} Approximation for dt = {ConfigData['step_size'][model_mode, trial_index]}") plt.xlabel("t/PI") plt.ylabel("y(t)") """ Run the Adaptive Stepper Test # TODO: set up separate output and step tables so this can be called earlier """ fill_config(AdaptiveTrialConfig) fill_state(AdaptiveTrialHistory) # AdaptiveTrialHistory.step_size.fill(0.3) compute_adaptive_trials() """ Analyze and plot the adaptive time steppers """ if PLOT_ADAPTIVE_RK: AdaptiveTrialConfigData = AdaptiveTrialConfig.to_numpy() AdaptiveTrialsData = AdaptiveTrialHistory.to_numpy() AdaptiveOutputData = AdaptiveTrialOutputs.to_numpy() model_mode = 0 tol_series = AdaptiveTrialConfigData["tol"][model_mode, :] t_plt = np.linspace(0, STOPPING_POINT,2000) fig_approx = plt.figure() trial_index = 0 final_step = AdaptiveOutputData['step'][model_mode, trial_index] step_data = AdaptiveTrialsData["t"][model_mode, trial_index, 0:final_step+1] approximation_data = AdaptiveTrialsData["Y"][model_mode, trial_index, 0:final_step+1, 0] plt.plot(step_data / np.pi, approximation_data, linewidth=1, alpha=0.9) plt.scatter(step_data / np.pi, approximation_data, s=2) plt.plot(t_plt / np.pi, np.cos(t_plt), linewidth=0.5, alpha=0.9) plt.title(f"Adaptive Stepper, tolerance = {AdaptiveTrialConfigData['tol'][model_mode, trial_index]}") plt.xlabel("t/PI") plt.ylabel("y(t)") fig_approx = plt.figure() trial_index = int(N_ADAPTIVE_TOLERANCE_TRIALS/2) final_step = AdaptiveOutputData['step'][model_mode, trial_index] step_data = AdaptiveTrialsData["t"][model_mode, trial_index, 0:final_step+1] approximation_data = AdaptiveTrialsData["Y"][model_mode, trial_index, 0:final_step+1, 0] plt.plot(step_data / np.pi, approximation_data, linewidth=1, alpha=0.9) plt.scatter(step_data / np.pi, approximation_data, s=2) plt.plot(t_plt / np.pi, np.cos(t_plt), linewidth=0.5, alpha=0.9) plt.title(f"Adaptive Stepper, tolerance = {AdaptiveTrialConfigData['tol'][model_mode, trial_index]}") plt.xlabel("t/PI") plt.ylabel("y(t)") fig_approx = plt.figure() trial_index = N_ADAPTIVE_TOLERANCE_TRIALS - 1 final_step = AdaptiveOutputData['step'][model_mode, trial_index] step_data = AdaptiveTrialsData["t"][model_mode, trial_index, 0:final_step+1] approximation_data = AdaptiveTrialsData["Y"][model_mode, trial_index, 0:final_step+1, 0] plt.plot(step_data / np.pi, approximation_data, linewidth=1, alpha=0.9) plt.scatter(step_data / np.pi, approximation_data, s=2) plt.plot(t_plt / np.pi, np.cos(t_plt), linewidth=0.5, alpha=0.9) plt.title(f"Adaptive Stepper, tolerance = {AdaptiveTrialConfigData['tol'][model_mode, trial_index]}") plt.xlabel("t/PI") plt.ylabel("y(t)") fig_avg_step = plt.figure() avg_step_data = AdaptiveOutputData["step_size"][model_mode,:] plt.scatter(tol_series, avg_step_data, s=2) plt.plot(tol_series, avg_step_data, alpha=0.2) plt.title(f"Average step sizes vs Adaptive Stepper Tolerance") plt.xlabel("Tolerance (difference between Y(t+dt) and y(t+dt/2+dt/2))") plt.ylabel("Average step sizes (radians)") fig_avg_step_error = plt.figure() avg_step_error = AdaptiveOutputData["err_total"][model_mode,:] plt.scatter(tol_series, avg_step_error, s=2) plt.plot(tol_series, avg_step_error, alpha=0.2) plt.title(f"Average Abs. Local Error vs Adaptive Stepper Tolerance") plt.xlabel("Tolerance (difference between Y(t+dt) and y(t+dt/2+dt/2))") plt.ylabel("Average Abs. Local Error") fig_step_error_final = plt.figure() step_error_final = AdaptiveOutputData["err"][model_mode,:] plt.scatter(tol_series, step_error_final, s=2) plt.plot(tol_series, step_error_final, alpha=0.2) plt.title(f"Final Error vs Adaptive Stepper Tolerance") plt.xlabel("Tolerance (difference between Y(t+dt) and y(t+dt/2+dt/2))") plt.ylabel("Final Error") """ Show the plots """ plt.show()