def full_leapfrog_ps_step(f, x, kx, v, kv, dv, t, dt, e, driver_function): """ Takes a step forward in time for f and e Uses leapfrog scheme 1 - spatial advection for 0.5 dt 2a - field solve 2b - velocity advection for dt 3 - spatial advection for 0.5 dt :param f: :param x: :param kx: :param v: :param kv: :param dv: :param t: :param dt: :param e: :param driver_function: :return: """ f = vlasov.update_velocity_adv_spectral(f, kv, e, 0.5 * dt) f = vlasov.update_spatial_adv_spectral(f, kx, v, dt) e = field.get_total_electric_field(driver_function(x, t + dt), f=f, dv=dv, kx=kx) f = vlasov.update_velocity_adv_spectral(f, kv, e, 0.5 * dt) return e, f
def start_run(all_params, pulse_dictionary, diagnostics, name="test", mlflow_path=None): """ This is the highest level function that calls the time integration loop MLFlow is initialized here. Domain configuration is also performed here. All file storage is initialized here. :param all_params: :param pulse_dictionary: :param diagnostics: :param name: :param mlflow_path: :return: """ if mlflow_path is not None: mlflow.set_tracking_uri(mlflow_path) mlflow.set_experiment(name) with mlflow.start_run(): # Log desired parameters params_to_log_dict = {} for param in diagnostics.params_to_log: if param in ["a0", "k0", "w0"]: params_to_log_dict[param] = pulse_dictionary["first pulse"][ param] else: params_to_log_dict[param] = all_params[param] mlflow.log_params(params_to_log_dict) # Initialize machinery nx = all_params["nx"] nv = all_params["nv"] nu = all_params["nu"] tmax = all_params["tmax"] nt = all_params["nt"] # Distribution function f = step.initialize(nx, nv) # Spatial Grid # Fixed to single wavenumber domains xmax = all_params["xmax"] xmin = all_params["xmin"] dx = (xmax - xmin) / nx x = np.linspace(xmin + dx / 2.0, xmax - dx / 2.0, nx) kx = np.fft.fftfreq(x.size, d=dx) * 2.0 * np.pi # Velocity grid vmax = all_params["vmax"] dv = 2 * vmax / nv v = np.linspace(-vmax + dv / 2.0, vmax - dv / 2.0, nv) kv = np.fft.fftfreq(v.size, d=dv) * 2.0 * np.pi t = np.linspace(0, tmax, nt) dt = t[1] - t[0] def driver_function(x, tt): total_field = np.zeros_like(x) for this_pulse in list(pulse_dictionary.keys()): kk = pulse_dictionary[this_pulse]["k0"] ww = pulse_dictionary[this_pulse]["w0"] envelope = get_pulse_coefficient( pulse_profile_dictionary=pulse_dictionary[this_pulse], tt=tt) if np.abs(envelope) > 0.0: total_field += envelope * np.cos(kk * x - ww * tt) return total_field e = field.get_total_electric_field(driver_function(x=x, tt=t[0]), f=f, dv=dv, kx=kx) # Storage temp_path = os.path.join(os.getcwd(), "temp-" + str(uuid.uuid4())[-6:]) os.makedirs(temp_path, exist_ok=True) storage_manager = storage.StorageManager(x, v, t, temp_path, store_f=diagnostics.f_rules) storage_manager.write_parameters_to_file(all_params, "all_parameters") storage_manager.write_parameters_to_file(pulse_dictionary, "pulses") # Matrix representing collision operator leftside = lenard_bernstein.make_philharmonic_matrix(vax=v, nv=nv, nu=nu, dt=dt, dv=dv, v0=1.0) # Time Loop for it in range(nt): e, f = step.full_PEFRL_ps_step(f, x, kx, v, kv, dv, t[it], dt, e, driver_function) if nu > 0.0: f = lenard_bernstein.take_collision_step( leftside=leftside, f=f, ) # All storage stuff here storage_manager.temp_update(tt=t[it], f=f, e=e, driver=driver_function(x=x, tt=t[it])) # Diagnostics diagnostics(storage_manager) # Log storage_manager.close() mlflow.log_artifacts(temp_path) # Cleanup shutil.rmtree(temp_path)
def full_PEFRL_ps_step(f, x, kx, v, kv, dv, t, dt, e, driver_function): """ Takes a step forward in time for f and e using the Performance-Extended Forest-Ruth-Like algorithm This is a 4th order symplectic integrator. http://physics.ucsc.edu/~peter/242/leapfrog.pdf :param f: :param x: :param kx: :param dx: :param v: :param kv: :param dv: :param t: :param dt: :param e: :param driver_function: :return: """ xsi = 0.1786178958448091 lambd = -0.2123418310626054 chi = -0.6626458266981849e-1 dt1 = xsi * dt dt2 = chi * dt dt3 = (1.0 - 2.0 * (chi + xsi)) * dt dt4 = dt2 dt5 = dt1 # x1 f = vlasov.update_spatial_adv_spectral(f, kx, v, dt1) e = field.get_total_electric_field(driver_function(x, t + dt1), f=f, dv=dv, kx=kx) # v1 f = vlasov.update_velocity_adv_spectral(f, kv, e, 0.5 * (1.0 - 2.0 * lambd) * dt) # x2 f = vlasov.update_spatial_adv_spectral(f, kx, v, dt2) e = field.get_total_electric_field(driver_function(x, t + dt1 + dt2), f=f, dv=dv, kx=kx) # v2 f = vlasov.update_velocity_adv_spectral(f, kv, e, lambd * dt) # x3 f = vlasov.update_spatial_adv_spectral(f, kx, v, dt3) e = field.get_total_electric_field(driver_function(x, t + dt1 + dt2 + dt3), f=f, dv=dv, kx=kx) # v3 f = vlasov.update_velocity_adv_spectral(f, kv, e, lambd * dt) # x4 f = vlasov.update_spatial_adv_spectral(f, kx, v, dt4) e = field.get_total_electric_field(driver_function( x, t + dt1 + dt2 + dt3 + dt4), f=f, dv=dv, kx=kx) # v4 f = vlasov.update_velocity_adv_spectral(f, kv, e, 0.5 * (1.0 - 2.0 * lambd) * dt) # x5 f = vlasov.update_spatial_adv_spectral(f, kx, v, dt5) e = field.get_total_electric_field(driver_function( x, t + dt1 + dt2 + dt3 + dt4 + dt5), f=f, dv=dv, kx=kx) return e, f
def test_full_leapfrog_ps_step_landau_damping(): nx = 64 nv = 512 f = step.initialize(nx, nv) k0 = 0.3 w_complex = get_roots_to_electrostatic_dispersion(1.0, 1.0, k0) w0 = np.real(w_complex) actual_decay_rate = np.imag(w_complex) xmax = 2 * np.pi / k0 xmin = 0.0 dx = (xmax - xmin) / nx x = np.linspace(xmin + dx / 2.0, xmax - dx / 2.0, nx) kx = np.fft.fftfreq(x.size, d=dx) * 2.0 * np.pi vmax = 6.0 dv = 2 * vmax / nv v = np.linspace(-vmax + dv / 2.0, vmax - dv / 2.0, nv) kv = np.fft.fftfreq(v.size, d=dv) * 2.0 * np.pi nt = 1000 tmax = 100 t = np.linspace(0, tmax, nt) dt = t[1] - t[0] def driver_function(x, t): """ t is 0D x is 1D """ a0 = 4e-4 envelope = np.exp(-((t - 8) ** 8.0) / 4.0 ** 8.0) return envelope * a0 * np.cos(k0 * x + w0 * t) field_store = np.zeros([nt, nx]) dist_store = np.zeros([nt, nx, nv]) t_store = np.zeros(nt) e = field.get_total_electric_field(driver_function(x, t[0]), f=f, dv=dv, kx=kx) it = 0 t_store[it] = t[it] dist_store[it] = f field_store[it] = e for it in range(1, nt): e, f = step.full_leapfrog_ps_step( f, x, kx, v, kv, dv, t[it], dt, e, driver_function ) t_store[it] = t[it] dist_store[it] = f field_store[it] = e t_ind = 600 ek = np.fft.fft(field_store, axis=1) ek_mag = np.array([np.abs(ek[it, 1]) for it in range(nt)]) decay_rate = np.mean(np.gradient(np.log(ek_mag[-t_ind:]), 0.1)) np.testing.assert_almost_equal(decay_rate, actual_decay_rate, decimal=2) ekw = np.fft.fft2(field_store[nt // 2 :,]) ek1w = np.abs(ekw[:, 1]) wax = np.fft.fftfreq(ek1w.shape[0], d=dt) * 2 * np.pi np.testing.assert_almost_equal(wax[ek1w.argmax()], w0, decimal=1)
def test_full_pefrl_ps_step_zero(): nx = 32 nv = 512 f = step.initialize(nx, nv) # f - defined # w0 = 1.1056 # k0 = 0.25 w0 = 1.1598 k0 = 0.3 # w0 = 1.2850 # k0 = 0.4 xmax = 2 * np.pi / k0 xmin = 0.0 dx = (xmax - xmin) / nx x = np.linspace(xmin + dx / 2.0, xmax - dx / 2.0, nx) kx = np.fft.fftfreq(x.size, d=dx) * 2.0 * np.pi vmax = 6.0 dv = 2 * vmax / nv v = np.linspace(-vmax + dv / 2.0, vmax - dv / 2.0, nv) kv = np.fft.fftfreq(v.size, d=dv) * 2.0 * np.pi nt = 400 tmax = 40.0 t = np.linspace(0, tmax, nt) dt = t[1] - t[0] def driver_function(x, t): """ t is 0D x is 1D """ a0 = 1e-6 envelope = np.exp(-((t - 8) ** 8.0) / 4.0 ** 8.0) return 0 * envelope * a0 * np.cos(k0 * x + w0 * t) field_store = np.zeros([nt, nx]) dist_store = np.zeros([nt, nx, nv]) t_store = np.zeros(nt) e = field.get_total_electric_field(driver_function(x, t[0]), f=f, dv=dv, kx=kx) it = 0 t_store[it] = t[it] dist_store[it] = f field_store[it] = e for it in range(1, nt): e, f = step.full_PEFRL_ps_step( f, x, kx, v, kv, dv, t[it], dt, e, driver_function ) t_store[it] = t[it] dist_store[it] = f field_store[it] = e np.testing.assert_almost_equal(field_store[0], field_store[-1], decimal=2) np.testing.assert_almost_equal(dist_store[0], dist_store[-1], decimal=2)
def start_run(nx, nv, nt, tmax, nu, w0, k0, a0, diagnostics, name="test", mlflow_path=None): """ End to end mlflow and xarray storage!! :param temp_path: :param nx: :param nv: :param nt: :param tmax: :param w0: :param k0: :param a0: :param name: :return: """ if mlflow_path is None: mlflow_client = mlflow.tracking.MlflowClient() else: mlflow_client = mlflow.tracking.MlflowClient(tracking_uri=mlflow_path) mlflow.set_experiment(name) with mlflow.start_run(): # Log initial conditions params_dict = { "nx": nx, "nv": nv, "w0": w0, "k0": k0, "nt": nt, "tmax": tmax, "a0": a0, "nu": nu, } mlflow.log_params(params_dict) # Initialize machinery # Distribution function f = step.initialize(nx, nv) # Spatial Grid # Fixed to single wavenumber domains xmax = 2 * np.pi / k0 xmin = 0.0 dx = (xmax - xmin) / nx x = np.linspace(xmin + dx / 2.0, xmax - dx / 2.0, nx) kx = np.fft.fftfreq(x.size, d=dx) * 2.0 * np.pi # Velocity grid vmax = 6.0 dv = 2 * vmax / nv v = np.linspace(-vmax + dv / 2.0, vmax - dv / 2.0, nv) kv = np.fft.fftfreq(v.size, d=dv) * 2.0 * np.pi t = np.linspace(0, tmax, nt) dt = t[1] - t[0] def driver_function(x, t): envelope = np.exp(-((t - 8)**8.0) / 4.0**8.0) return envelope * a0 * np.cos(k0 * x + w0 * t) e = field.get_total_electric_field(driver_function(x, t[0]), f=f, dv=dv, kx=kx) # Storage temp_path = os.path.join(os.getcwd(), "temp-" + str(uuid.uuid4())[-6:]) os.makedirs(temp_path, exist_ok=True) if nt // 4 < 100: t_store = 100 else: t_store = nt // 4 temp_field_store = np.zeros([t_store, nx]) temp_dist_store = np.zeros([t_store, nx, nv]) temp_t_store = np.zeros(t_store) it_store = 0 storage_manager = storage.StorageManager(x, v, t, temp_path) # Matrix representing collision operator A = lenard_bernstein.make_philharmonic_matrix(vax=v, nv=nv, nu=nu, dt=dt, dv=dv, v0=1.0) # Time Loop for it in range(nt): e, f = step.full_leapfrog_ps_step(f, x, kx, v, kv, dv, t[it], dt, e, driver_function) if nu > 0.0: for ix in range(nx): f[ix, ] = lenard_bernstein.take_collision_step(leftside=A, f=f[ix]) # All storage stuff here temp_t_store[it_store] = t[it] temp_dist_store[it_store] = f temp_field_store[it_store] = e it_store += 1 if it_store == t_store: storage_manager.batched_write_to_file(temp_t_store, temp_field_store, temp_dist_store) it_store = 0 # Diagnostics diagnostics(storage_manager) # Log storage_manager.close() mlflow.log_artifacts(temp_path) # Cleanup shutil.rmtree(temp_path)