gstaichi 0.1.18.dev1__cp310-cp310-win_amd64.whl

taichi/shaders/Circles_vk.vert

@@ -0,0 +1,45 @@
+#version 450
+
+layout(location = 0) in vec3 in_position;
+layout(location = 1) in vec3 in_normal;
+layout(location = 2) in vec2 in_texcoord;
+layout(location = 3) in vec4 in_color;
+
+layout(binding = 0) uniform UBO {
+  vec3 color;
+  int use_per_vertex_color;
+  int use_per_vertex_radius;
+  float radius;
+  float window_width;
+  float window_height;
+}
+ubo;
+
+layout(location = 1) out vec3 selected_color;
+layout(location = 2) out vec2 pos_2d;
+
+const vec2 offsets[6] = {
+  vec2(-1.0f, 1.0f),
+  vec2(1.0f, -1.0f),
+  vec2(-1.0f, -1.0f),
+  vec2(-1.0f, 1.0f),
+  vec2(1.0f, 1.0f),
+  vec2(1.0f, -1.0f),
+};
+
+void main() {
+  float x = in_position.x * 2.0 - 1.0;
+  float y = -(in_position.y * 2.0 - 1.0);
+
+  pos_2d = offsets[gl_VertexIndex % 6];
+  float radius_1d = ubo.use_per_vertex_radius == 0 ? ubo.radius : in_normal.x;
+
+  gl_Position = vec4(x, y, 0.0, 1.0);
+  gl_Position.xy += pos_2d * vec2(radius_1d / ubo.window_width * ubo.window_height, radius_1d) * 2.0;
+
+  if (ubo.use_per_vertex_color == 0) {
+    selected_color = ubo.color;
+  } else {
+    selected_color = in_color.rgb;
+  }
+}

taichi/shaders/Lines_vk.frag

@@ -0,0 +1,9 @@
+#version 450
+
+layout(location = 0) in vec3 color;
+
+layout(location = 0) out vec4 out_color;
+
+void main() {
+  out_color = vec4(color, 1);
+}

taichi/shaders/Lines_vk.vert

@@ -0,0 +1,11 @@
+#version 450
+
+layout(location = 0) in vec2 in_position;
+layout(location = 1) in uint in_color_encoded;
+
+layout(location = 0) out vec3 frag_color;
+
+void main() {
+  gl_Position = vec4(in_position.x * 2.0 - 1.0, -(in_position.y * 2.0 - 1.0), 0.0, 1.0);
+  frag_color = unpackUnorm4x8(in_color_encoded).rgb;
+}

taichi/shaders/Mesh_vk.frag

@@ -0,0 +1,71 @@
+#version 450
+
+layout(location = 0) in vec3 frag_pos;
+layout(location = 1) in vec3 frag_normal;
+layout(location = 2) in vec2 frag_texcoord;
+
+layout(location = 0) out vec4 out_color;
+
+struct SceneUBO {
+  vec3 camera_pos;
+  mat4 view;
+  mat4 projection;
+  vec3 ambient_light;
+  int point_light_count;
+};
+
+layout(binding = 0) uniform UBORenderable {
+  vec3 color;
+  int use_per_vertex_color;
+  int two_sided;
+  float has_attribute;
+}
+ubo_renderable;
+
+layout(binding = 1) uniform UBOScene {
+  SceneUBO scene;
+  float window_width;
+  float window_height;
+  float tan_half_fov;
+  float aspect_ratio;
+}
+ubo_scene;
+
+struct PointLight {
+  vec3 pos;
+  vec3 color;
+};
+
+layout(binding = 2, std430) buffer SSBO {
+  PointLight point_lights[];
+}
+ssbo;
+
+layout(location = 3) in vec4 selected_color;
+
+vec3 lambertian() {
+  vec3 ambient = ubo_scene.scene.ambient_light * selected_color.rgb;
+  vec3 result = ambient;
+
+  for (int i = 0; i < ubo_scene.scene.point_light_count; ++i) {
+    vec3 light_color = ssbo.point_lights[i].color;
+
+    vec3 light_dir = normalize(ssbo.point_lights[i].pos - frag_pos);
+    vec3 normal = normalize(frag_normal);
+    float factor = 0.0;
+    if(ubo_renderable.two_sided != 0){
+      factor = abs(dot(light_dir, normal));
+    }
+    else{
+      factor = max(dot(light_dir, normal), 0);
+    }
+    vec3 diffuse = factor * selected_color.rgb * light_color;
+    result += diffuse;
+  }
+
+  return result;
+}
+
+void main() {
+  out_color = vec4(lambertian(), selected_color.a);
+}

taichi/shaders/Mesh_vk.vert

@@ -0,0 +1,68 @@
+#version 450
+
+layout(location = 0) in vec3 in_position;
+layout(location = 1) in vec3 in_normal;
+layout(location = 2) in vec2 in_texcoord;
+layout(location = 3) in vec4 in_color;
+
+layout(location = 0) out vec3 frag_pos;
+layout(location = 1) out vec3 frag_normal;
+layout(location = 2) out vec2 frag_texcoord;
+layout(location = 3) out vec4 selected_color;
+
+struct SceneUBO {
+  vec3 camera_pos;
+  mat4 view;
+  mat4 projection;
+  vec3 ambient_light;
+  int point_light_count;
+};
+
+struct PointLight {
+  vec3 pos;
+  vec3 color;
+};
+
+struct MeshAttribute {
+  mat4 model;
+};
+
+layout(binding = 0) uniform UBORenderable {
+  vec3 color;
+  int use_per_vertex_color;
+  int two_sided;
+  float has_attribute;
+}
+ubo_renderable;
+
+layout(binding = 1) uniform UBOScene {
+  SceneUBO scene;
+  float window_width;
+  float window_height;
+  float tan_half_fov;
+  float aspect_ratio;
+}
+ubo_scene;
+
+layout(binding = 3, std430) buffer MeshAttributeBuffer {
+  MeshAttribute mesh_attr[];
+}
+mesh_attr_buffer;
+
+void main() {
+  mat4 model_tmp = transpose(mesh_attr_buffer.mesh_attr[gl_InstanceIndex].model);
+  // if mesh attributes not given, then use Identity as model matrix
+  mat4 model = mat4(1.0) * (1.0 - ubo_renderable.has_attribute) + model_tmp * ubo_renderable.has_attribute;
+
+  gl_Position = ubo_scene.scene.projection * ubo_scene.scene.view * model * vec4(in_position, 1.0);
+  gl_Position.y *= -1.0;
+  frag_texcoord = in_texcoord;
+  frag_pos = in_position;
+  frag_normal = in_normal;
+
+  if (ubo_renderable.use_per_vertex_color == 0) {
+    selected_color = vec4(ubo_renderable.color, 1.0);
+  } else {
+    selected_color = in_color;
+  }
+}

taichi/shaders/Particles_vk.frag

@@ -0,0 +1,95 @@
+#version 450
+
+struct SceneUBO {
+  vec3 camera_pos;
+  mat4 view;
+  mat4 projection;
+  vec3 ambient_light;
+  int point_light_count;
+};
+
+layout(binding = 0) uniform UBORenderable {
+  vec3 color;
+  int use_per_vertex_color;
+  int use_per_vertex_radius;
+  float radius;
+}
+ubo_renderable;
+
+layout(binding = 1) uniform UBOScene {
+  SceneUBO scene;
+  float window_width;
+  float window_height;
+  float tan_half_fov;
+  float aspect_ratio;
+}
+ubo_scene;
+
+struct PointLight {
+  vec3 pos;
+  vec3 color;
+};
+
+layout(binding = 2, std430) buffer SSBO {
+  PointLight point_lights[];
+}
+ssbo;
+
+layout(location = 0) out vec4 out_color;
+
+layout(location = 0) in vec4 pos_camera_space;
+layout(location = 1) in vec4 selected_color;
+layout(location = 2) in vec2 pos_2d;
+layout(location = 3) in float selected_radius;
+
+float project_z(float view_z) {
+  vec4 projected = ubo_scene.scene.projection * vec4(0, 0, view_z, 1);
+  return projected.z / projected.w;
+}
+
+vec3 to_camera_space(vec3 pos) {
+  vec4 temp = ubo_scene.scene.view * vec4(pos, 1.0);
+  return temp.xyz / temp.w;
+}
+
+// operates in camera space !!
+vec3 lambertian(vec3 frag_pos, vec3 frag_normal) {
+  vec3 ambient = ubo_scene.scene.ambient_light * selected_color.rgb;
+  vec3 result = ambient;
+
+  for (int i = 0; i < ubo_scene.scene.point_light_count; ++i) {
+    vec3 light_color = ssbo.point_lights[i].color;
+
+    vec3 light_dir =
+        normalize(to_camera_space(ssbo.point_lights[i].pos) - frag_pos);
+    vec3 normal = normalize(frag_normal);
+    vec3 diffuse =
+        max(dot(light_dir, normal), 0.0) * selected_color.rgb * light_color;
+
+    result += diffuse;
+  }
+
+  return result;
+}
+
+void main() {
+  vec2 coord2D = pos_2d;
+
+  if (length(coord2D) >= 1.0) {
+    discard;
+  }
+
+  float z_in_sphere = sqrt(1 - coord2D.x * coord2D.x - coord2D.y * coord2D.y);
+  vec3 coord_in_sphere = vec3(coord2D, z_in_sphere);
+
+  vec3 frag_pos =
+      pos_camera_space.xyz / pos_camera_space.w + coord_in_sphere * selected_radius;
+  vec3 frag_normal = coord_in_sphere;
+  vec3 color = lambertian(frag_pos, frag_normal);
+  out_color = vec4(color, selected_color.a);
+
+  float depth =
+      (pos_camera_space.z / pos_camera_space.w) + z_in_sphere * selected_radius;
+
+  gl_FragDepth = project_z(depth);
+}

taichi/shaders/Particles_vk.vert

@@ -0,0 +1,73 @@
+#version 450
+#extension GL_KHR_vulkan_glsl : enable
+
+layout(location = 0) in vec3 in_position;
+layout(location = 1) in vec3 in_normal;
+layout(location = 2) in vec2 in_texcoord;
+layout(location = 3) in vec4 in_color;
+
+struct SceneUBO {
+  vec3 camera_pos;
+  mat4 view;
+  mat4 projection;
+  vec3 ambient_light;
+  int point_light_count;
+};
+
+layout(binding = 0) uniform UBORenderable {
+  vec3 color;
+  int use_per_vertex_color;
+  int use_per_vertex_radius;
+  float radius;
+}
+ubo_renderable;
+
+layout(binding = 1) uniform UBOScene {
+  SceneUBO scene;
+  float window_width;
+  float window_height;
+  float tan_half_fov;
+  float aspect_ratio;
+}
+ubo_scene;
+layout(location = 0) out vec4 pos_camera_space;
+layout(location = 1) out vec4 selected_color;
+layout(location = 2) out vec2 pos_2d;
+layout(location = 3) out float selected_radius;
+
+const vec2 offsets[6] = {
+  vec2(-1.0f, 1.0f),
+  vec2(1.0f, -1.0f),
+  vec2(-1.0f, -1.0f),
+  vec2(-1.0f, 1.0f),
+  vec2(1.0f, 1.0f),
+  vec2(1.0f, -1.0f),
+};
+
+void main() {
+  float distance = length(in_position - ubo_scene.scene.camera_pos);
+
+  if (ubo_renderable.use_per_vertex_radius == 0) {
+    selected_radius = ubo_renderable.radius;
+  } else {
+    selected_radius = in_normal.x;
+  }
+
+  float hsize = selected_radius / (ubo_scene.tan_half_fov * distance);
+
+  pos_camera_space = ubo_scene.scene.view * vec4(in_position, 1.0);
+
+  pos_2d = offsets[gl_VertexIndex % 6];
+
+  vec4 pos_proj = ubo_scene.scene.projection * pos_camera_space;
+  pos_proj.xy += pos_2d * vec2(hsize, hsize * ubo_scene.window_width / ubo_scene.window_height) * pos_proj.w;
+
+  gl_Position = pos_proj;
+  gl_Position.y *= -1.0;
+
+  if (ubo_renderable.use_per_vertex_color == 0) {
+    selected_color = vec4(ubo_renderable.color, 1.0);
+  } else {
+    selected_color = in_color;
+  }
+}

taichi/shaders/SceneLines_vk.frag

@@ -0,0 +1,9 @@
+#version 450
+
+layout(location = 0) in vec3 color;
+
+layout(location = 0) out vec4 out_color;
+
+void main() {
+  out_color = vec4(color, 1);
+}

taichi/shaders/SceneLines_vk.vert

@@ -0,0 +1,12 @@
+#version 450
+
+layout(location = 0) in vec4 in_position;
+layout(location = 1) in vec4 in_color;
+
+layout(location = 0) out vec3 frag_color;
+
+void main() {
+  gl_Position = in_position;
+  gl_Position.y = -gl_Position.y;
+  frag_color = in_color.rgb;
+}

taichi/shaders/SetImage_vk.frag

@@ -0,0 +1,21 @@
+#version 450
+
+layout(binding = 0) uniform sampler2D texSampler;
+
+layout(location = 0) in vec2 frag_texcoord;
+
+layout(location = 0) out vec4 out_color;
+
+layout(binding = 1) uniform UBO {
+  vec2 lower_bound;
+  vec2 upper_bound;
+  float x_factor;
+  float y_factor;
+  int is_transposed;
+} ubo;
+
+void main() {
+  vec2 coord = frag_texcoord * vec2(ubo.x_factor,ubo.y_factor);
+  coord = clamp(coord, ubo.lower_bound, ubo.upper_bound);
+  out_color = textureLod(texSampler, ubo.is_transposed != 0 ? coord.yx : coord, 0);
+}

taichi/shaders/SetImage_vk.vert

@@ -0,0 +1,15 @@
+#version 450
+
+// layout(binding = 0) uniform UniformBufferObject {} ubo;
+
+layout(location = 0) in vec3 in_position;
+layout(location = 1) in vec3 in_normal;
+layout(location = 2) in vec2 in_texcoord;
+layout(location = 3) in vec4 in_color;
+
+layout(location = 0) out vec2 frag_texcoord;
+
+void main() {
+  gl_Position = vec4(in_position.xy, 0.0, 1.0);
+  frag_texcoord = in_texcoord;
+}

taichi/shaders/Triangles_vk.frag

@@ -0,0 +1,16 @@
+#version 450
+
+layout(location = 0) in vec2 frag_texcoord;
+layout(location = 1) in vec3 selected_color;
+
+layout(location = 0) out vec4 out_color;
+
+layout(binding = 0) uniform UniformBufferObject {
+  vec3 color;
+  int use_per_vertex_color;
+}
+ubo;
+
+void main() {
+  out_color = vec4(selected_color, 1);
+}

taichi/shaders/Triangles_vk.vert

@@ -0,0 +1,29 @@
+#version 450
+
+layout(location = 0) in vec3 in_position;
+layout(location = 1) in vec3 in_normal;
+layout(location = 2) in vec2 in_texcoord;
+layout(location = 3) in vec4 in_color;
+
+layout(location = 0) out vec2 frag_texcoord;
+layout(location = 1) out vec3 selected_color;
+
+layout(binding = 0) uniform UniformBufferObject {
+  vec3 color;
+  int use_per_vertex_color;
+}
+ubo;
+
+void main() {
+  float x = in_position.x * 2.0 - 1.0;
+  float y = -(in_position.y * 2.0 - 1.0);
+
+  gl_Position = vec4(x, y, 0.0, 1.0);
+  frag_texcoord = in_texcoord;
+
+  if (ubo.use_per_vertex_color == 0) {
+    selected_color = ubo.color;
+  } else {
+    selected_color = in_color.rgb;
+  }
+}

taichi/sparse/__init__.py

@@ -0,0 +1,3 @@
+# type: ignore
+
+from ._sparse_grid import *

taichi/sparse/_sparse_grid.py

@@ -0,0 +1,77 @@
+# type: ignore
+
+from taichi.lang.impl import grouped, root, static
+from taichi.lang.kernel_impl import kernel
+from taichi.lang.misc import ij, ijk
+from taichi.lang.snode import is_active
+from taichi.lang.struct import Struct
+from taichi.types.annotations import template
+from taichi.types.primitive_types import f32
+
+
+def grid(field_dict, shape):
+    """Creates a 2D/3D sparse grid with each element is a struct. The struct is placed on a bitmasked snode.
+
+    Args:
+        field_dict (dict): a dict, each item is like `name: type`.
+        shape (Tuple[int]): shape of the field.
+    Returns:
+        x: the created sparse grid, which is a bitmasked `ti.Struct.field`.
+
+    Examples::
+        # create a 2D sparse grid
+        >>> grid = ti.sparse.grid({'pos': ti.math.vec2, 'mass': ti.f32, 'grid2particles': ti.types.vector(20, ti.i32)}, shape=(10, 10))
+
+        # access
+        >>> grid[0, 0].pos = ti.math.vec2(1.0, 2.0)
+        >>> grid[0, 0].mass = 1.0
+        >>> grid[0, 0].grid2particles[2] = 123
+
+        # print the usage of the sparse grid, which is in [0,1]
+        >>> print(ti.sparse.usage(grid))
+        # 0.009999999776482582
+    """
+    x = Struct.field(field_dict)
+    if len(shape) == 2:
+        snode = root.bitmasked(ij, shape)
+        snode.place(x)
+    elif len(shape) == 3:
+        snode = root.bitmasked(ijk, shape)
+        snode.place(x)
+    else:
+        raise Exception("Only 2D and 3D sparse grids are supported")
+    return x
+
+
+@kernel
+def usage(x: template()) -> f32:
+    """
+    Get the usage of the sparse grid, which is in [0,1]
+
+    Args:
+        x(struct field): the sparse grid to be checked.
+    Returns:
+        usage(f32): the usage of the sparse grid, which is in [0,1]
+
+    Examples::
+        >>> grid = ti.sparse.grid({'pos': ti.math.vec2, 'mass': ti.f32, 'grid2particles': ti.types.vector(20, ti.i32)}, shape=(10, 10))
+        >>> grid[0, 0].mass = 1.0
+        >>> print(ti.sparse.usage(grid))
+        # 0.009999999776482582
+    """
+    cnt = 0
+    for I in grouped(x.parent()):
+        if is_active(x.parent(), I):
+            cnt += 1
+    total = 1.0
+    if static(len(x.shape) == 2):
+        total = x.shape[0] * x.shape[1]
+    elif static(len(x.shape) == 3):
+        total = x.shape[0] * x.shape[1] * x.shape[2]
+    else:
+        raise ValueError("The dimension of the sparse grid should be 2 or 3")
+    res = cnt / total
+    return res
+
+
+__all__ = ["grid", "usage"]

taichi/tools/__init__.py

@@ -0,0 +1,12 @@
+# type: ignore
+
+"""Taichi utility module.
+
+- `image` submodule for image io.
+- `video` submodule for exporting results to video files.
+- `diagnose` submodule for printing system environment information.
+"""
+
+from taichi.tools.diagnose import *
+from taichi.tools.np2ply import *
+from taichi.tools.vtk import *