PyMPDATA_examples.Olesik_et_al_2022.analysis

  1from collections import namedtuple
  2from copy import deepcopy
  3
  4import numpy as np
  5from joblib import Parallel, delayed, parallel_backend
  6from PyMPDATA_examples.Olesik_et_al_2022.coordinates import x_id, x_log_of_pn, x_p2
  7from PyMPDATA_examples.Olesik_et_al_2022.equilibrium_drop_growth import PdfEvolver
  8from PyMPDATA_examples.Olesik_et_al_2022.settings import (
  9    Settings,
 10    default_mixing_ratios_g_kg,
 11)
 12from PyMPDATA_examples.Olesik_et_al_2022.simulation import Simulation
 13from PyMPDATA_examples.utils.discretisation import discretised_analytical_solution
 14from PyMPDATA_examples.utils.error_norms import L2
 15
 16from PyMPDATA import Options
 17
 18
 19def analysis(settings, grid_layout, psi_coord, options_dict, GC_max):
 20    options_str = str(options_dict)
 21    options = Options(**options_dict)
 22    simulation = Simulation(settings, grid_layout, psi_coord, options, GC_max)
 23    result = {"n": [], "n_analytical": [], "error_norm_L2": [], "wall_time": []}
 24    last_step = 0
 25    for n_steps in simulation.out_steps:
 26        steps = n_steps - last_step
 27        wall_time = simulation.step(steps) if steps > 0 else 0
 28        last_step += steps
 29        result["n"].append(simulation.n_of_r.copy())
 30        result["wall_time"].append(wall_time)
 31    result["r"] = simulation.r.copy()
 32    result["rh"] = simulation.rh.copy()
 33    result["dx"] = simulation.dx
 34    return Result(
 35        grid_layout_str=grid_layout.__class__.__name__,
 36        option_str=options_str,
 37        result=result,
 38        out_steps=simulation.out_steps,
 39        dt=simulation.dt,
 40    )
 41
 42
 43def compute_figure_data(
 44    *,
 45    nr,
 46    GC_max,
 47    psi_coord=x_id(),
 48    grid_layouts=(x_id(), x_p2(), x_log_of_pn(r0=1, n=1)),
 49    opt_set=({"n_iters": 1},),
 50    mixing_ratios_g_kg=default_mixing_ratios_g_kg
 51):
 52    settings = Settings(nr=nr, mixing_ratios_g_kg=mixing_ratios_g_kg)
 53    with parallel_backend("threading", n_jobs=-2):
 54        results = Parallel(verbose=10)(
 55            delayed(analysis)(settings, grid_layout, psi_coord, options, GC_max)
 56            for grid_layout in grid_layouts
 57            for options in deepcopy(opt_set)
 58        )
 59
 60    cases = {}
 61    for result in results:
 62        case = Case(result)
 63        if case.grid_layout_str not in cases:
 64            cases[case.grid_layout_str] = case
 65
 66    output = {}
 67    for coord, case in cases.items():
 68        output[coord] = {"numerical": {}, "wall_time": {}}
 69        for result in results:
 70            if coord == result.grid_layout_str:
 71                opts = result.option_str
 72                data = result.result
 73                rh = data.pop("rh")
 74                r = data.pop("r")
 75                dx = data.pop("dx")
 76
 77                if "grid" not in output[coord]:
 78                    output[coord]["grid"] = {
 79                        "rh": rh,
 80                        "r": r,
 81                        "dx": dx,
 82                        "dt": case.dt,
 83                        "out_steps": case.out_steps,
 84                    }
 85                output[coord]["numerical"][opts] = data["n"]
 86
 87    for coord, case in cases.items():
 88        analytical = []
 89        for t in [case.dt * nt for nt in case.out_steps]:
 90            pdf_t = PdfEvolver(settings.pdf, settings.drdt, t)
 91            rh = output[coord]["grid"]["rh"]
 92            analytical.append(
 93                discretised_analytical_solution(
 94                    rh.magnitude,
 95                    lambda r, pdf_t=pdf_t, rh=rh: pdf_t(r * rh.units).magnitude,
 96                    midpoint_value=True,
 97                    r=output[coord]["grid"]["r"].magnitude,
 98                )
 99                * pdf_t(rh[0]).units
100            )
101        output[coord]["analytical"] = analytical
102
103    for coord, case in cases.items():
104        error_L2 = {}
105        analytical = output[coord]["analytical"]
106        for opts in output[coord]["numerical"]:
107            numerical = output[coord]["numerical"][opts]
108            error_L2[opts] = L2(
109                numerical[-1].magnitude, analytical[-1].magnitude, case.out_steps[-1]
110            )
111        output[coord]["error_L2"] = error_L2
112
113    for coord, case in cases.items():
114        for result in results:
115            data = result.result
116            for opts in output[coord]["numerical"]:
117                output[coord]["wall_time"][opts] = data["wall_time"]
118
119    return output, settings
120
121
122class Result(
123    namedtuple("Result", ("dt", "out_steps", "grid_layout_str", "option_str", "result"))
124):
125    __slots__ = ()
126
127
128def Case(result: Result):
129    return namedtuple("Case", ("dt", "out_steps", "grid_layout_str"))(
130        dt=result.dt, out_steps=result.out_steps, grid_layout_str=result.grid_layout_str
131    )
132
133
134def rel_disp(r, psi, psi_coord):
135    mom0 = 0
136    mom1 = 0
137    mom2 = 0
138    for i, psi_i in enumerate(psi):
139        dp_i = psi_coord.moment_of_r_integral(
140            psi_coord.x(r[i + 1]), 0
141        ) - psi_coord.moment_of_r_integral(psi_coord.x(r[i]), 0)
142        A_i = psi_coord.moment_of_r_integral(
143            psi_coord.x(r[i + 1]), 1
144        ) - psi_coord.moment_of_r_integral(psi_coord.x(r[i]), 1)
145        B_i = psi_coord.moment_of_r_integral(
146            psi_coord.x(r[i + 1]), 2
147        ) - psi_coord.moment_of_r_integral(psi_coord.x(r[i]), 2)
148        bin_mom0 = psi_i * dp_i
149        bin_mom1 = psi_i * A_i
150        bin_mom2 = psi_i * B_i
151        mom0 += bin_mom0
152        mom1 += bin_mom1
153        mom2 += bin_mom2
154    mu = mom1 / mom0
155    std = np.sqrt(mom2 / mom0 - mu**2)
156    d = std / mu
157    return d
158
159
160def third_moment(r, psi, psi_coord, normalised=True):
161    mom0 = 0
162    mom3 = 0
163    for i, psi_i in enumerate(psi):
164        dp_i = psi_coord.moment_of_r_integral(
165            psi_coord.x(r[i + 1]), 0
166        ) - psi_coord.moment_of_r_integral(psi_coord.x(r[i]), 0)
167        integral_i = psi_coord.moment_of_r_integral(
168            psi_coord.x(r[i + 1]), 3
169        ) - psi_coord.moment_of_r_integral(psi_coord.x(r[i]), 3)
170        bin_mom0 = psi_i * dp_i
171        bin_mom3 = psi_i * integral_i
172        mom0 += bin_mom0
173        mom3 += bin_mom3
174    if normalised:
175        mom3 /= mom0
176    return mom3