Spaces:
Sleeping
Sleeping
| # Copyright (c) 2022 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. | |
| import warp as wp | |
| from .model import PARTICLE_FLAG_ACTIVE | |
| def particle_force(n: wp.vec3, v: wp.vec3, c: float, k_n: float, k_d: float, k_f: float, k_mu: float): | |
| vn = wp.dot(n, v) | |
| jn = c * k_n | |
| jd = min(vn, 0.0) * k_d | |
| # contact force | |
| fn = jn + jd | |
| # friction force | |
| vt = v - n * vn | |
| vs = wp.length(vt) | |
| if vs > 0.0: | |
| vt = vt / vs | |
| # Coulomb condition | |
| ft = wp.min(vs * k_f, k_mu * wp.abs(fn)) | |
| # total force | |
| return -n * fn - vt * ft | |
| def eval_particle_forces_kernel( | |
| grid: wp.uint64, | |
| particle_x: wp.array(dtype=wp.vec3), | |
| particle_v: wp.array(dtype=wp.vec3), | |
| particle_radius: wp.array(dtype=float), | |
| particle_flags: wp.array(dtype=wp.uint32), | |
| k_contact: float, | |
| k_damp: float, | |
| k_friction: float, | |
| k_mu: float, | |
| k_cohesion: float, | |
| max_radius: float, | |
| # outputs | |
| particle_f: wp.array(dtype=wp.vec3), | |
| ): | |
| tid = wp.tid() | |
| # order threads by cell | |
| i = wp.hash_grid_point_id(grid, tid) | |
| if i == -1: | |
| # hash grid has not been built yet | |
| # if tid == 0: | |
| # print("hash grid has not been built yet") | |
| return | |
| if (particle_flags[i] & PARTICLE_FLAG_ACTIVE) == 0: | |
| return | |
| x = particle_x[i] | |
| v = particle_v[i] | |
| radius = particle_radius[i] | |
| f = wp.vec3() | |
| # particle contact | |
| query = wp.hash_grid_query(grid, x, radius + max_radius + k_cohesion) | |
| index = int(0) | |
| count = int(0) | |
| while wp.hash_grid_query_next(query, index): | |
| if (particle_flags[index] & PARTICLE_FLAG_ACTIVE) != 0 and index != i: | |
| # compute distance to point | |
| n = x - particle_x[index] | |
| d = wp.length(n) | |
| err = d - radius - particle_radius[index] | |
| count += 1 | |
| if err <= k_cohesion: | |
| n = n / d | |
| vrel = v - particle_v[index] | |
| f = f + particle_force(n, vrel, err, k_contact, k_damp, k_friction, k_mu) | |
| particle_f[i] = f | |
| def eval_particle_forces(model, state, forces): | |
| if model.particle_max_radius > 0.0: | |
| wp.launch( | |
| kernel=eval_particle_forces_kernel, | |
| dim=model.particle_count, | |
| inputs=[ | |
| model.particle_grid.id, | |
| state.particle_q, | |
| state.particle_qd, | |
| model.particle_radius, | |
| model.particle_flags, | |
| model.particle_ke, | |
| model.particle_kd, | |
| model.particle_kf, | |
| model.particle_mu, | |
| model.particle_cohesion, | |
| model.particle_max_radius, | |
| ], | |
| outputs=[forces], | |
| device=model.device, | |
| ) | |