<jupyter_start><jupyter_text><jupyter_code>import jax.numpy as jnp
from jax import jit, vmap, grad
import jax
import matplotlib.pyplot as plt
def render_pixel(uvt):
  uv = uvt[0:2]
  time = uvt[2]

  def vec3(x, y, z):
    return jnp.array([x, y, z])

  def mag(v):
    return jnp.sqrt(v.dot(v))

  def normalize(v):
    return v / mag(v)

  def sdf_sphere(v, r):
    return mag(v) - r

  def scene(p):
    return sdf_sphere(p, 0.6)

  sdf_normal = grad(scene)
  ray_dir = normalize(vec3(*uv, 1.0))
  ray_pos = vec3(0.0, 0.0, -2.0)
  eps = 0.001
  max_dist = 100
  bg_color = vec3(0.0, 0.0, 0.0)

  def raymarch_step(d):
    d_s = scene(ray_pos + ray_dir * d[0] * 0.95)
    return d[0] + d_s, d_s

  run_condition = lambda d : jnp.where((d[1] > eps) & (d[1] < max_dist), True, False)

  dist_i = scene(ray_pos)
  dist_o, dist_s = jax.lax.while_loop(run_condition, raymarch_step, (dist_i, dist_i))

  intersect_mask = (dist_s < eps).astype(jnp.float32)

  def color(p):
    light_dir = normalize(vec3(0.2,-0.4,-1.0))
    normal = sdf_normal(p)
    light = jnp.maximum(normal.dot(light_dir), 0.0)
    return jnp.array([0.8, 0.1, 0.15]) * light

  intersect_pos = ray_pos + ray_dir * dist_o
  obj_color = color(intersect_pos)

  return intersect_mask * obj_color + (1.0 - intersect_mask) * bg_color

img_res = 512
ax_coords = jnp.arange(img_res)/img_res-0.5
img_coords = jnp.meshgrid(ax_coords, ax_coords)
img_coords = jnp.stack([*img_coords, jnp.full((img_res, img_res), 1.0)], axis=2)
# x = jnp.array([-0.4, -0.4, -2.0])
# x.dot(x)
render_pixel(jnp.array([0.0, 0.0, 1.0]))
render_pixels = jit(vmap(vmap(render_pixel)))
test_img = render_pixels(img_coords)
plt.imshow(test_img)
img_coords[0][0]
test_img.min()
img_coords.shape<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1630558722919042, 'fraction_numerical': 0.04503990877993159, 'mean_word_length': 3.875, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '243638', 'license': 'MIT', 'n_tokens_mistral': 852, 'n_tokens_neox': 803, 'n_words': 153, 'path': 'differentiable_rendering_sdf.ipynb', 'repo_name': 'PWhiddy/jax-experiments', 'size': '1754'}