3D : Lift

#!/usr/bin/env python3

import os
import vedo

import numpy as np
import pygranite as pg # trajectory library
import pgt # pygranite toolkit

from netCDF4 import Dataset

class LiftLoader(pgt.netCDFLoader3D):
    def __init__(self, *args, **kwargs):
        pgt.netCDFLoader3D.__init__(self, *args, **kwargs)

    def uplift(self):
        return np.array([[0, 0, np.random.rand(1)[0] * .5] for _ in range(1024)])


if __name__ == "__main__":
    root = os.path.abspath(os.path.dirname(__file__))

    # initialize windfield loader
    directory = "fields_scaled/"
    loader = LiftLoader([
        "matt_ref_u20_shear_compr_lt30_tot90_uinnew_uvwtavchmstat1.nc"
    ], u_key="utav", v_key="vtav", w_key="wtav")
    loader.PrintProgress = True

    # simulation properties
    settings = pg.IntegratorSettings()
    settings.Space = pg.Space3D
    settings.MinimumAliveParticles = 1
    settings.WindfieldMode = pg.WindfieldMode.Dynamic
    settings.Integrator = pg.Integrator.ClassicRungeKutta
    settings.DeltaT = 1
    settings.DataTimeDistance = 1
    settings.MaximumSimulationTime = 16000
    settings.GridScale = [25, 25, 25]
    settings.SaveInterval = 4
    settings.UpLiftMode = pg.UpLiftMode.Constant
    settings.AbortMode = pg.AbortMode.Length
    settings.MaximumLength = 500

    # generate start positions
    spawn_min = np.array([100, 2500, 150])
    spawn_max = np.array([150, 2600, 170])

    inp = np.random.rand(1024, 3)
    for i in range(inp.shape[0]):
        inp[i] = ((spawn_max - spawn_min) * inp[i] + spawn_min)

    start_set = pg.TrajectorySet(inp)

    # computation
    integrator = pg.TrajectoryIntegrator(settings, loader, start_set)
    # with pgt.Timer():
    result_set: pg.TrajectorySet = integrator.compute()

    traj = result_set.trajectories()
    traj_start = traj[:,0, 0]
    for i, p in enumerate(inp):
        print(p[0], traj_start[i])

    # visualization using vedo
    points = result_set.cloud()  # .trajectory(0) # .cloud()
    colors = [
        i / result_set.numberTrajectories()
        for i in range(result_set.numberTrajectories())
        for _ in range(result_set.lengthTrajectories())
    ]

    pts = vedo.Points(points).legend("trajectory points") \
        .cmap("viridis", input_array=np.array(colors))

    domain = [
        settings.GridScale[0] * 320,
        settings.GridScale[1] * 192,
        settings.GridScale[2] * 120
    ]
    world = vedo.Box(list(map(lambda x: x / 2, domain)),  # box origin = box center
                     size=tuple(domain)).wireframe().color([1, 0, 0]).lineWidth(1).flat()

    plt = vedo.show(world, pts, axes=1)