def test_lenard_bernstein_velocity_zero(): """ tests if the 1st moment of f is (approximately) 0 :return: """ vmax = 6.0 nv = 256 dv = 2 * vmax / nv v = np.linspace(-vmax + dv / 2.0, vmax - dv / 2.0, nv) nu = 5e-2 dt = 0.1 v0 = 1.0 leftside = lenard_bernstein.make_philharmonic_matrix(v, nv, nu, dt, dv, v0) f = np.exp(-((v - 0.25)**2.0) / 2.0 / v0) f = f / np.sum(f * dv) f_out = f.copy() for it in range(500): f_out = lenard_bernstein.take_collision_step(leftside, f_out) temp_in = np.sum(f * v) * dv temp_out = np.sum(f_out * v) * dv np.testing.assert_almost_equal(temp_out, 0., decimal=1)
def test_lenard_bernstein_density_conservation(): """ tests if the 0th moment of f is conserved :return: """ vmax = 6.0 nv = 512 dv = 2 * vmax / nv v = np.linspace(-vmax + dv / 2.0, vmax - dv / 2.0, nv) nu = 1e-3 dt = 0.1 v0 = 1.0 leftside = lenard_bernstein.make_philharmonic_matrix(v, nv, nu, dt, dv, v0) f = np.exp(-((v - 0.5)**2.0) / 2.0 / v0) f = f / np.sum(f * dv) f_out = f.copy() for it in range(32): f_out = lenard_bernstein.take_collision_step(leftside, f_out) temp_in = np.sum(f) * dv temp_out = np.sum(f_out) * dv np.testing.assert_almost_equal(temp_out, temp_in, decimal=3)
def test_lenard_bernstein_maxwellian_solution(): """ tests if df/dt = 0 if f = maxwellian :return: """ vmax = 6.0 nv = 256 dv = 2 * vmax / nv v = np.linspace(-vmax + dv / 2.0, vmax - dv / 2.0, nv) nu = 1e-3 dt = 0.1 v0 = 1.0 leftside = lenard_bernstein.make_philharmonic_matrix(v, nv, nu, dt, dv, v0) f = np.exp(-(v**2.0) / 2.0 / v0) / np.sum( np.exp(-(v**2.0) / 2.0 / v0) * dv) f_out = f.copy() for it in range(32): f_out = lenard_bernstein.take_collision_step(leftside, f_out) np.testing.assert_almost_equal(f, f_out, decimal=4)
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 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)