Skip to content

2D : Visualization

Ronja Schnur edited this page Jan 17, 2023 · 3 revisions

This example shows the usage of pygranite in a 2D scenario - specifically the computation of trajectories arround the globe.

#!/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
    loader = pgt.netCDFLoader2D(["testdata_2020april0115.nc"])
    loader.PrintProgress = False
    
    radius = 6371e3

    # simulation properties
    settings = pg.IntegratorSettings()
    settings.Space = pg.Space2D
    settings.MinimumAliveParticles = 1
    settings.WindfieldMode = pg.WindfieldMode.Dynamic
    settings.Integrator = pg.Integrator.ExplicitEuler
    settings.DeltaT = 100
    settings.DataTimeDistance = 3 * 60 * 60
    settings.MaximumSimulationTime = 3 * 60 * 60 * 5 * 8
    settings.GridScale = [loader.ScaleU, loader.ScaleV]
    settings.Offset = [loader.OffsetU, loader.OffsetV]
    settings.SphereRadius = radius
    settings.BorderMode = pg.BorderMode.LoopXY
    settings.SaveInterval = 1

    start_points = [
        [np.random.rand(1)[0] * 360, np.random.rand(1)[0] * 180 - 90] for _ in range(100)
    ]
    start_set = pg.TrajectorySet(np.array(start_points))

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

    # visualization using vedo

    # get points and transform them to cartesian space
    # (-90, as this p[1] goes from -90 to 90 but 0-180 is expected)
    points = result_set.cloud() 
    points = np.array([
        [
            radius * np.cos(np.pi/180 * p[0]) * np.sin(np.pi/180 * (p[1] - 90)),
            radius * np.sin(np.pi/180 * p[0]) * np.sin(np.pi/180 * (p[1] - 90)),
            radius * np.cos(np.pi/180 * (p[1] - 90))
        ] for p in points
    ])


    # assign color values
    colors = [
        i / result_set.numberTrajectories()
        for i in range(result_set.numberTrajectories())
        for _ in range(result_set.lengthTrajectories())
    ]

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

    world = vedo.Earth(r =radius)
    vedo.show(pts, world, axes=1)