Spaces:
Sleeping
Sleeping
GarmentCode / NvidiaWarp-GarmentCode /exts /omni.warp /omni /warp /nodes /_impl /OgnSamplePrimFlocking.py
| # Copyright (c) 2023 NVIDIA CORPORATION. All rights reserved. | |
| # NVIDIA CORPORATION and its licensors retain all intellectual property | |
| # and proprietary rights in and to this software, related documentation | |
| # and any modifications thereto. Any use, reproduction, disclosure or | |
| # distribution of this software and related documentation without an express | |
| # license agreement from NVIDIA CORPORATION is strictly prohibited. | |
| """Sample node that simulates flocking behaviors by animating prim attributes.""" | |
| import math | |
| import traceback | |
| import numpy as np | |
| import carb.settings | |
| import omni.kit.app | |
| import omni.graph.core as og | |
| import omni.usd | |
| import usdrt | |
| import warp as wp | |
| import omni.warp.nodes | |
| from omni.warp.nodes.ogn.OgnSamplePrimFlockingDatabase import OgnSamplePrimFlockingDatabase | |
| # device used for flocking simulation | |
| MAIN_DEVICE = "cuda:0" | |
| # device used for updating colors | |
| COLOR_DEVICE = "cpu" | |
| # Kernels | |
| # ----------------------------------------------------------------------------- | |
| class Boid: | |
| vel: wp.vec3f | |
| wander_angles: wp.vec2f | |
| mass: float | |
| group: int | |
| class Obstacle: | |
| pos: wp.vec3f | |
| radius: float | |
| class World: | |
| lower: wp.vec3f | |
| upper: wp.vec3f | |
| grid: wp.uint64 | |
| seed: int | |
| biases: wp.mat33f | |
| obstacles: wp.array(dtype=Obstacle) | |
| def copy_positions(dst: wp.array(dtype=wp.vec3f), src: wp.fabricarray(dtype=wp.vec3d)): | |
| tid = wp.tid() | |
| pos = src[tid] | |
| dst[tid] = wp.vec3f(float(pos[0]), float(pos[1]), float(pos[2])) | |
| def assign_colors( | |
| glows: wp.array(dtype=float), | |
| groups: wp.array(dtype=int), | |
| color_ramps: wp.array2d(dtype=wp.vec3f), | |
| colors: wp.fabricarrayarray(dtype=wp.vec3f) | |
| ): | |
| tid = wp.tid() | |
| glow = glows[tid] | |
| group = groups[tid] | |
| if glow < 0.4: | |
| alpha = glow / 0.4 | |
| colors[tid][0] = (1.0 - alpha) * color_ramps[group, 0] + alpha * color_ramps[group, 1] | |
| elif glow < 0.8: | |
| alpha = (glow - 0.4) / 0.4 | |
| colors[tid][0] = (1.0 - alpha) * color_ramps[group, 1] + alpha * color_ramps[group, 2] | |
| else: | |
| alpha = (glow - 0.8) / 0.2 | |
| colors[tid][0] = (1.0 - alpha) * color_ramps[group, 2] + alpha * color_ramps[group, 3] | |
| def intersect_ray_sphere(origin: wp.vec3f, dir: wp.vec3f, center: wp.vec3f, radius: float): | |
| to_sphere = center - origin | |
| tc = wp.dot(to_sphere, dir) | |
| if tc < 0.0: | |
| return tc | |
| d = wp.sqrt(wp.length_sq(to_sphere) - tc * tc) | |
| if d < 0.0: | |
| return -999999.0 | |
| ts = wp.sqrt(radius * radius - d * d) | |
| return tc - ts | |
| def boids( | |
| boids: wp.array(dtype=Boid), | |
| world: World, | |
| dt: float, | |
| positions: wp.fabricarray(dtype=wp.vec3d), | |
| orientations: wp.fabricarray(dtype=wp.quatf), | |
| glows: wp.array(dtype=float), | |
| ): | |
| tid = wp.tid() | |
| boid = boids[tid] | |
| old_pos = positions[tid] | |
| old_rot = orientations[tid] | |
| pos = wp.vec3(float(old_pos[0]), float(old_pos[1]), float(old_pos[2])) | |
| vel = boid.vel | |
| forward = wp.quat_rotate(old_rot, wp.vec3f(1.0, 0.0, 0.0)) | |
| force = wp.vec3f(0.0) | |
| # obstacle avoidance | |
| depenetration_force = 100.0 | |
| avoidance_dist = 20.0 | |
| avoidance_force = 20.0 | |
| obstacles = world.obstacles | |
| num_obstacles = obstacles.shape[0] | |
| for i in range(num_obstacles): | |
| obstacle = obstacles[i] | |
| to_obstacle = obstacle.pos - pos | |
| # use padded radius | |
| radius = obstacle.radius + 2.0 | |
| if wp.length(to_obstacle) < radius: | |
| # depenetration | |
| force += depenetration_force * wp.normalize(-to_obstacle) | |
| else: | |
| # avoidance | |
| t = intersect_ray_sphere(pos, forward, obstacle.pos, radius) | |
| if t > 0.0 and t < avoidance_dist: | |
| intersection_point = pos + t * forward | |
| out = intersection_point - obstacle.pos | |
| force += avoidance_force * (1.0 - t / avoidance_dist) * wp.normalize(out) | |
| # wander | |
| r = 10.0 | |
| s0 = wp.sin(boid.wander_angles[0]) | |
| c0 = wp.cos(boid.wander_angles[0]) | |
| s1 = wp.sin(boid.wander_angles[1]) | |
| c1 = wp.cos(boid.wander_angles[1]) | |
| p = wp.vec3f(r * s0 * s1, r * s0 * c1, r * c0) | |
| offset = r + 1.0 | |
| target = pos + wp.quat_rotate(old_rot, wp.vec3f(offset, 0.0, 0.0) + p) | |
| wander_force = 7.0 | |
| force += wander_force * wp.normalize(target - pos) | |
| state = wp.rand_init(world.seed, tid) | |
| angle0 = boid.wander_angles[0] + wp.pi * (0.1 - 0.2 * wp.randf(state)) | |
| angle1 = boid.wander_angles[1] + wp.pi * (0.1 - 0.2 * wp.randf(state)) | |
| boid.wander_angles = wp.vec2f(angle0, angle1) | |
| cohesion_radius = 15.0 | |
| cohesion_force = 20.0 | |
| separation_radius = 10.0 | |
| separation_force = 100.0 | |
| # flocking | |
| query = wp.hash_grid_query(world.grid, pos, cohesion_radius) | |
| num_neighbors = int(0) | |
| num_align_neighbors = int(0) | |
| num_cohesion_neighbors = float(0) | |
| num_decohesion_neighbors = float(0) | |
| cohesion_pos_sum = wp.vec3f(0.0) | |
| decohesion_pos_sum = wp.vec3f(0.0) | |
| vel_sum = wp.vec3f(0.0) | |
| for index in query: | |
| if index != tid: | |
| other = boids[index] | |
| other_pos64 = positions[index] | |
| other_pos = wp.vec3f(float(other_pos64[0]), float(other_pos64[1]), float(other_pos64[2])) | |
| dist = wp.length(pos - other_pos) | |
| if dist < cohesion_radius: | |
| to_other = wp.normalize(other_pos - pos) | |
| # separation | |
| if dist < separation_radius: | |
| force -= separation_force * (1.0 - dist / separation_radius) * to_other | |
| # cohesion | |
| bias = world.biases[boid.group, other.group] | |
| if bias > 0.0: | |
| cohesion_pos_sum += bias * other_pos | |
| num_cohesion_neighbors += bias | |
| else: | |
| decohesion_pos_sum -= bias * other_pos | |
| num_decohesion_neighbors -= bias | |
| # alignment | |
| if other.group == boid.group: | |
| vel_sum += bias * other.vel | |
| num_align_neighbors += 1 | |
| num_neighbors += 1 | |
| # align | |
| if num_align_neighbors > 0: | |
| vel_avg = vel_sum / float(num_align_neighbors) | |
| force += vel_avg - vel | |
| # cohere | |
| if num_cohesion_neighbors > 0.0: | |
| cohesion_pos_avg = cohesion_pos_sum / float(num_cohesion_neighbors) | |
| force += cohesion_force * wp.normalize(cohesion_pos_avg - pos) | |
| # decohere (group separation) | |
| if num_decohesion_neighbors > 0.0: | |
| decohesion_pos_avg = decohesion_pos_sum / float(num_decohesion_neighbors) | |
| force += cohesion_force * wp.normalize(pos - decohesion_pos_avg) | |
| # boundaries | |
| boundary_force = 20.0 | |
| if pos[0] >= world.upper[0]: | |
| force += wp.vec3f(-boundary_force, 0.0, 0.0) | |
| if pos[0] <= world.lower[0]: | |
| force += wp.vec3f(boundary_force, 0.0, 0.0) | |
| if pos[1] >= world.upper[1]: | |
| force += wp.vec3f(0.0, -0.5 * boundary_force, 0.0) | |
| if pos[1] <= world.lower[1]: | |
| force += wp.vec3f(0.0, 5.0 * boundary_force, 0.0) | |
| if pos[2] >= world.upper[2]: | |
| force += wp.vec3f(0.0, 0.0, -boundary_force) | |
| if pos[2] <= world.lower[2]: | |
| force += wp.vec3f(0.0, 0.0, boundary_force) | |
| vel += dt * force / boid.mass | |
| # clamp speed | |
| max_speed = 15.0 | |
| speed_sq = wp.length_sq(vel) | |
| if speed_sq > max_speed * max_speed: | |
| vel = max_speed * wp.normalize(vel) | |
| # update position | |
| pos += dt * vel | |
| positions[tid] = wp.vec3d(wp.float64(pos[0]), wp.float64(pos[1]), wp.float64(pos[2])) | |
| # update orientation | |
| dq = wp.quat_between_vectors(forward, vel) | |
| orientations[tid] = wp.normalize(dq * orientations[tid]) | |
| # save velocity | |
| boid.vel = vel | |
| boids[tid] = boid | |
| # update glow as an exponentially weighted moving average to keep it smooth | |
| glow = wp.min(1.0, float(num_neighbors) / 25.0) | |
| glow_alpha = 0.25 | |
| glows[tid] = glow_alpha * glow + (1.0 - glow_alpha) * glows[tid] | |
| # Internal State | |
| # ------------------------------------------------------------------------------ | |
| class InternalState: | |
| """Internal state for the node.""" | |
| def __init__(self) -> None: | |
| self.initialized = False | |
| def initialize(self, device): | |
| # requirement checks | |
| ext_mgr = omni.kit.app.get_app().get_extension_manager() | |
| # make sure USDRT is enabled | |
| usdrt_ext_name = "usdrt.scenegraph" | |
| if not ext_mgr.is_extension_enabled(usdrt_ext_name): | |
| raise RuntimeError(f"This sample requires the '{usdrt_ext_name}' extension to be enabled") | |
| # check USDRT version to make sure we have a working SelectPrims() | |
| usdrt_ext_id = ext_mgr.get_enabled_extension_id(usdrt_ext_name) | |
| usdrt_version_string = ext_mgr.get_extension_dict(usdrt_ext_id)["package"]["version"] | |
| usdrt_version = tuple(int(v) for v in usdrt_version_string.split(".")) | |
| if usdrt_version < (7, 3, 0): | |
| raise RuntimeError(f"USDRT version 7.3.0 is required, found {usdrt_version_string}. Please update to a newer version of Kit to run this sample.") | |
| # check if FSD is enabled | |
| settings = carb.settings.get_settings() | |
| is_fsd_enabled = settings.get_as_bool("/app/useFabricSceneDelegate") | |
| if not is_fsd_enabled: | |
| print("***") | |
| print("*** Flocking demo warning: The Fabric Scene Delegate is not enabled.") | |
| print("*** Some features, like color animation, may not work.") | |
| print("*** You can enable FSD in Preferences->Rendering.") | |
| print("***") | |
| stage_id = omni.usd.get_context().get_stage_id() | |
| usdrt_stage = usdrt.Usd.Stage.Attach(stage_id) | |
| # import to Fabric | |
| for prim in usdrt_stage.Traverse(): | |
| pass | |
| # set up for Fabric interop | |
| boid_root = usdrt_stage.GetPrimAtPath(usdrt.Sdf.Path("/World/Boids")) | |
| boid_prims = boid_root.GetChildren() | |
| for prim in boid_prims: | |
| pos = prim.GetAttribute("xformOp:translate").Get() | |
| prim.CreateAttribute("_worldPosition", usdrt.Sdf.ValueTypeNames.Double3, True).Set(pos) | |
| prim.CreateAttribute("_worldOrientation", usdrt.Sdf.ValueTypeNames.Quatf, True).Set(usdrt.Gf.Quatf(1, 0, 0, 0)) | |
| # create a custom tag for the boids (could use applied schema too) | |
| prim.CreateAttribute("BoidTag", usdrt.Sdf.ValueTypeNames.AppliedSchemaTypeTag, True) | |
| num_boids = len(boid_prims) | |
| self.stage = usdrt_stage | |
| self.require_schemas = ["BoidTag"] | |
| self.transform_attrs = [ | |
| (usdrt.Sdf.ValueTypeNames.Double3, "_worldPosition", usdrt.Usd.Access.ReadWrite), | |
| (usdrt.Sdf.ValueTypeNames.Quatf, "_worldOrientation", usdrt.Usd.Access.ReadWrite), | |
| ] | |
| self.color_attrs = [ | |
| (usdrt.Sdf.ValueTypeNames.Float3Array, "primvars:_emissive", usdrt.Usd.Access.ReadWrite), | |
| ] | |
| npboids = np.zeros(num_boids, dtype=Boid.numpy_dtype()) | |
| angles = math.pi - 2 * math.pi * np.random.rand(num_boids) | |
| vx = 20 * np.sin(angles) | |
| vz = 20 * np.cos(angles) | |
| npboids["vel"][:, 0] = vx | |
| npboids["vel"][:, 2] = vz | |
| npboids["wander_angles"][:, 0] = math.pi * np.random.rand(num_boids) | |
| npboids["wander_angles"][:, 1] = 2 * math.pi * np.random.rand(num_boids) | |
| min_mass = 1.0 | |
| max_mass = 2.0 | |
| npboids["mass"][:] = min_mass + (max_mass - min_mass) * np.random.rand(num_boids) | |
| # we can have up to 3 groups currently, but that can be easily extended | |
| self.num_groups = 2 | |
| npboids["group"] = np.random.randint(self.num_groups, size=num_boids) | |
| num_obstacles = 3 | |
| npobstacles = np.zeros(num_obstacles, dtype=Obstacle.numpy_dtype()) | |
| npobstacles["pos"][0] = (-20, 30, -40) | |
| npobstacles["radius"][0] = 40 | |
| npobstacles["pos"][1] = (90, 30, 30) | |
| npobstacles["radius"][1] = 30 | |
| npobstacles["pos"][2] = (-100, 30, 60) | |
| npobstacles["radius"][2] = 25 | |
| self.grid = wp.HashGrid(dim_x=32, dim_y=32, dim_z=32, device=device) | |
| biases = wp.mat33f(-1.0) | |
| for i in range(self.num_groups): | |
| biases[i, i] = 1.0 | |
| world = World() | |
| world.lower = (-120, 20, -90) | |
| world.upper = (120, 40, 90) | |
| world.grid = self.grid.id | |
| world.seed = 0 | |
| world.biases = biases | |
| world.obstacles = wp.array(npobstacles, dtype=Obstacle, device=device) | |
| self.world = world | |
| self.num_boids = num_boids | |
| self.boids = wp.array(npboids, dtype=Boid, device=device) | |
| # color ramps per group | |
| color_ramps = [ | |
| [[0.3, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 0.5, 0.0], [1.0, 1.0, 0.5]], | |
| [[0.0, 0.0, 0.3], [0.0, 0.0, 1.0], [0.0, 0.5, 1.0], [0.5, 1.0, 1.0]], | |
| [[0.0, 0.3, 0.0], [0.0, 1.0, 0.0], [0.0, 1.0, 0.5], [0.8, 1.0, 0.8]], | |
| ] | |
| # copy of positions used for updating the spatial grid | |
| self.positions = wp.zeros(num_boids, dtype=wp.vec3f, device=device) | |
| # color ramps are only used on the COLOR_DEVICE | |
| self.color_ramps_c = wp.array(color_ramps, dtype=wp.vec3f, device=COLOR_DEVICE) | |
| # keep a copy of group assignments on the COLOR_DEVICE | |
| self.groups_c = wp.array(npboids["group"], device=COLOR_DEVICE) | |
| # if we use different devices, the glow array must be copied on each update | |
| if COLOR_DEVICE == device: | |
| # use the same glow array on each device, no copying needed | |
| self.glows_c = wp.zeros(num_boids, dtype=float, device=device) | |
| self.glows_m = self.glows_c | |
| elif COLOR_DEVICE == "cpu" or device == "cpu": | |
| # use a pinned host array for async copying glows between devices | |
| glows_h = wp.zeros(num_boids, dtype=float, device="cpu", pinned=True) | |
| if COLOR_DEVICE == "cpu": | |
| self.glows_c = glows_h | |
| self.glows_m = wp.zeros_like(glows_h, device=device) | |
| else: | |
| self.glows_c = wp.zeros_like(glows_h, device=COLOR_DEVICE) | |
| self.glows_m = glows_h | |
| else: | |
| # two different CUDA devices | |
| self.glows_c = wp.zeros(num_boids, dtype=float, device=COLOR_DEVICE) | |
| self.glows_m = wp.zeros(num_boids, dtype=float, device=device) | |
| # ...but that's currently not supported in Kit | |
| raise ValueError("Multiple GPUs not supported yet") | |
| self.time = 0.0 | |
| self.min_group_think = 3.0 | |
| self.max_group_think = 10.0 | |
| self.next_group_think = self.min_group_think + (self.max_group_think - self.min_group_think) * np.random.rand() | |
| self.frameno = 0 | |
| self.initialized = True | |
| # Compute | |
| # ------------------------------------------------------------------------------ | |
| def compute(db: OgnSamplePrimFlockingDatabase) -> None: | |
| """Evaluates the node.""" | |
| state = db.internal_state | |
| device = wp.get_device() | |
| if not state.initialized: | |
| state.initialize(device) | |
| state.frameno += 1 | |
| # get transform attributes | |
| selection = state.stage.SelectPrims( | |
| require_applied_schemas=state.require_schemas, | |
| require_attrs=state.transform_attrs, | |
| device=str(device) | |
| ) | |
| fpos = wp.fabricarray(data=selection, attrib="_worldPosition") | |
| frot = wp.fabricarray(data=selection, attrib="_worldOrientation") | |
| # use fixed dt for stability | |
| dt = 1.0 / 60.0 | |
| state.time += dt | |
| # copy positions to a contiguous array and convert to vec3f so they can be used to update the spatial grid | |
| wp.launch(copy_positions, dim=state.num_boids, inputs=[state.positions, fpos]) | |
| # grid cell radius should be a bit bigger than query radius | |
| cell_radius = 20.0 | |
| state.grid.build(state.positions, cell_radius) | |
| state.world.seed = state.frameno | |
| # step the flocking simulation | |
| wp.launch(boids, dim=state.num_boids, inputs=[state.boids, state.world, dt, fpos, frot, state.glows_m]) | |
| # async copy from main device and remember the stream so we can sync later | |
| if COLOR_DEVICE != device: | |
| if device.is_cuda: | |
| work_stream = device.stream | |
| else: | |
| work_stream = wp.get_stream(COLOR_DEVICE) | |
| wp.copy(state.glows_c, state.glows_m, stream=work_stream) | |
| else: | |
| work_stream = None | |
| # get color attributes | |
| color_selection = state.stage.SelectPrims( | |
| require_applied_schemas=state.require_schemas, | |
| require_attrs=state.color_attrs, | |
| device=COLOR_DEVICE | |
| ) | |
| fcolor = wp.fabricarray(data=color_selection, attrib="primvars:_emissive") | |
| # occasionally update group biases (whether they are attracted or repelled from each other) | |
| if state.num_groups > 1 and state.time >= state.next_group_think: | |
| # pick two random groups | |
| group0 = np.random.randint(state.num_groups) | |
| group1 = np.random.randint(state.num_groups) | |
| while group0 == group1: | |
| group1 = np.random.randint(state.num_groups) | |
| # bias towards intra-group separation, but also allow attraction | |
| state.world.biases[group0, group1] = 1.0 - 5.0 * np.random.rand() | |
| state.world.biases[group1, group0] = 1.0 - 5.0 * np.random.rand() | |
| state.next_group_think += state.min_group_think + (state.max_group_think - state.min_group_think) * np.random.rand() | |
| if work_stream is not None: | |
| # wait for async GPU work to complete | |
| wp.synchronize_stream(work_stream) | |
| # update colors | |
| wp.launch(assign_colors, dim=state.num_boids, inputs=[state.glows_c, state.groups_c, state.color_ramps_c, fcolor], device=COLOR_DEVICE) | |
| # Node Entry Point | |
| # ------------------------------------------------------------------------------ | |
| class OgnSamplePrimFlocking: | |
| """Node.""" | |
| def internal_state() -> InternalState: | |
| return InternalState() | |
| def compute(db: OgnSamplePrimFlockingDatabase) -> None: | |
| device = wp.get_device(MAIN_DEVICE) | |
| try: | |
| with wp.ScopedDevice(device): | |
| compute(db) | |
| except Exception: | |
| db.log_error(traceback.format_exc()) | |
| return | |
| # Fire the execution for the downstream nodes. | |
| db.outputs.execOut = og.ExecutionAttributeState.ENABLED | |