PyMPDATA_examples.Olesik_et_al_2022.wall_time

  1import pathlib
  2
  3import numpy as np
  4from PyMPDATA_examples.Olesik_et_al_2022.coordinates import x_id, x_log_of_pn
  5from PyMPDATA_examples.Olesik_et_al_2022.settings import (
  6    Settings,
  7    default_GC_max,
  8    default_nr,
  9    default_opt_set,
 10)
 11from PyMPDATA_examples.Olesik_et_al_2022.simulation import Simulation
 12
 13from PyMPDATA import Options
 14
 15grid_layout_set = (x_log_of_pn(r0=1, base=2),)
 16opt_set = default_opt_set.values()
 17
 18RTOL = 1.5
 19
 20
 21def test_wall_time(n_runs=3, mrats=(20,), generate=False, print_tab=True, rtol=RTOL):
 22    settings = Settings(nr=default_nr * 10, mixing_ratios_g_kg=np.array(mrats))
 23    table_data = {"opts": [], "values": []}
 24    for grid in grid_layout_set:
 25        for opts in opt_set:
 26            i = 0
 27            minimum_values = []
 28            while i < n_runs:
 29                result = make_data(settings, grid, opts)
 30                wall_times = np.asarray(result["wall_time"])
 31                minimal = np.nanmin(wall_times[wall_times > 0])
 32                minimum_values.append(minimal)
 33                i += 1
 34            selected_value = np.min(minimum_values)
 35            if opts == {"n_iters": 1}:
 36                norm = selected_value
 37            table_data["opts"].append(str(opts) + "(" + grid.__class__.__name__ + ")")
 38            table_data["values"].append(
 39                np.nan if norm == 0 else round(selected_value / norm, 1)
 40            )
 41    make_textable(data=table_data, generate=generate, print_tab=print_tab)
 42    compare_refdata(data=table_data, rtol=rtol, generate=generate)
 43
 44
 45def make_data(settings, grid, opts):
 46    options = Options(**opts)
 47    simulation = Simulation(
 48        settings=settings,
 49        grid_layout=grid,
 50        psi_coord=x_id(),
 51        opts=options,
 52        GC_max=default_GC_max,
 53    )
 54    result = {"wall_time": []}
 55    last_step = 0
 56    for n_steps in simulation.out_steps:
 57        steps = n_steps - last_step
 58        wall_time_per_timestep = simulation.step(steps)
 59        last_step += steps
 60        result["wall_time"].append(wall_time_per_timestep)
 61    return result
 62
 63
 64def make_textable(data, generate=False, print_tab=False):
 65    latex_data = (
 66        r"\hline" + " Variant  & Elapsed Real Time (wrt upwind) " + r"\\ \hline" + "\n"
 67    )
 68    for opt, value in zip(data["opts"], data["values"]):
 69        latex_data += r"\hline" + f" {opt} & {value} " + r"\\ \hline" + "\n"
 70    latex_start = r"\begin{table}[]" + "\n" + r"\begin{tabular}{| l | l |}" + "\n"
 71    latex_end = r"\end{tabular}" + "\n" + r"\end{table}"
 72    latex_table = latex_start + latex_data + latex_end
 73    if print_tab:
 74        print(latex_table)
 75    with open(
 76        pathlib.Path(__file__).parent.joinpath("wall_time_textable.txt"),
 77        "w+" if generate else "r",
 78        encoding="utf-8",
 79    ) as f:
 80        if generate:
 81            f.write(latex_table)
 82
 83
 84def compare_refdata(data, rtol, generate=False):
 85    delimiter = ";"
 86    path = pathlib.Path(__file__).parent.joinpath("wall_time_refdata.txt")
 87    if generate:
 88        table = (
 89            np.char.array(np.concatenate([data["opts"], data["values"]]))
 90            .reshape(2, len(data["values"]))
 91            .T
 92        )
 93        np.savetxt(path, table, delimiter=delimiter, fmt="%s")
 94    else:
 95        table = np.loadtxt(path, delimiter=delimiter, dtype=str)
 96        np.testing.assert_allclose(
 97            actual=data["values"],
 98            desired=np.array(table[:, 1].astype(float)),
 99            rtol=rtol,
100        )
101        np.testing.assert_array_equal(data["opts"], table[:, 0])