<commit_before>import taichi_lang as ti
import math
import numpy as np
import cv2
import os
import matplotlib.pyplot as plt

real = ti.f32
ti.set_default_fp(real)

max_steps = 4096
vis_interval = 2
output_vis_interval = 2
steps = 2048
assert steps * 2 <= max_steps

vis_resolution = 1024

scalar = lambda: ti.var(dt=real)
vec = lambda: ti.Vector(2, dt=real)

loss = scalar()

ti.cfg.arch = ti.cuda

x = vec()
v = vec()
rotation = scalar()
# angular velocity
omega = scalar()
halfsize = vec()

inverse_mass = scalar()
inverse_inertia = scalar()

v_inc = vec()
omega_inc = scalar()

n_objects = 1
# target_ball = 0
elasticity = 0.0
ground_height = 0.1
gravity = -9.8
friction = 0.5
# penalty = 1e4
penalty = 0.0
damping = 0


@ti.layout
def place():
  ti.root.dense(ti.l, max_steps).dense(ti.i, n_objects).place(x, v, rotation,
                                                              omega, v_inc,
                                                              omega_inc)
  ti.root.dense(ti.i, n_objects).place(halfsize, inverse_mass, inverse_inertia)
  ti.root.place(loss)
  ti.root.lazy_grad()


dt = 0.001
learning_rate = 0.02


@ti.func
def rotation_matrix(r):
  return ti.Matrix([[ti.cos(r), -ti.sin(r)], [ti.sin(r), ti.cos(r)]])


@ti.kernel
def initialize_properties():
  for i in range(n_objects):
    inverse_mass[i] = 1.0 / (4 * halfsize[i][0] * halfsize[i][1])
    inverse_inertia[i] = 1.0 / (4 / 3 * halfsize[i][0] * halfsize[i][1] * (
        halfsize[i][0] * halfsize[i][0] + halfsize[i][1] * halfsize[i][1]))
    ti.print(inverse_mass[i])
    ti.print(inverse_inertia[i])


@ti.func
def cross(a, b):
  return a[0] * b[1] - a[1] * b[0]


@ti.kernel
def apply_gravity_and_collide(t: ti.i32):
  for i in range(n_objects):
    for k in ti.static(range(4)):
      # the corner for collision detection
      hs = halfsize[k]
      offset_scale = ti.Vector([k % 2 * 2 - 1, k // 2 % 2 * 2 - 1])
      rot = rotation[t, i]
      rot_matrix = rotation_matrix(rot)
      
      rela_pos = rot_matrix @ (offset_scale * hs)
      rela_v = omega[t, i] * ti.Vector([-rela_pos[1], rela_pos[0]])
      
      corner_x = x[t, i] + rela_pos
      corner_v = v[t, i] + rela_v + dt * gravity * ti.Vector([0.0, 1.0])
      
      # Apply impulse so that there's no sinking
      normal = ti.Vector([0.0, 1.0])
      tao = ti.Vector([1.0, 0.0])
      
      rn = cross(rela_pos, normal)
      rt = cross(rela_pos, tao)
      impulse_contribution = inverse_mass[i] + ti.sqr(rn) * \
                             inverse_inertia[i]
      timpulse_contribution = inverse_mass[i] + ti.sqr(rt) * \
                              inverse_inertia[i]
      
      # ti.print(impulse_contribution)
      
      rela_v_ground = normal.dot(corner_v)
      
      impulse = 0.0
      timpulse = 0.0
      if rela_v_ground < 0 and corner_x[1] < ground_height:
        impulse = -(1 + elasticity) * rela_v_ground / impulse_contribution
        if impulse > 0:
          # friction
          timpulse = -rela_v.dot(tao) / timpulse_contribution
          timpulse = ti.min(friction * impulse,
                            ti.max(-friction * impulse, timpulse))

      ti.print(impulse)
      ti.print(timpulse)
      
      if corner_x[1] < ground_height:
        # apply penalty
        impulse = impulse - dt * penalty * (
              corner_x[1] - ground_height) / impulse_contribution
      
      ti.atomic_add(v_inc[t + 1, i],
                    (impulse * normal + timpulse * tao) * inverse_mass[i])
      ti.atomic_add(omega_inc[t + 1, i],
                    (impulse * rn + timpulse * rt) * inverse_inertia[i])


@ti.kernel
def advance(t: ti.i32):
  for i in range(n_objects):
    s = math.exp(-dt * damping)
    v[t, i] = s * v[t - 1, i] + v_inc[t, i] + dt * gravity * ti.Vector([0.0, 1.0])
    x[t, i] = x[t - 1, i] + dt * v[t, i]
    omega[t, i] = s * omega[t - 1, i] + omega_inc[t, i]
    rotation[t, i] = rotation[t - 1, i] + dt * omega[t, i]


@ti.kernel
def compute_loss(t: ti.i32):
  pass


def forward(output=None):
  # initialize()
  
  for i in range(n_objects):
    x[0, i] = [0.5, 0.5]
    halfsize[i] = [0.1, 0.05]
    rotation[0, i] = 0.0
    omega[0, i] = 1
  
  initialize_properties()
  
  interval = vis_interval
  if output:
    interval = output_vis_interval
    os.makedirs('rigid_body/{}/'.format(output), exist_ok=True)
  
  for t in range(1, steps):
    apply_gravity_and_collide(t - 1)
    advance(t)
    
    if (t + 1) % interval == 0:
      img = np.ones(shape=(vis_resolution, vis_resolution, 3),
                    dtype=np.float32) * 0.8
      
      color = (0.3, 0.5, 0.8)
      for i in range(n_objects):
        points = []
        for k in range(4):
          offset_scale = [[-1, -1], [1, -1], [1, 1], [-1, 1]][k]
          rot = rotation[t, i]
          rot_matrix = np.array(
            [[math.cos(rot), -math.sin(rot)], [math.sin(rot), math.cos(rot)]])
          <commit_msg>reproduce wierd sliding<commit_after>import taichi_lang as ti
import math
import numpy as np
import cv2
import os
import matplotlib.pyplot as plt

real = ti.f32
ti.set_default_fp(real)

max_steps = 4096
vis_interval = 2
output_vis_interval = 2
steps = 2048
assert steps * 2 <= max_steps

vis_resolution = 1024

scalar = lambda: ti.var(dt=real)
vec = lambda: ti.Vector(2, dt=real)

loss = scalar()

ti.cfg.arch = ti.cuda

x = vec()
v = vec()
rotation = scalar()
# angular velocity
omega = scalar()
halfsize = vec()

inverse_mass = scalar()
inverse_inertia = scalar()

v_inc = vec()
omega_inc = scalar()

n_objects = 1
# target_ball = 0
elasticity = 0.5
ground_height = 0.1
gravity = -9.8
friction = 0.5
# penalty = 1e4
penalty = 0.0
damping = 0


@ti.layout
def place():
  ti.root.dense(ti.l, max_steps).dense(ti.i, n_objects).place(x, v, rotation,
                                                              omega, v_inc,
                                                              omega_inc)
  ti.root.dense(ti.i, n_objects).place(halfsize, inverse_mass, inverse_inertia)
  ti.root.place(loss)
  ti.root.lazy_grad()


dt = 0.001
learning_rate = 0.02


@ti.func
def rotation_matrix(r):
  return ti.Matrix([[ti.cos(r), -ti.sin(r)], [ti.sin(r), ti.cos(r)]])


@ti.kernel
def initialize_properties():
  for i in range(n_objects):
    inverse_mass[i] = 1.0 / (4 * halfsize[i][0] * halfsize[i][1])
    inverse_inertia[i] = 1.0 / (4 / 3 * halfsize[i][0] * halfsize[i][1] * (
        halfsize[i][0] * halfsize[i][0] + halfsize[i][1] * halfsize[i][1]))
    ti.print(inverse_mass[i])
    ti.print(inverse_inertia[i])


@ti.func
def cross(a, b):
  return a[0] * b[1] - a[1] * b[0]


@ti.kernel
def apply_gravity_and_collide(t: ti.i32):
  for i in range(n_objects):
    for k in ti.static(range(4)):
      # the corner for collision detection
      hs = halfsize[k]
      offset_scale = ti.Vector([k % 2 * 2 - 1, k // 2 % 2 * 2 - 1])
      rot = rotation[t, i]
      rot_matrix = rotation_matrix(rot)
      
      rela_pos = rot_matrix @ (offset_scale * hs)
      rela_v = omega[t, i] * ti.Vector([-rela_pos[1], rela_pos[0]])
      
      corner_x = x[t, i] + rela_pos
      corner_v = v[t, i] + rela_v + dt * gravity * ti.Vector([0.0, 1.0])
      
      # Apply impulse so that there's no sinking
      normal = ti.Vector([0.0, 1.0])
      tao = ti.Vector([1.0, 0.0])
      
      rn = cross(rela_pos, normal)
      rt = cross(rela_pos, tao)
      impulse_contribution = inverse_mass[i] + ti.sqr(rn) * \
                             inverse_inertia[i]
      timpulse_contribution = inverse_mass[i] + ti.sqr(rt) * \
                              inverse_inertia[i]
      
      # ti.print(impulse_contribution)
      
      rela_v_ground = normal.dot(corner_v)
      
      impulse = 0.0
      timpulse = 0.0
      if rela_v_ground < 0 and corner_x[1] < ground_height:
        impulse = -(1 + elasticity) * rela_v_ground / impulse_contribution
        if impulse > 0:
          # friction
          timpulse = -rela_v.dot(tao) / timpulse_contribution
          timpulse = ti.min(friction * impulse,
                            ti.max(-friction * impulse, timpulse))

      ti.print(impulse)
      ti.print(timpulse)
      
      if corner_x[1] < ground_height:
        # apply penalty
        impulse = impulse - dt * penalty * (
              corner_x[1] - ground_height) / impulse_contribution
      
      ti.atomic_add(v_inc[t + 1, i],
                    (impulse * normal + timpulse * tao) * inverse_mass[i])
      ti.atomic_add(omega_inc[t + 1, i],
                    (impulse * rn + timpulse * rt) * inverse_inertia[i])


@ti.kernel
def advance(t: ti.i32):
  for i in range(n_objects):
    s = math.exp(-dt * damping)
    v[t, i] = s * v[t - 1, i] + v_inc[t, i] + dt * gravity * ti.Vector([0.0, 1.0])
    x[t, i] = x[t - 1, i] + dt * v[t, i]
    omega[t, i] = s * omega[t - 1, i] + omega_inc[t, i]
    rotation[t, i] = rotation[t - 1, i] + dt * omega[t, i]


@ti.kernel
def compute_loss(t: ti.i32):
  pass


def forward(output=None):
  # initialize()
  
  for i in range(n_objects):
    x[0, i] = [0.5, 0.5]
    halfsize[i] = [0.1, 0.05]
    rotation[0, i] = 0.4
    omega[0, i] = 1
  
  initialize_properties()
  
  interval = vis_interval
  if output:
    interval = output_vis_interval
    os.makedirs('rigid_body/{}/'.format(output), exist_ok=True)
  
  for t in range(1, steps):
    apply_gravity_and_collide(t - 1)
    advance(t)
    
    if (t + 1) % interval == 0:
      img = np.ones(shape=(vis_resolution, vis_resolution, 3),
                    dtype=np.float32) * 0.8
      
      color = (0.3, 0.5, 0.8)
      for i in range(n_objects):
        points = []
        for k in range(4):
          offset_scale = [[-1, -1], [1, -1], [1, 1], [-1, 1]][k]
          rot = rotation[t, i]
          rot_matrix = np.array(
            [[math.cos(rot), -math.sin(rot)], [math.sin(rot), math.cos(rot)]])
          
{'fraction_non_alphanumeric': 0.14759771986970685, 'fraction_numerical': 0.02992671009771987, 'mean_word_length': 2.1972014318255777, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '199371', 'n_tokens_mistral': 3979, 'n_tokens_neox': 3714, 'n_words': 842}