3D : Reverse Computation

#!/usr/bin/env python3

import os
import vedo

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

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 = 600
    settings.GridScale = [25, 25, 25]
    settings.SaveInterval = 1
    settings.Reverse = False

    # 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)
    forward_set = integrator.compute() # contains starting points and computed onex

    print("Finished forward computation.")
    settings.Reverse = 1
    loader.IncrementAmount = -1
    loader._index -= 1

    integrator = pg.TrajectoryIntegrator(settings, loader, forward_set)
    with pgt.Timer():
        result_set = 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, axes=1)