3D : Topography Test

#!/usr/bin/env python3

import os
import vedo

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

def heightmap(csvf):
    return np.genfromtxt(csvf, delimiter=",")

def buildPoints(hm, scale):
    pts = []
    for x in range(hm.shape[1]):
        for y in range(hm.shape[0]):
            pts.append([x, y, hm[y][x] / scale[2]])
    return np.array(pts)

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

    # initialize windfield loader
    directory = root + "fields_scaled/"
    loader = pgt.netCDFLoader3D([
        directory + f"{i}.nc" for i in range(len(os.listdir(directory)))
    ])
    loader.PrintProgress = True

    # simulation properties
    settings = pg.IntegratorSettings()
    settings.Space = pg.Space3D
    settings.MinimumAliveParticles = 16
    settings.WindfieldMode = pg.WindfieldMode.Dynamic
    settings.Integrator = pg.Integrator.ClassicRungeKutta
    settings.DeltaT = 1 
    settings.DataTimeDistance = 60
    settings.MaximumSimulationTime = 6000
    settings.GridScale = [25, 25, 25]
    settings.SaveInterval = 1

    # load heightmap
    hm = heightmap(root + "pyr_1000.csv")
    topography_v = vedo.Points(buildPoints(hm, settings.GridScale), r=.5).scale(settings.GridScale)

    settings.Topography = hm 

    # generate start positions
    spawn_min = np.array([100, 2500, 150])
    spawn_max = np.array([500, 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()

    # 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()

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