taichi interpolation

interp.py

@ti.func
def _bilateral_interp_coeffs(self, 
    uv_grid:tm.vec2, 
    probe_elevation:float, 
    grid_height:int, 
    grid_width:int, 
    probe_height:int, 
    probe_width:int,
    cascade:int) -> (tm.ivec4, tm.vec4):
    """Returns pixel indices and weights for interpolation"""
    pos = tm.vec2(grid_width * uv_grid.x - 0.5, grid_height * uv_grid.y - 0.5)
    x0, y0, x1, y1 = int(tm.floor(pos.x)), int(tm.floor(pos.y)), 0, 0
    dx = (pos.x + 1.) - (float(x0) + 1.)
    dy = (pos.y + 1.) - (float(y0) + 1.)

    if self.wrap_mode == WrapMode.REPEAT:
        x0, y0 = x0%grid_width, y0%grid_height
        x1, y1 = (x0+1)%grid_width, (y0+1)%grid_height
    elif self.wrap_mode == WrapMode.CLAMP:
        x0, y0 = tm.clamp(x0, 0, grid_width-1), tm.clamp(y0, 0, grid_height-1)
        x1, y1 = tm.min(x0+1, grid_width-1), tm.min(y0+1, grid_height-1)

    q00 = int((y0 * grid_width) + x0)
    q01 = int((y1 * grid_width) + x0)
    q10 = int((y0 * grid_width) + x1)
    q11 = int((y1 * grid_width) + x1)
    indices = tm.ivec4(q00, q01, q10, q11)

    depth_bias = self.depth_bias
    canvas_shape = tm.vec2(float(self.shape[1]), float(self.shape[0]))
    uv00 = tm.vec2(x0 * probe_width + probe_width * 0.5, y0 * probe_height + probe_height * 0.5)/canvas_shape.xy + self.eps
    uv01 = tm.vec2(x0 * probe_width + probe_width * 0.5, y1 * probe_height + probe_height * 0.5)/canvas_shape.xy + self.eps
    uv10 = tm.vec2(x1 * probe_width + probe_width * 0.5, y0 * probe_height + probe_height * 0.5)/canvas_shape.xy + self.eps
    uv11 = tm.vec2(x1 * probe_width + probe_width * 0.5, y1 * probe_height + probe_height * 0.5)/canvas_shape.xy + self.eps

    d00 = self.positions.sample_bilinear(uv00).z
    d01 = self.positions.sample_bilinear(uv01).z
    d10 = self.positions.sample_bilinear(uv10).z
    d11 = self.positions.sample_bilinear(uv11).z

    weights = tm.vec4(0.);
    weights[0] = (1. - dx) * (1. - dy)
    weights[1] = (1. - dx) * dy
    weights[2] = dx * (1. - dy)
    weights[3] = dx * dy
    weights[0] *= tm.exp(-ti.abs(probe_elevation - d00) * depth_bias)
    weights[1] *= tm.exp(-ti.abs(probe_elevation - d01) * depth_bias)
    weights[2] *= tm.exp(-ti.abs(probe_elevation - d10) * depth_bias)
    weights[3] *= tm.exp(-ti.abs(probe_elevation - d11) * depth_bias)
    weights /= tm.dot(tm.vec4(1.), weights)

    return indices, weights