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)