Skip to content
Extraits de code Groupes Projets
run2D.py 4,77 ko
Newer Older
  • Learn to ignore specific revisions
  • import datetime as dt
    import pandas as pd
    import getopt
    import os
    import sys
    
    import slim
    import slimPre
    
    import param
    
    partition_nb = slimPre.partition_nb()
    partition_id = slimPre.partition_id()
    
    if partition_id == "0":
        print("      _/_/_/  _/        _/_/_/  _/      _/")
        print("   _/        _/          _/    _/_/  _/_/")
        print("    _/_/    _/          _/    _/  _/  _/")
        print("       _/  _/          _/    _/      _/")
        print("_/_/_/    _/_/_/_/  _/_/_/  _/      _/\n")
    
    # --------- LOAD COMMAND LINE OPTIONS  -----------------------------------------
    opts = None
    args = None
    try:
        opts, args = getopt.getopt(
            sys.argv[1:], "m:n:r:", longopts=["nonudging", "noslope", "bulkdrag", "smago="]
        )
    except getopt.GetoptError:
        print("Argument error!")
        sys.exit(1)
    
    mesh_setup = "gbr_styx"
    simu_name = ""
    restart_index = None
    
    with_nudging = True
    with_slope = True
    bulkdrag = False
    smago_param = None
    
    for opt, arg in opts:
        if opt == "-m":
            mesh_setup = arg
        elif opt == "-n":
            simu_name = arg
        elif opt == "-r":
            restart_index = int(arg)
        elif opt == "--nonudging":
            with_nudging = False
        elif opt == "--noslope":
            with_slope = False
        elif opt == "--bulkdrag":
            bulkdrag = True
        elif opt == "--smago":
            smago_param = float(arg)
    
    
    p = param.parameters(mesh_setup)
    
    output_directory = f"{p.output_directory}/{simu_name}"
    
    if partition_id == "0":
        p.print_info()
        os.makedirs(output_directory, exist_ok=True)
    # ------------------------------------------------------------------------------
    
    
    # --------- CONFIGURE HYDRO ----------------------------------------------------
    pre_data_dir = p.prepro_dir
    if os.path.isdir(pre_data_dir + "/%s" % slimPre.partition_id()):
        pre_data_dir = pre_data_dir + "/%s" % slimPre.partition_id()
    
    if slimPre.partition_nb() == "1":
        mesh_file_name = p.mesh_file
    else:
        mesh_file_name = p.mesh_file[:-4] + "_" + slimPre.partition_nb() + ".msh"
    
    domain = slim.Domain(
        mesh_file_name, (p.prepro_dir + "/bathymetry_smooth.nc", "bathymetry")
    )
    
    eq = slim.ShallowWater2d(
        domain, "implicit", initial_time=p.initial_time, final_time=p.final_time
    )
    
    if smago_param is None:
        eq.set_viscosity("smagorinsky")
    else:
        eq.set_viscosity("smagorinsky", smagorinsky_coefficient=smago_param)
    
    if bulkdrag:
        eq.set_dissipation("quadratic", (p.prepro_dir + "/reef.nc", "bulk_0025_05"))
    else:
        eq.set_dissipation("manning", (p.prepro_dir + "/reef.nc", "manning_025_25"))
    
    if with_slope:
        eq.set_dissipation("slope", (p.prepro_dir + "/slope.nc", "cs5"))
    
    eq.set_coriolis((p.prepro_dir + "/coriolis.nc", "coriolis"))
    eq.set_wind_stress(
        "speed",
        (pre_data_dir + "/eReefs_wind.nc", "u"),
        (pre_data_dir + "/eReefs_wind.nc", "v"),
        density_air=1.204,
    )
    
    if with_nudging:
        eq.set_nudging(
            (p.prepro_dir + "/nudging.nc", "gamma"),
            external_ux=(pre_data_dir + "/mercator_and_tides.nc", "u"),
            external_uy=(pre_data_dir + "/mercator_and_tides.nc", "v"),
            transport_flux=False,
            ramp_period=5 * p.dt_export,
        )
    eq.set_boundary_coast(p.closed_tags, slip=False)
    eq.set_boundary_open(
        p.open_tags,
        sse=(pre_data_dir + "/mercator_and_tides.nc", "h"),
        ux=(pre_data_dir + "/mercator_and_tides.nc", "u"),
        uy=(pre_data_dir + "/mercator_and_tides.nc", "v"),
        transport_flux=False,
        ramp_period=5 * p.dt_export,
    )
    streamflow_df = pd.read_csv(
        f"{p.local_base_dir}/spatial/geoflows_australia_on_coast_utm56.csv"
    )
    river_discharge_dir = p.prepro_dir + "/discharge/"
    for s in streamflow_df.COMID:
        eq.set_boundary_open(
            str(s),
            discharge=(
                river_discharge_dir + "river_discharge_" + str(s) + ".nc",
                "river_discharge",
            ),
        )
    # ------------------------------------------------------------------------------
    
    # --------- RUN HYDROS ---------------------------------------------------------
    loop = slim.Loop(
        maximum_time_step=p.dt,
        export_time=p.dt_export,
        path=output_directory,
        export_format="abin",
    )
    loop.add_equation(eq)
    loop.set_export_options(
        float_size=4,
        continuous=True,
        checkpoint_ratio=p.checkpoint_ratio,
        keep_only_last_checkpoint=True,
    )
    
    if restart_index is not None:
        fmt = "%Y-%m-%d %H:%M:%S"
        restart_time = dt.datetime.strptime(p.initial_time, fmt) + dt.timedelta(
            seconds=restart_index * p.dt_export
        )
        if partition_id == "0":
            print(f"Restarting on {restart_time}")
        restart_time = dt.datetime.strftime(restart_time, fmt)
        loop.restart(
            index=restart_index,
            time=restart_time,
            use_safe_iterate=True,
            reload_dir=output_directory + "/checkpoint/",
            append=True,
        )
    
    loop.run(use_safe_iterate=True)
    
    slim.exit()
    # ------------------------------------------------------------------------------