def test_storage_init_files_exist():
xax = np.linspace(0, 1, 16)
vax = np.linspace(0, 1, 24)
tax = np.linspace(0, 1, 32)
dirname = os.path.join(os.getcwd(), str(uuid.uuid4()))
os.makedirs(dirname, exist_ok=True)
st = storage.StorageManager(xax=xax, vax=vax, tax=tax, base_path=dirname)
assert os.path.exists(os.path.join(dirname, "electric_field_vs_time.nc"))
assert os.path.exists(os.path.join(dirname, "dist_func_vs_time.nc"))
shutil.rmtree(dirname)
def test_storage_init_shape():
xax = np.linspace(0, 1, 16)
vax = np.linspace(0, 1, 24)
tax = np.linspace(0, 1, 32)
dirname = os.path.join(os.getcwd(), str(uuid.uuid4()))
os.makedirs(dirname, exist_ok=True)
st = storage.StorageManager(xax=xax, vax=vax, tax=tax, base_path=dirname)
np.testing.assert_equal(st.efield_arr.coords["space"].size, xax.size)
np.testing.assert_equal(st.efield_arr.coords["time"].size, tax.size)
np.testing.assert_equal(st.f_arr.coords["space"].size, xax.size)
np.testing.assert_equal(st.f_arr.coords["time"].size, tax.size)
np.testing.assert_equal(st.f_arr.coords["velocity"].size, vax.size)
shutil.rmtree(dirname)
def __initialize_base_storage_stuff__(td, xax, vax, rules_to_store_f=None):
if rules_to_store_f is None:
rules_to_store_f = {"space": "all", "time": "all"}
stuff_for_time_loop = {
"f": helpers.__initialize_f__(nx=xax.size, v=vax, v0=1.0, vshift=0.0),
"e": np.ones(xax.size),
"v": vax,
"x": xax,
"nx": xax.size,
"nv": vax.size,
"nu": 0.0,
"kv": vax,
"kx": xax,
"one_over_kx": xax,
"dt": 0.1,
"dv": 0.1,
"pulse_dictionary": {}
}
st = storage.StorageManager(
xax=xax,
vax=vax,
f=stuff_for_time_loop["f"],
base_path=td,
rules_to_store_f=rules_to_store_f,
all_params={},
pulse_dictionary={},
num_steps_in_one_loop=2,
)
sim_config = outer_loop.get_arrays_for_inner_loop(
stuff_for_time_loop=stuff_for_time_loop,
nt_in_loop=2,
store_f_rules=rules_to_store_f,
this_np=np,
)
st.batch_update(sim_config=sim_config)
return st
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(all_params, pulse_dictionary, diagnostics, uris, name="test"):
"""
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: (dictionary) contains the input parameters of the simulation
:param pulse_dictionary: (dictionary) contains the parameters for the ponderomotive force driver
:param diagnostics: (vlapy.Diagnostics) contains the diagnostics routine for this particular simulation
:param uris: (string) the location of the mlflow server
:param name: (string) the name of the MLFlow experiment
:return: (Mlflow.Run) returns the completed Run object
"""
t0 = time()
print_to_screen.print_startup_message(name, all_params, pulse_dictionary,
uris)
if "local" not in uris["tracking"].casefold():
mlflow.set_tracking_uri(uris["tracking"])
exp_id = mlflow.set_experiment(name)
with mlflow.start_run(experiment_id=exp_id) as run:
with tempfile.TemporaryDirectory() as temp_path:
if diagnostics.rules_to_store_f["space"] == "all":
x_storage_multiplier = all_params["nx"]
elif diagnostics.rules_to_store_f["space"][0] == "k0":
x_storage_multiplier = 2 * len(
diagnostics.rules_to_store_f["space"])
mem_f_store = x_storage_multiplier * all_params["nv"]
mem_field_store = all_params["nx"] * 8
# Initialize loop parameters
steps_in_loop = int(1e9 *
all_params["backend"]["max_GB_for_device"] /
(6 * (mem_f_store + mem_field_store) * 8))
if steps_in_loop > all_params["nt"]:
steps_in_loop = int(all_params["nt"] / 1.25)
n_loops = all_params["nt"] // steps_in_loop + 1
actual_num_steps = n_loops * steps_in_loop
all_params["steps_in_loop"] = steps_in_loop
all_params["n_loops"] = n_loops
all_params["actual_num_steps"] = actual_num_steps
# Get numpy arrays of the simulation configuration
stuff_for_time_loop = outer_loop.get_everything_ready_for_outer_loop(
diagnostics=diagnostics,
all_params=all_params,
pulse_dictionary=pulse_dictionary,
overall_num_steps=actual_num_steps,
)
# Initialize the storage manager -- folders, parameters, etc.
storage_manager = storage.StorageManager(
xax=stuff_for_time_loop["x"],
vax=stuff_for_time_loop["v"],
f=stuff_for_time_loop["f"],
base_path=temp_path,
rules_to_store_f=diagnostics.rules_to_store_f,
num_steps_in_one_loop=steps_in_loop,
all_params=all_params,
pulse_dictionary=pulse_dictionary,
)
mlflow.log_metrics(metrics={"startup_time": time() - t0}, step=0)
t0 = time()
sim_config, do_inner_loop = outer_loop.get_sim_config_and_inner_loop_step(
all_params=all_params,
stuff_for_time_loop=stuff_for_time_loop,
nt_in_loop=steps_in_loop,
store_f_rules=diagnostics.rules_to_store_f,
)
mlflow.log_metrics(metrics={"compile_time": time() - t0}, step=0)
t0 = time()
# TODO: Could support resume here
it_start = 0
# We run a higher level loop and a lower level loop.
# The higher level loop contains:
# 1 - Get the driver for all time-steps involved in this iteration of the lower level loop
# 2 - Perform the lower level loop
# 3 - Write the output to file
#
# The lower level loop is a loop over `steps_in_loop` timesteps. It is entirely executed on the
# accelerator. The size of this loop can be controlled by the `MAX_DOUBLES_IN_FILE` parameter.
# The goal was to allow that parameter to control the amount of memory needed on the accelerator
for it in tqdm(
range(it_start, n_loops * steps_in_loop, steps_in_loop)):
curr_time_index = np.arange(it, it + steps_in_loop)
# Get driver and time array for the duration of the lower level loop
driver_array = np.array(
stuff_for_time_loop["driver"][curr_time_index])
time_array = np.array(
stuff_for_time_loop["t"][curr_time_index])
# Perform lower level loop
sim_config = do_inner_loop(
temp_storage=sim_config,
driver_array=driver_array,
time_array=time_array,
)
mlflow.log_metrics(metrics={
"calculation_time": (time() - t0) / steps_in_loop
},
step=it)
t0 = time()
# Perform a batched data update with the lower level loop output
storage_manager.batch_update(sim_config)
mlflow.log_metrics(metrics={"batch_update_time": time() - t0},
step=it)
t0 = time()
# Run the diagnostics on the simulation so far
diagnostics(storage_manager)
mlflow.log_metrics(metrics={"diagnostic_time": time() - t0},
step=it)
mlflow.log_metrics(
metrics={
"diagnostic_time_averaged_over_sim_time":
(time() - t0) / ((it + 1) * steps_in_loop)
},
step=it,
)
t0 = time()
# Log the artifacts
storage_manager.log_artifacts()
mlflow.log_metrics(metrics={"logging_time": time() - t0},
step=it)
t0 = time()
storage_manager.close()
del storage_manager
return run
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)