3D : GPU Custom Computations

Note that these features are still experimental.

#!/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, vdisp, *args):
        pgt.netCDFLoader3D.__init__(self, *args)
        self.vdisp = vdisp

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

    def constants(self):
        return {
            "first": np.array([42 for _ in range(1024)]),
            "second": np.array([42 for _ in range(1024)]),
            "third": np.array([42 for _ in range(1024)]),
            "fourth": np.array([42 for _ in range(1024)])
        }

    def additionalVolumes(self):
        return {
            "vdisp": self.vdisp
        }


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

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

    # additional volume data
    ds = Dataset(root + "matt_iso_3000_rh3000_u20_stat.nc", 'r')
    tracer = ds.variables['stat'][1, 1]
    zz = np.array(list(range(3000, 6000, 25)))
    vdisp = np.empty(tracer.shape)
    for iz in range(zz.shape[0]):
        vdisp[iz, :, :] = zz[iz] - tracer[iz, :, :]
    vdisp -= 3000
    vol = vedo.Volume(vdisp).isosurface([570]).alpha(0.4)
    vol.c('gray').scale(25)

    # simulation properties
    settings = pg.IntegratorSettings()
    settings.Space = pg.Space3D
    settings.MinimumAliveParticles = 0
    settings.WindfieldMode = pg.WindfieldMode.Dynamic
    settings.Integrator = pg.Integrator.ClassicRungeKutta
    settings.DeltaT = 1
    settings.DataTimeDistance = 60
    settings.MaximumSimulationTime = 1000
    settings.GridScale = [25, 25, 25]
    settings.SaveInterval = 1
    settings.UpLiftMode = pg.UpLiftMode.Constant
    settings.AdditionalVolumeMode = pg.AdditionalVolumeMode.Constant
    settings.ConstantsMode = pg.ConstantsMode.Constant

    settings.AdditionalCompute = [
        "first * 2",
        "second * 3",
        "third * 4",
        "fourth * 5 + vdisp"
    ]

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

    print(result_set.computeInfo(0))
    print(result_set.computeInfo(1))
    print(result_set.computeInfo(2))
    print(result_set.computeInfo(3))

    # 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, vol, axes=1)