Untitled

import taichi
import taichi.math
from ray import Ray

taichi.init(arch=taichi.cuda)

# Image
aspectRatio = 16.0 / 9.0
imageWidth = 400
imageHeight = int(imageWidth / aspectRatio)
pixels = taichi.field(float, shape=(imageWidth, imageHeight))

# Camera
viewportHeight = 2.0
viewportWidth = aspectRatio * viewportHeight
focalLength = 1.0

origin = taichi.math.vec3(0, 0, 0)
horizontal = taichi.math.vec3(viewportWidth, 0, 0)
vertical = taichi.math.vec3(0, viewportHeight, 0)
lowerLeftCorner = origin - horizontal / 2 - vertical / 2 - taichi.math.vec3(0, 0, focalLength)


@taichi.func
def rayColor(r):
    unitDirection = taichi.math.vec3(r.direction)
    unitDirection.norm()
    t = 0.5 * (unitDirection.y + 1.0)
    return (1.0 - t) * taichi.math.vec3(1.0, 1.0, 1.0) + t * taichi.math.vec3(0.5, 0.7, 1.0)


@taichi.kernel
def main():
    # Render
    for i, j in pixels:
        u = i / (imageWidth - 1)
        v = j / (imageHeight - 1)
        r = Ray(origin, lowerLeftCorner + u * horizontal + v * vertical - origin)
        pixelColor = rayColor(r)

        ir = 255.999 * pixelColor.x
        ig = 255.999 * pixelColor.y
        ib = 255.999 * pixelColor.z

        pixels[i, j] = taichi.math.vec3(ir, ig, ib)

    taichi.tools.imwrite(pixels, "demo.png")

if __name__ == "__main__":
    main()