hw2d.run2

  1import fire
  2import numpy as np
  3import matplotlib.pyplot as plt
  4from tqdm import tqdm
  5import os
  6import time
  7from typing import Iterable
  8
  9# Local imports
 10from hw2d.model import HW
 11from hw2d.initializations.fourier_noise import get_fft_noise
 12from hw2d.initializations.sine import get_2d_sine
 13from hw2d.utils.io import (
 14    get_save_params,
 15    create_appendable_h5,
 16    save_to_buffered_h5,
 17    append_h5,
 18    continue_h5_file,
 19)
 20from hw2d.utils.namespaces import Namespace
 21from hw2d.utils.plot.movie import create_movie
 22from hw2d.utils.run_properties import calculate_properties
 23from hw2d.utils.plot.timetrace import plot_timetraces, plot_timetrace_comparison
 24from hw2d.utils.downsample import fourier_downsample
 25from hw2d.utils.helpers import format_print_dict
 26from hw2d.physical_properties.numpy_properties import (
 27    get_gamma_n,
 28    get_gamma_n_spectrally,
 29    get_gamma_c,
 30    get_energy,
 31    get_energy_N_spectrally,
 32    get_energy_V_spectrally,
 33    get_enstrophy,
 34    get_enstrophy_phi,
 35    get_DE,
 36    get_DU,
 37    get_dE_dt,
 38    get_dU_dt,
 39)
 40
 41
 42property_fncs = {
 43    "gamma_n": lambda n, p, dx, **kwargs: get_gamma_n(n=n, p=p, dx=dx),
 44    "gamma_n_spectral": lambda n, p, dx, **kwargs: get_gamma_n_spectrally(
 45        n=n, p=p, dx=dx
 46    ),
 47    "gamma_c": lambda n, p, dx, c1, **kwargs: get_gamma_c(n=n, p=p, c1=c1, dx=dx),
 48    "energy": lambda n, p, dx, **kwargs: get_energy(n=n, phi=p, dx=dx),
 49    "thermal_energy": lambda n, **kwargs: get_energy_N_spectrally(n=n),
 50    "kinetic_energy": lambda p, dx, **kwargs: get_energy_V_spectrally(p=p, dx=dx),
 51    "enstrophy": lambda n, o, dx, **kwargs: get_enstrophy(n=n, omega=o, dx=dx),
 52    "enstrophy_phi": lambda n, p, dx, **kwargs: get_enstrophy_phi(n=n, phi=p, dx=dx),
 53    "DE": lambda n, p, Dn, Dp, **kwargs: get_DE(n=n, p=p, Dn=Dn, Dp=Dp),
 54    "DU": lambda n, o, Dn, Dp, **kwargs: get_DU(n=n, o=o, Dn=Dn, Dp=Dp),
 55    "dE_dt": lambda gamma_n, gamma_c, DE, **kwargs: get_dE_dt(gamma_n=gamma_n, gamma_c=gamma_c, DE=DE),
 56    "dU_dt": lambda gamma_n, DU, **kwargs: get_dU_dt(gamma_n=gamma_n, DU=DU),
 57}
 58
 59
 60def run(
 61    # Physics & Numerics
 62    step_size: float = 0.025,
 63    adaptive_step_size: bool = False,
 64    end_time: float = 1_000,
 65    grid_pts: int = 512,
 66    k0: float = 0.15,
 67    N: int = 3,
 68    nu: float = 5.0e-09,
 69    c1: float = 1,
 70    kappa_coeff: float = 1.0,
 71    poisson_bracket_coeff: float = 1.0,
 72    zonal: bool = False,
 73    # Running
 74    show_property: str = "gamma_n",
 75    # Initialization
 76    seed: int or None = None,
 77    init_type: str = "normal",
 78    init_scale: float = 1 / 100,
 79    # Saving
 80    output_path: str = "c1=1.0_nu=5e-09_dt=0.025_v=1.h5",
 81    recording_start_time: float = 0,
 82    continue_file: bool = False,
 83    buffer_length: int = 100,
 84    snaps: int = 1,
 85    chunk_size: int = 100,
 86    downsample_factor: float = 16,
 87    add_last_state: bool = True,
 88    # Movie
 89    movie: bool = False,
 90    min_fps: int = 10,
 91    dpi: int = 75,
 92    speed: int = 10,
 93    # Properties
 94    properties: Iterable[str] = [
 95        "gamma_n",
 96        "gamma_n_spectral",
 97        "gamma_c",
 98        "energy",
 99        "thermal_energy",
100        "kinetic_energy",
101        "enstrophy",
102        "enstrophy_phi",
103    ],
104    # Plotting
105    t0_std=500,
106    plot_properties: Iterable[str] = (
107        "gamma_c",
108        "gamma_n",
109        "gamma_n_spectral",
110        "energy",
111        "kinetic_energy",
112        "thermal_energy",
113        "enstrophy",
114        "enstrophy_phi"
115    ),
116    # Other
117    debug: bool = False,
118    force_recompute: bool = False,
119):
120    """
121    Run the simulation with the given parameters.
122
123    Args:
124        step_size (float, optional): Incremental step for simulation progression. Defaults to 0.025.
125        end_time (float, optional): Duration till the simulation should run. Defaults to 1_000.
126        grid_pts (int, optional): Grid points. Suggested: 128 for coarse, 1024 for fine. Defaults to 512.
127        k0 (float, optional): Determines k-focus. Suggested: 0.15 for high-k, 0.0375 for low-k. Defaults to 0.15.
128        N (int, optional): Dissipation exponent's half value. Defaults to 3.
129        nu (float, optional): Viscosity. Suggested: 5e-10 for coarse-large, 1e-4 for fine-small. Defaults to 5.0e-08.
130        c1 (float, optional): Transition scale between hydrodynamic and adiabatic. Suggested values: 0.1, 1, 5. Defaults to 1.0.
131        kappa_coeff (float, optional): Coefficient of d/dy phi. Defaults to 1.0.
132        poisson_bracket_coeff (float, optional): Coefficient of Poisson bracket [A,B] implemented with Arakawa Scheme. Defaults to 1.0.
133        zonal (bool, ooptional): If True, uses make_zonal_func to simulate mHW2D. Defaults to False.
134        seed (int or None, optional): Seed for random number generation. Defaults to None.
135        init_type (str, optional): Initialization method. Choices: 'fourier', 'sine', 'random', 'normal'. Defaults to 'normal'.
136        init_scale (float, optional): Scaling factor for initialization. Defaults to 0.01.
137        output_path (str, optional): Where to save the simulation data. Defaults to ''.
138        recording_start_time (float, optional): Time (t) from which onwards the data is recorded. Defaults to 0.
139        continue_file (bool, optional): If True, continue with existing file. Defaults to False.
140        buffer_length (int, optional): Size of buffer for storage. Defaults to 100.
141        chunks (int, optional): Chunks of the h5 file. Defaults to 100.
142        downsample_factor (float, optional): Factor along each axis that it is downsampled with for saving. Defaults to 1.
143        add_last_state (bool, optional): Whether the last high-resolution frame is saved. Turns True if downsampling on. Defaults to False.
144        snaps (int, optional): Snapshot intervals for saving. Defaults to 1.
145        movie (bool, optional): If True, generate a movie out of simulation. Defaults to True.
146        min_fps (int, optional): Parameter for movie generation. Defaults to 10.
147        dpi (int, optional): Parameter for movie generation. Defaults to 75.
148        speed (int, optional): Parameter for movie generation. Defaults to 5.
149        properties (Iterable[str], optional): List of properties to calculate for the saved file.
150        plot_properties (Iterable[str], optional): List of properties to plot a timetrace for.
151        debug (bool, optional): Enable or disable debug mode. Defaults to False.
152
153    Returns:
154        None: The function saves simulation data or generates a movie as specified.
155    """
156    step_size = float(step_size)
157    nu = float(nu)
158
159    # Unpacking
160    y = grid_pts
161    x = grid_pts
162    L = 2 * np.pi / k0  # Physical box size
163    dx = L / x  # Grid resolution
164    steps = int(end_time / step_size)  # Number of Steps until end_time
165    snap_count = steps // snaps + 1  # number of snapshots
166    field_list = ("density", "omega", "phi")
167    initial_time = np.float64(0)
168    iteration_count = 0
169    np.random.seed(seed)
170    if downsample_factor != 1:
171        add_last_state = True
172    if show_property not in properties:
173        properties += [show_property]
174
175    # Define Initializations
176    noise = {
177        "fourier": lambda y, x: get_fft_noise(
178            resolution=[y, x],
179            size=L,
180            scale=1,
181            min_wavelength=dx * 12,
182            max_wavelength=dx * grid_pts * 100,
183            factor=2,
184        ),
185        "sine": lambda y, x: get_2d_sine((y, x), L),
186        "random": lambda y, x: np.random.rand(y, x).astype(np.float64),
187        "normal": lambda y, x: np.random.normal(size=(y, x)).astype(np.float64),
188    }
189    downsample_fnc = fourier_downsample
190
191    # Physics
192    # TODO: Move all to Tensors along batch dimension for simulations!
193    physics_params = dict(
194        dx=dx,
195        N=N,
196        c1=c1,
197        nu=nu,
198        k0=k0,
199        poisson_bracket_coeff=poisson_bracket_coeff,
200        kappa_coeff=kappa_coeff,
201        zonal=zonal
202    )
203    # Initialize Plasma
204    plasma = Namespace(
205        density=noise[init_type](y, x) * init_scale,
206        omega=noise[init_type](y, x) * init_scale,
207        phi=noise[init_type](y, x) * init_scale,
208        age=0,
209    )
210
211    # File Handling
212    if output_path:
213        y_save = y
214        x_save = x
215        if downsample_factor != 1:
216            y_save = int(round(y / downsample_factor))
217            x_save = int(round(x / downsample_factor))
218            print(f"Downsample by {downsample_factor}x: {y}x{x}->{y_save}x{x_save}")
219        # Output data from this run
220        buffer = {
221            field: np.zeros((buffer_length, y_save, x_save), dtype=np.float32)
222            for field in field_list
223        }
224        # Property buffer
225        buffer[f"time"] = np.zeros((buffer_length, 1), dtype=np.float32)
226        for prop_name in properties:
227            buffer[f"fullres_{prop_name}"] = np.zeros(
228                (buffer_length, 1), dtype=np.float32
229            )
230        # Other output parameters
231        output_params = {
232            "buffer": buffer,
233            "buffer_index": 0,
234            "output_path": output_path,
235        }
236        # Load Data
237        if os.path.isfile(output_path) and not force_recompute:
238            if continue_file:
239                plasma, physics_params = continue_h5_file(output_path, field_list)
240                save_params = get_save_params(
241                    physics_params,
242                    step_size,
243                    snaps,
244                    x,
245                    y,
246                    x_save=x_save,
247                    y_save=y_save,
248                    recording_start_time=recording_start_time,
249                    adaptive_step_size=adaptive_step_size,
250                )
251                initial_time = np.float64(plasma.age)
252                print(
253                    f"Loaded: {output_path} (age={plasma.age}, time={initial_time})\n{physics_params}"
254                )
255                # Check if broken
256                for field in plasma.keys():
257                    # Not NaNs
258                    if np.isnan(np.sum(plasma[field])):
259                        print(f"Input file is broken: NaNs in {field}")
260                        raise BaseException(f"Input file is broken: NaNs in {field}")
261                    # Not zeros
262                    if field in ("density", "omega"):
263                        if np.all(plasma[field] == 0):
264                            print(f"FAILED {field} is zero")
265                            raise BaseException(f"Input File is zero: {field}")
266            else:
267                print(f"File already exists.")
268                return
269        # Create Data
270        else:
271            save_params = get_save_params(
272                physics_params,
273                step_size,
274                snaps,
275                x,
276                y,
277                x_save=x_save,
278                y_save=y_save,
279                recording_start_time=recording_start_time,
280                adaptive_step_size=adaptive_step_size,
281            )
282            # Create file
283            create_appendable_h5(
284                output_path,
285                save_params,
286                properties=["time"]
287                + [f"fullres_{prop_name}" for prop_name in properties],
288                dtype=np.float32,
289                chunk_size=chunk_size,
290                add_last_state=add_last_state,
291            )
292            # Save initial values
293            if not recording_start_time:
294                new_val = Namespace(
295                    **{
296                        k: v
297                        for k, v in plasma.items()
298                        if k in ("phi", "omega", "density")
299                    }
300                )
301                last_state = {}
302                for k, v in plasma.items():
303                    if k in ("phi", "omega", "density"):
304                        if downsample_factor != 1:
305                            new_val[k] = downsample_fnc(v, downsample_factor)
306                        if add_last_state:
307                            last_state[f"state_{k}"] = v
308                for prop_name in properties:
309                    new_val[f"fullres_{prop_name}"] = property_fncs[prop_name](
310                        n=plasma["density"],
311                        p=plasma["phi"],
312                        o=plasma["omega"],
313                        dx=dx,
314                        c1=c1,
315                    )
316                new_val["time"] = initial_time
317                output_params["buffer_index"] = save_to_buffered_h5(
318                    new_val=new_val,
319                    last_state=last_state,
320                    buffer_length=buffer_length,
321                    **output_params,
322                )
323
324    # Display runtime parameters
325    # save_params["output_path"] = output_path  # No support for strings
326    format_print_dict(save_params)
327
328    # Setup Simulation
329    hw = HW(**physics_params, debug=debug)
330    plasma["phi"] = hw.get_phi(plasma.omega, physics_params["dx"])
331
332    # Setup Progress Bar
333    print("Running simulation...")
334    iteration_count = 0
335    bar_format = "{rate_fmt} | {desc} {percentage:>6.2f}%|{bar}| {n:.2f}/{total:.2f} [{elapsed}<{remaining}]"
336    # bar_format = "{percentage:>6.2f}%|{bar}| {n:.4f}/{total_fmt} [{elapsed}<{remaining}] ({rate_fmt})"
337    used_step_size = step_size
338    used_time_steps = []
339    current_time = initial_time
340    new_val = {}
341
342    try:
343        with tqdm(
344            total=end_time - initial_time,
345            unit="t",
346            bar_format=bar_format,
347        ) as pbar:
348            # Run Simulation
349            while round(current_time, 4) < end_time:
350                # Progress one step, alternatively: hw.euler_step()
351                successful = False
352                while not successful:
353                    try:
354                        plasma = hw.rk4_step(plasma, dt=step_size, dx=dx)
355                        used_step_size = step_size
356                        successful = True
357                    except Exception as e:
358                        print(e)
359                        time.sleep(10)
360                iteration_count += 1
361
362                # Improved accuracy of time tracking
363                if adaptive_step_size:
364                    used_time_steps.append(np.float64(used_step_size))
365                    current_time = initial_time + np.sum(used_time_steps)
366                else:
367                    current_time = initial_time + iteration_count * step_size
368
369                # Batched Saving
370                # TODO: Figure out whether it should be iteration_count % (snaps - 1) == 0
371                if (current_time >= recording_start_time) and (
372                    iteration_count % snaps == 0
373                ):
374                    new_val = Namespace(
375                        **{
376                            k: v.copy()
377                            for k, v in plasma.items()
378                            if k in ("phi", "omega", "density")
379                        }
380                    )
381                    # Add properties if they were selected
382                    # TODO: Faster to do before saving through trivial parallelization
383                    new_val["time"] = current_time
384                    for prop_name in properties:
385                        new_val[f"fullres_{prop_name}"] = property_fncs[prop_name](
386                            n=plasma["density"],
387                            p=plasma["phi"],
388                            o=plasma["omega"],
389                            dx=dx,
390                            c1=c1,
391                        )
392
393                    # Save to records
394                    if output_path:
395                        last_state = {}
396                        # Save downsampled & state
397                        for k, v in plasma.items():
398                            if k in ("phi", "omega", "density"):
399                                # Downsample for saving
400                                if downsample_factor != 1:
401                                    new_val[k] = downsample_fnc(v.copy(), downsample_factor)
402                                # Add last state as high-res for continuing simulation
403                                if add_last_state:
404                                    last_state[f"state_{k}"] = v
405                        # Save to buffer/file
406                        output_params["buffer_index"] = save_to_buffered_h5(
407                            new_val=new_val,
408                            last_state=last_state,
409                            buffer_length=buffer_length,
410                            **output_params,
411                        )
412
413                    # Check for breaking
414                    if np.isnan(np.sum(plasma.density)):
415                        print(f"FAILED @ {iteration_count:,} steps ({plasma.age:,})")
416                        break
417
418                # Calculate iterations per second, handling division by zero
419                try:
420                    iter_per_sec = iteration_count / pbar.format_dict["elapsed"]
421                except ZeroDivisionError:
422                    iter_per_sec = float("inf")  # or set to 0 or any default value
423
424                # Update progress
425                pbar.update(used_step_size)
426                new_description = f"{iter_per_sec:.2f}it/s"
427                if show_property:
428                    new_description += f" | Γn = {new_val.get(show_property, 0):.2g}"
429                pbar.set_description(new_description)
430
431    except Exception as e:
432        print(e)
433        raise e
434        # os.remove(output_path)
435        # print(f"removed {output_path}")
436        return
437
438    # If output_path is defined, flush any remaining data in the buffer
439    if (
440        output_path
441        and (current_time > recording_start_time)
442        and (output_params["buffer_index"] > 0)
443    ):
444        append_h5(**output_params)
445
446    # Get Performance stats
447    hw.print_log()
448
449    # Generate Movie from saved file
450    if movie and output_path:
451        print(f"Generating movie...")
452        create_movie(
453            input_filename=output_path,
454            output_filename=output_path.replace(".h5", ""),
455            t0=0,
456            t1=None,
457            plot_order=field_list,
458            min_fps=min_fps,
459            dpi=dpi,
460            speed=speed,
461        )
462
463    if properties and output_path:
464        print(f"Calculating properties...")
465        calculate_properties(
466            file_path=output_path,
467            batch_size=buffer_length,
468            property_list=properties,
469            force_recompute=True,
470            is_debug=False,
471        )
472
473    if plot_properties and output_path:
474        print(f"Plotting properties...")
475        plot_timetrace_comparison(
476            file_path=output_path,
477            out_path=None,
478            properties=plot_properties,
479            t0=0,
480            t0_std=t0_std,
481        )
482
483
484if __name__ == "__main__":
485    fire.Fire(run)