cfno/scripts/01_simu.py at main · Parallel-in-Time/cfno · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import sys
import argparse
import numpy as np
sys.path.insert(2, os.getcwd())
from cfno.simulation.rbc2d import runSim, runSimPySDC, MPI_RANK
from cfno.simulation.post import OutputFiles
from cfno.utils import readConfig

# -----------------------------------------------------------------------------
# Script parameters
# -----------------------------------------------------------------------------
parser = argparse.ArgumentParser(
    description='Generate Dedalus simulation data for 2D RBC',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
    "--dataDir", default="simuData", help="directory containing simulation data")
parser.add_argument(
    "--Rayleigh", default=1e7, type=float, help="Rayleigh number")
parser.add_argument(
    "--tInit", default=100, type=float, help="initial simulation time (before generating data)")
parser.add_argument(
    "--resFactor", default=1, type=int, help="resolution factor for space grid (256*resFactor, 64*resFactor)")
parser.add_argument(
    "--nSimu", default=10, type=int, help="number of simulation with different seeds")
parser.add_argument(
    "--dtSimu", default=0.001, type=float, help="base time-step for simulation")
parser.add_argument(
    "--dtData", default=0.1, type=float, help="time-step for simulation data output")
parser.add_argument(
    "--tEnd", default=100, type=float, help="simulation time for data generation")
parser.add_argument(
    "--seed", default=12345678, type=int, help="starting seed to generate all simulation seeds")
parser.add_argument(
    "--solver", default="dedalus", choices=["dedalus", "pySDC"], help="solver used to generate data")
parser.add_argument(
    "--config", default=None, help="config file, overwriting all parameters specified in it")
args = parser.parse_args()

if args.config is not None:
    config = readConfig(args.config)
    assert "simu" in config, "config file needs a simu section"
    args.__dict__.update(**config.simu)

dataDir = args.dataDir
Rayleigh = args.Rayleigh
tInit = args.tInit
resFactor = args.resFactor
nSimu = args.nSimu
dtSimu = args.dtSimu
dtData = args.dtData
tEnd = args.tEnd
solver = args.solver
np.random.seed(args.seed)

# -----------------------------------------------------------------------------
# Script execution
# -----------------------------------------------------------------------------
def log(msg):
    if MPI_RANK == 0:
        print(msg)

seeds = [int(s*1e6) for s in np.random.rand(nSimu)]

for seed in seeds:

    # Setup simulation folder names
    simDir = f"{dataDir}/simu_{seed}"

    initRunDir = f"{simDir}/run_init"
    dtInit = 1e-2/2
    os.makedirs(initRunDir, exist_ok=True)

    dataRunDir = f"{simDir}/run_data"
    os.makedirs(dataRunDir, exist_ok=True)

    if solver == "dedalus":
        # -- run initial simulation
        log(f" -- running initial dedalus simulation with dt={dtInit:1.1e} in {initRunDir}")
        runSim(initRunDir, Rayleigh, resFactor, baseDt=dtInit, useSDC=False, tEnd=tInit,
            dtWrite=tInit, writeFull=True, tBeg=0,seed=seed)

        # -- extract initial field for data generation
        initFiles = OutputFiles(initRunDir)
        initFields = initFiles.file(0)['tasks']

        # -- generate simulation data
        log(f" -- generating dedalus simulation data with dt={dtData:1.1e} (dtSimu={dtSimu:1.1e}) in {dataRunDir}")
        runSim(dataRunDir, Rayleigh, resFactor, baseDt=dtSimu, useSDC=False,
            tEnd=tEnd, dtWrite=dtData, initFields=initFields)

    elif solver == "pySDC":
        # -- run initial simulation
        log(f" -- running initial pySDC simulation with dt={dtInit:1.1e} in {initRunDir}")
        runSimPySDC(initRunDir, Rayleigh, resFactor, baseDt=dtInit, tBeg=0, tEnd=tInit,
            dtWrite=tInit, seed=seed)
        initFile = f"{initRunDir}/outputs.pysdc"

        # -- generate simulation data
        log(f" -- generating pySDC simulation data with dt={dtData:1.1e} (dtSimu={dtSimu:1.1e}) in {dataRunDir}")
        runSimPySDC(dataRunDir, Rayleigh, resFactor, baseDt=dtSimu,
            tEnd=tEnd, dtWrite=dtData, restartFile=initFile)