def test_initial_density():
nx = 64
nv = 1024
vmax = 6.0
dv = 2 * vmax / nv
f = step.initialize(nx, nv)
np.testing.assert_almost_equal(f[0, ].sum() * dv, np.ones(nx), decimal=3)
def test_initial_temperature():
nx = 64
nv = 1024
vmax = 6.0
dv = 2 * vmax / nv
v = np.linspace(-vmax + dv / 2.0, vmax - dv / 2.0, nv)
f = step.initialize(nx, nv)
np.testing.assert_almost_equal(
np.array([((f[ix, 1:-1] * v[1:-1]**2.0).sum() + 0.5 *
(f[ix, 0] * v[0] + f[ix, -1]) * v[-1]) * dv
for ix in range(nx)]),
np.ones(nx),
decimal=3,
)
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 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)