LPT_quasi3D_sediments.py

import getopt
import os
import sys
import time
import datetime

import dgpy
import geopandas as gpd
import numpy as np
import pandas as pd
import slimParticle
import slimPre
from scipy import io

import LPT_functions
import param
import sediment_functions as sf

# ------------------------------------------------------------------------------

# ----- Set parameters --------------------------------------------------------_
number_particles = int(5e3)
simu_name = "one_month_release_with_rivers_constant_seed_by_site"
mesh_name = "gbr_styx_with_rivers"
close_rivers = False

nu = 1e-2

min_size = 4  # min size in micrometers
max_size = 8  # max size in micrometers

opts = None
args = None
try:
    opts, args = getopt.getopt(sys.argv[1:], "cl:h:p:n:m:i:")
except getopt.GetoptError:
    print("Argument error!")
    sys.exit(1)


for opt, arg in opts:
    if opt == "-l":
        min_size = int(arg)
    elif opt == "-h":
        max_size = int(arg)
    elif opt == "-p":
        number_particles = int(arg)
    elif opt == "-n":
        simu_name = arg
    elif opt == "-m":
        mesh_name = arg
    elif opt == "-c":
        close_rivers = True
    elif opt == "-i":
        nu = float(arg)


LPT_initial_time = "2021-01-05 00:00:00"
LPT_final_time = "2021-03-31 12:00:00"
release_initial_time = LPT_initial_time
release_final_time = "2021-02-05 00:00:00"
# ------------------------------------------------------------------------------

# ----- Parametrize simu -------------------------------------------------------
p = param.parameters(mesh_name)
hydro_dir = p.output_directory
hydro_initial_time = p.initial_time
hydro_time_step = p.dt_export
mesh_file = p.mesh_file
if "with_rivers" in mesh_file and close_rivers:
    mesh_file = mesh_file.replace("_with_rivers", "")
mesh_proj = p.mesh_proj
open_tags = p.open_tags
bath_file = (
    p.prepro_dir + "/bathymetry_smooth.nc",
    "bathymetry",
)

output_dir = (
    f"{hydro_dir}/LPT_sediments/3D/{simu_name}_{number_particles}/{min_size}_{max_size}"
)
print(f"Info \t: Exporting in {output_dir}")
stats_dir = f"{output_dir}/stats/"
raster_dir = f"{output_dir}/raster/"
os.makedirs(stats_dir, exist_ok=True)
os.makedirs(raster_dir, exist_ok=True)

dt_lpt = 200.0  # time step for particle tracker
dt_seed = 3600.0  # time step between two releases
dt_print = 3600.0  # time step between two solution prints
dt_export = 86400.0  # time step between two full exports

release_shp = p.local_base_dir + "/spatial/sediments_release_polygons.gpkg"
zoi_shp = p.local_base_dir + "/spatial/SMA_dugongs_simplified.gpkg"

# export grid parameters
minx = 114000  # utm 56s coordinates
maxx = 230000
miny = 7491000
maxy = 7640000

resol = 100  # resolution in meters

nx = int(np.ceil((maxx - minx) / resol))
ny = int(np.ceil((maxy - miny) / resol))

# Ensure compatibility between bounds and resolution
maxx += resol - (maxx - minx) % resol
maxy += resol - (maxy - miny) % resol
# ------------------------------------------------------------------------------

# ----- Functions --------------------------------------------------------------
def seed_particles(groups, tracker, n_sites):
    """add particles to the simu
    arguments:
        * groups       : dgpy.dgGroupCollection object
        * tracker      : slimParticle.Tracker object
        * n_sites      : list of site ids

    """
    print("Info\t: seeding particles...")
    particles = dgpy.slimParticleTracker2D(groups)

    gdf_points = LPT_functions.gdf_seeding_points(particles_per_timestep, gdf_poly)
    xx = [p.x for p in gdf_points.geometry]
    yy = [p.y for p in gdf_points.geometry]
    allx = np.column_stack([xx, yy])

    particles.add_particles(allx, True)
    pos = particles.position()
    p = gdf_points.ID.values
    n_points = p.size

    D, rho_s = sf.get_sed_size_density(min_size, max_size, n_points)

    Dgr = sf.get_Dgr(rho_w, rho_s, nu, D)
    w_s = sf.get_ws(Dgr, D, nu)
    z_nb = w_s * dt_lpt

    # Initialize fields
    data = pd.DataFrame(
        data={
            "origin": p,  # origin
            "z": np.repeat(0, n_points),
            "h": np.repeat(0, n_points),
            "release_time": np.repeat(rk.time, n_points),
            "reach_zoi": np.repeat(0, n_points),
            "sedimented": np.repeat(0, n_points),
            "sedimented_zoi": np.repeat(0, n_points),
            "time_reach_zoi": np.repeat(0, n_points),
            "time_sedimented_zoi": np.repeat(0, n_points),
            "duration_sedimented": np.repeat(0, n_points),
            "number_times_sedimented": np.repeat(0, n_points),
            "D": D,
            "rho_s": rho_s,
            "Dgr": Dgr,
            "z_nb": z_nb,
            "w_s": w_s,
        }
    )

    tracker.add_particles(pos, data.values, 1, True)
    toc = time.time()
    print("Info\t: %5i particles seeded in %f seconds" % (pos.shape[0], toc - tic))
    S = np.array([np.size(p[p == k + 1]) for k in range(n_sites)])
    return S


def get_velocity_profile(u2d, z):
    """Compute the horizontal velocity profile from the depth-averaged velocity using a logarithmic law."""
    hydro.load(tracker, rk.time)

    h = hydro.H.ravel()
    D = tracker.field("D")[:, 0]
    z_nb = tracker.field("z_nb")[:, 0]
    rho_s = tracker.field("rho_s")[:, 0]
    Dgr = tracker.field("Dgr")[:, 0]
    w_s = tracker.field("w_s")[:, 0]

    s = rho_s / rho_w  # TODO check
    zPTM = np.minimum(np.maximum(h[:] - z[:], 0.0), h[:])
    che_skin = sf.get_chezy_coeff(h, 3 * D)
    ust_skin = sf.get_bed_shear_velocity(u2d, che_skin)
    u_log = ust_skin * (5.75 * np.log10(zPTM / (3 * D)) + 8.5)
    u_log[zPTM < 3 * D] = 0.0
    u_znb = ust_skin * (5.75 * np.log10(z_nb / (3 * D)) + 8.5)
    u_znb[z_nb < 3 * D] = 0.0
    tau = sf.get_skin_roughness_stress(u_log, h, rho_w, D)
    theta = sf.get_theta(tau, rho_w, rho_s, D)
    # tau_crit = sf.get_skin_roughness_stress(tracker.field("u_crit")[:, 0], h, rho_w, D90) #TODO use u_crit ?
    theta_crit = sf.get_thetac(Dgr)

    c_delta = np.minimum(
        z_nb
        / 1.4
        / w_s
        / D
        * 3.3e-4
        * np.maximum((theta - theta_crit) / theta_crit, 0.0) ** 1.5
        * ((s - 1) * g * D ** 3 / nu ** 2) ** 0.1
        * np.sqrt((s - 1) * g * D),
        1,
    )

    u3d = np.copy(u_log)

    index = zPTM <= z_nb
    u3d[index] = c_delta[index] * u_log[index] + (zPTM[index] - (z_nb[index] / 1.4)) / (
        0.4 / 1.4 * z_nb[index]
    ) * (u_znb[index] - c_delta[index] * u_log[index])

    index = zPTM <= z_nb / 1.4
    u3d[index] = c_delta[index] * u_log[index]

    return u3d


def compute_dx():
    hydro.load(tracker, rk.time)
    h = hydro.H.ravel()
    D = tracker.field("D")[:, 0]
    rho_s = tracker.field("rho_s")[:, 0]
    Dgr = tracker.field("Dgr")[:, 0]

    U2D = np.maximum(np.hypot(hydro.uv[:, 0], hydro.uv[:, 1])[:, None], 1e-6)
    u_dir = hydro.uv / U2D
    tau_skin = sf.get_skin_roughness_stress(U2D, h, rho_w, D)
    che_skin = sf.get_chezy_coeff(h, 3 * D)
    ust_skin = sf.get_bed_shear_velocity(U2D, che_skin)
    theta = sf.get_theta(np.ravel(tau_skin), rho_w, rho_s, D)
    thetac = sf.get_thetac(Dgr)
    M = theta / thetac
    M[M < 0] = 0
    ub = sf.get_ub(np.ravel(ust_skin), M)
    ul = np.copy(ub)

    z = tracker.field("z")[:, 0]
    z[z > h] = h[z > h]
    U3D = get_velocity_profile(U2D.reshape(-1), z)
    index = (h - z) > 3 * D
    ul[index] = U3D[index]

    dx = np.reshape(ul, (-1, 1)) * u_dir * rk.dt
    return dx, M, z


# ------------------------------------------------------------------------------


# ----- Initialize simu --------------------------------------------------------
rho_w = 1034.0
g = 9.81

init_release = slimParticle.parse_time(release_initial_time)
end_release = slimParticle.parse_time(release_final_time)
init = slimParticle.parse_time(LPT_initial_time)
end = slimParticle.parse_time(LPT_final_time)
particles_per_timestep = int(number_particles * dt_seed / (end_release - init_release))

mesh = slimPre.Mesh(mesh_file, mesh_proj)
if not os.path.isfile(p.prepro_dir + "/okubo_map.msh"):
    slimPre.create_okubo_map(mesh, p.prepro_dir + "/okubo_map", 2e-4)
tracker = slimParticle.Tracker(
    mesh,
    [
        "origin",
        "z",
        "h",
        "release_time",
        "reach_zoi",
        "sedimented",
        "sedimented_zoi",
        "time_reach_zoi",
        "time_sedimented_zoi",
        "duration_sedimented",
        "number_times_sedimented",
        "D",
        "rho_s",
        "Dgr",
        "z_nb",
        "w_s",
    ],
)
tracker.set_boundary_open(open_tags)
hydro = slimParticle.HydroReader(
    mesh, hydro_dir, bath_file, hydro_initial_time, hydro_time_step, input_type="abin"
)

print("Info\t: Intersecting mesh and polygons...")
# load sites and generate seeding points inside them
gdf_poly = gpd.read_file(release_shp)
n_sites = gdf_poly.shape[0]
gdf_poly.to_crs(mesh_proj, inplace=True)
print(f"-> found {n_sites} sediment dumping sites")

loader_kappa = slimParticle.PreLoader(
    mesh, [p.prepro_dir + "/okubo_map.msh"], compute_gradient=True
)
rk = slimParticle.ERK44(LPT_initial_time, LPT_final_time, dt_lpt)
zoi_polys = LPT_functions.initialize_zoi_map(mesh_file, mesh_proj, zoi_shp, "OBJECTID")

print("==== START SIMULATION ====")
i_iter = 1
i_export = 0
n_export = int((end - init) / dt_print) - 1
n_iter = int((end - init) / dt_lpt)
n_out = 0
S = 0
tic_begin = toc = time.time()

# Initialize export matrix and arrays
export_stats = np.empty((n_export + 1, 7, n_sites))  # +1 because first export is "0"
timesteps = []

raster_grid = np.zeros((ny, nx), dtype=int)
floating_inside_grid = np.copy(raster_grid)
floating_outside_grid = np.copy(raster_grid)
sedimented_inside_grid = np.copy(raster_grid)
sedimented_outside_grid = np.copy(raster_grid)
# ------------------------------------------------------------------------------

# ----- Loop -------------------------------------------------------------------
while not rk.end:
    tic = toc
    # seed particles
    if (
        ((rk.time - init) % dt_seed == 0)
        and (rk.time >= init_release)
        and (rk.time < end_release)
    ):
        S += seed_particles(mesh._groups, tracker, n_sites)
        print("Info\t: Updating number of seeded particles...")

    # RK iteration
    for irk in range(rk.nsub):
        dx, M, z = compute_dx()
        rk.sub_time_step(tracker, dx)

    hydro.load(tracker, rk.time)
    h = hydro.H.ravel()
    U2D = np.maximum(np.hypot(hydro.uv[:, 0], hydro.uv[:, 1])[:, None], 1e-6)

    z = tracker.field("z")[:, 0]
    w_s = tracker.field("w_s")[:, 0]
    D = tracker.field("D")[:, 0]
    rho_s = tracker.field("rho_s")[:, 0]
    z_nb = tracker.field("z_nb")[:, 0]

    U3D = get_velocity_profile(U2D.reshape(-1), z)

    Rn = np.random.standard_normal(size=np.shape(z))
    Kv = 8.59e-3 * U3D * z * z * (h - z) * (h - z) / h / h / h
    dz = w_s * rk.dt + (Rn - np.mean(Rn)) / np.sqrt(np.var(Rn)) * np.sqrt(
        2 * Kv * rk.dt
    )

    tau = sf.get_skin_roughness_stress(U3D, h, rho_w, D)
    theta = sf.get_theta(tau, rho_w, rho_s, D)
    Dgr = sf.get_Dgr(rho_w, rho_s, nu, D)
    theta_crit = sf.get_thetac(Dgr)

    M = theta / theta_crit

    for i in range(len(U2D)):
        # remove motion for particle at the bottom
        if z[i] >= h[i]:
            tracker.field("z")[i, 0] = h[i]
            if M[i] >= 1.0 and sf.resuspensionProbability(theta, theta_crit):
                dz[i] = -z_nb[i]
            else:
                dz[i] = 0.0
        elif z[i] + dz[i] >= h[i]:
            tracker.field("z")[i, 0] = h[i]
            dz[i] = 0.0
        elif z[i] + dz[i] < 0:
            tracker.field("z")[i, 0] = 0
            dz[i] = 0

    tracker.field("z")[:, 0] += dz[:]

    # Update fields and compute stats
    tracker.field("h")[:, 0] = h

    kappa = loader_kappa.load(tracker, rk.time)
    rk.end_time_step(tracker, hydro.diffusivity(kappa, rk.dt))
    on_boundary = tracker.is_on_open_boundary()
    rm = np.logical_or(on_boundary, np.isnan(M[:]))
    n_out_local = np.sum(rm)
    if n_out_local > 0:
        tracker.remove_particles(rm)
        n_out += n_out_local

    posx = tracker.position[:, 0]
    posy = tracker.position[:, 1]

    on_zoi = (
        np.asarray(tracker._tracker.in_polygon(zoi_polys)) >= 0
    )  # boolean (True if over zoi)

    sedimented_before = tracker.field("sedimented")[:, 0] == 1

    z_PTM = np.minimum(
        np.maximum(tracker.field("h")[:, 0] - tracker.field("z")[:, 0], 0.0),
        tracker.field("h")[:, 0][:],
    )
    sedimented = np.ravel(
        z_PTM < tracker.field("z_nb")[:, 0]
    )  # boolean (True if sedimented)

    tracker.field("sedimented")[sedimented] = 1
    tracker.field("sedimented")[np.logical_not(sedimented)] = 0

    tracker.field("duration_sedimented")[
        sedimented
    ] += dt_lpt  # update field (add timestep for all sedimented particles)
    tracker.field("number_times_sedimented")[
        sedimented & np.logical_not(sedimented_before)
    ] += 1  # update field (add one for new sedimentation)

    inside = tracker.field("origin")[:, 0] <= 4
    outside = tracker.field("origin")[:, 0] >= 5

    floating_inside_grid += LPT_functions.points_on_grid(
        posx[inside], posy[inside], nx, ny, minx, maxx, miny, maxy
    )
    floating_outside_grid += LPT_functions.points_on_grid(
        posx[outside], posy[outside], nx, ny, minx, maxx, miny, maxy
    )
    sedimented_inside_grid += LPT_functions.points_on_grid(
        posx[inside & sedimented],
        posy[inside & sedimented],
        nx,
        ny,
        minx,
        maxx,
        miny,
        maxy,
    )
    sedimented_outside_grid += LPT_functions.points_on_grid(
        posx[outside & sedimented],
        posy[outside & sedimented],
        nx,
        ny,
        minx,
        maxx,
        miny,
        maxy,
    )

    didnt_reached_zoi_before = np.ravel(
        tracker.field("reach_zoi")[:] != 1
    )  # boolean (True if never reached zoi before current timestep)
    didnt_sedimented_zoi_before = np.ravel(
        tracker.field("sedimented_zoi")[:] != 1
    )  # boolean (True if never sedimented on zoi before current timestep)

    tracker.field("reach_zoi")[on_zoi] = 1  # update field (1 if ever reached zoi)
    tracker.field("sedimented_zoi")[
        on_zoi & sedimented
    ] = 1  # update field (1 if ever sedimented on zoi)

    reached_zoi = np.ravel(
        tracker.field("reach_zoi")[:] == 1
    )  # boolean (True if ever reached zoi)
    sedimented_zoi = np.ravel(
        tracker.field("sedimented_zoi")[:] == 1
    )  # boolean (True if ever sedimented on zoi)

    tracker.field("time_reach_zoi")[
        didnt_reached_zoi_before & reached_zoi
    ] = rk.time  # update field (time first reach zoi)
    tracker.field("time_sedimented_zoi")[
        didnt_sedimented_zoi_before & sedimented_zoi
    ] = rk.time  # update field (time first sedimented on zoi)

    i_iter += 1

    # save result
    if (rk.time - init) % dt_print == 0:
        # tracker.write_abin(output_dir, rk.time, i_export)
        tracker.write_vtp(output_dir, rk.time, i_export)

        if i_export >= 0:
            for i in range(n_sites):
                from_site = np.ravel(tracker.field("origin")[:] == i + 1)

                n_from_site = np.count_nonzero(from_site)
                n_sedimented = np.count_nonzero(from_site & sedimented)

                n_sedimented_on_zoi = np.count_nonzero(from_site & sedimented & on_zoi)
                n_reached_zoi = np.count_nonzero(from_site & on_zoi)
                n_on_zoi = n_sedimented_on_zoi + n_reached_zoi
                n_ever_reached_zoi = np.count_nonzero(from_site & reached_zoi)
                n_ever_sedimented_on_zoi = np.count_nonzero(from_site & sedimented_zoi)

                export_stats[i_export, :, i] = [
                    n_from_site,
                    n_sedimented,
                    n_sedimented_on_zoi,
                    n_reached_zoi,
                    n_on_zoi,
                    n_ever_reached_zoi,
                    n_ever_sedimented_on_zoi,
                ]

            timesteps.append(rk.time)

        toc = time.time()

        print(
            "EXPORT %4i/%4i | %s | Current pace: %.2f s/it | Elapsed: %s | Total: %i | Sedimented: %i"
            % (
                i_export,
                n_export,
                slimPre.slim_private._format_time(rk.time),
                np.round(toc - tic, 2),
                str(datetime.timedelta(seconds=(toc - tic_begin))),
                len(posx),
                np.count_nonzero(sedimented),
            )
        )
        i_export += 1

    if (rk.time - init) % dt_export == 0:
        full_time_reach_sedimented = pd.DataFrame(
            data={
                "release_site": np.ravel(tracker.field("origin")[:]),
                "time_reach": np.ravel(tracker.field("time_reach_zoi")[:]),
                "time_sedimented": np.ravel(tracker.field("time_sedimented_zoi")[:]),
            }
        )
        full_time_reach_sedimented.to_csv(
            f"{stats_dir}full_reach_sedimented_{int(rk.time)}.csv"
        )

        # Export position tiff rasters
        LPT_functions.export_tiff(
            f"{raster_dir}floating_inside_cumul_{min_size}_{max_size}_{int(rk.time)}.tiff",
            minx,
            miny,
            resol,
            resol,
            floating_inside_grid,
        )
        LPT_functions.export_tiff(
            f"{raster_dir}floating_outside_cumul_{min_size}_{max_size}_{int(rk.time)}.tiff",
            minx,
            miny,
            resol,
            resol,
            floating_outside_grid,
        )
        LPT_functions.export_tiff(
            f"{raster_dir}sedimented_inside_cumul_{min_size}_{max_size}_{int(rk.time)}.tiff",
            minx,
            miny,
            resol,
            resol,
            sedimented_inside_grid,
        )
        LPT_functions.export_tiff(
            f"{raster_dir}sedimented_outside_cumul_{min_size}_{max_size}_{int(rk.time)}.tiff",
            minx,
            miny,
            resol,
            resol,
            sedimented_outside_grid,
        )

# Export final dataframes
print("Export final dataframes")
for i in range(n_sites):
    stats_df = pd.DataFrame(
        data=export_stats[:i_export, :, i],
        index=timesteps,
        columns=[
            "n_from_site",
            "n_sedimented",
            "n_sedimented_on_zoi",
            "n_reached_zoi",
            "n_on_zoi",
            "n_ever_reached_zoi",
            "n_ever_sedimented_on_zoi",
        ],
    )
    stats_df.index = pd.to_datetime(stats_df.index, unit="s", utc=True)
    stats_df.to_csv(f"{stats_dir}/stats_{min_size}_{max_size}_release_point_{i}.csv")

# Sedimation duration stats
duration = tracker.field("duration_sedimented")[:, 0]
nb_times = tracker.field("number_times_sedimented")[:, 0]
mean_duration = np.nan_to_num(duration / nb_times)
percentage_duration = duration / (rk.time - tracker.field("release_time")[:, 0]) * 100

sedimentation_df = pd.DataFrame(
    {
        "nb_time_sed": [
            np.mean(nb_times[inside]),
            np.mean(nb_times[outside]),
        ],
        "total_duration_sed": [
            np.mean(duration[inside]),
            np.mean(duration[outside]),
        ],
        "mean_duration_sed": [
            np.mean(mean_duration[inside]),
            np.mean(mean_duration[outside]),
        ],
        "percentage_time_sed": [
            np.mean(percentage_duration[inside]),
            np.mean(percentage_duration[outside]),
        ],
    },
    index=["Inside", "Outside"],
)

sedimentation_df.to_csv(f"{stats_dir}/stats_{min_size}_{max_size}_sedimentation.csv")