def matmul_(a0, a1): if not (a0.Size == a1.Size and (a0.IsMatrix or a0.IsVector) \ and (a1.IsMatrix or a1.IsVector)): raise Exception("matmul(): unsupported operand shape!") if a0.IsVector and a1.IsVector: return _ek.dot(a0, a1) elif a0.IsMatrix and a1.IsVector: ar = a0[0] * a1[0] for i in range(1, a1.Size): ar = _ek.fmadd(a0[i], a1[i], ar) return ar elif a0.IsVector and a1.IsMatrix: ar = a1.Value() for i in range(a1.Size): ar[i] = _ek.dot(a0, a1[i]) return ar else: # matrix @ matrix ar, sr = _check2(a0, a1) for j in range(a0.Size): accum = a0[0] * _ek.full(a0.Value, a1[0, j]) for i in range(1, a0.Size): accum = _ek.fmadd(a0[i], _ek.full(a0.Value, a1[i, j]), accum) ar[j] = accum return ar
def quat_to_euler(q): name = _ek.detail.array_name('Array', q.Type, [3], q.IsScalar) module = _modules.get(q.__module__) Array3f = getattr(module, name) sinp = 2 * _ek.fmsub(q.w, q.y, q.z * q.x) gimbal_lock = _ek.abs(sinp) > (1.0 - 5e-8) # roll (x-axis rotation) q_y_2 = _ek.sqr(q.y) sinr_cosp = 2 * _ek.fmadd(q.w, q.x, q.y * q.z) cosr_cosp = _ek.fnmadd(2, _ek.fmadd(q.x, q.x, q_y_2), 1) roll = _ek.select(gimbal_lock, 2 * _ek.atan2(q.x, q.w), _ek.atan2(sinr_cosp, cosr_cosp)) # pitch (y-axis rotation) pitch = _ek.select(gimbal_lock, _ek.copysign(0.5 * _ek.Pi, sinp), _ek.asin(sinp)) # yaw (z-axis rotation) siny_cosp = 2 * _ek.fmadd(q.w, q.z, q.x * q.y) cosy_cosp = _ek.fnmadd(2, _ek.fmadd(q.z, q.z, q_y_2), 1) yaw = _ek.select(gimbal_lock, 0, _ek.atan2(siny_cosp, cosy_cosp)) return Array3f(roll, pitch, yaw)
def test05_scalar(t): if not ek.is_array_v(t) or ek.array_size_v(t) == 0: return get_class(t.__module__) if ek.is_mask_v(t): assert ek.all_nested(t(True)) assert ek.any_nested(t(True)) assert ek.none_nested(t(False)) assert ek.all_nested(t(False) ^ t(True)) assert ek.all_nested(ek.eq(t(False), t(False))) assert ek.none_nested(ek.eq(t(True), t(False))) if ek.is_arithmetic_v(t): assert t(1) + t(1) == t(2) assert t(3) - t(1) == t(2) assert t(2) * t(2) == t(4) assert ek.min(t(2), t(3)) == t(2) assert ek.max(t(2), t(3)) == t(3) if ek.is_signed_v(t): assert t(2) * t(-2) == t(-4) assert ek.abs(t(-2)) == t(2) if ek.is_integral_v(t): assert t(6) // t(2) == t(3) assert t(7) % t(2) == t(1) assert t(7) >> 1 == t(3) assert t(7) << 1 == t(14) assert t(1) | t(2) == t(3) assert t(1) ^ t(3) == t(2) assert t(1) & t(3) == t(1) else: assert t(6) / t(2) == t(3) assert ek.sqrt(t(4)) == t(2) assert ek.fmadd(t(1), t(2), t(3)) == t(5) assert ek.fmsub(t(1), t(2), t(3)) == t(-1) assert ek.fnmadd(t(1), t(2), t(3)) == t(1) assert ek.fnmsub(t(1), t(2), t(3)) == t(-5) assert (t(1) & True) == t(1) assert (t(1) & False) == t(0) assert (t(1) | False) == t(1) assert ek.all_nested(t(3) > t(2)) assert ek.all_nested(ek.eq(t(2), t(2))) assert ek.all_nested(ek.neq(t(3), t(2))) assert ek.all_nested(t(1) >= t(1)) assert ek.all_nested(t(2) < t(3)) assert ek.all_nested(t(1) <= t(1)) assert ek.select(ek.eq(t(2), t(2)), t(4), t(5)) == t(4) assert ek.select(ek.eq(t(3), t(2)), t(4), t(5)) == t(5) t2 = t(2) assert ek.hsum(t2) == t.Value(2 * len(t2)) assert ek.dot(t2, t2) == t.Value(4 * len(t2)) assert ek.dot_async(t2, t2) == t(4 * len(t2)) value = t(1) value[ek.eq(value, t(1))] = t(2) value[ek.eq(value, t(3))] = t(5) assert value == t(2)
def polar_decomp(a, it=10): q = type(a)(a) for i in range(it): qi = _ek.inverse_transpose(q) gamma = _ek.sqrt(_ek.frob(qi) / _ek.frob(q)) s1, s2 = gamma * .5, (_ek.rcp(gamma) * .5) for i in range(a.Size): q[i] = _ek.fmadd(q[i], s1, qi[i] * s2) return q, transpose(q) @ a
def frob(a): if not _ek.is_matrix_v(a): raise Exception('Unsupported type!') result = _ek.sqr(a[0]) for i in range(1, a.Size): value = a[i] result = _ek.fmadd(value, value, result) return _ek.hsum(result)
def fmadd_(a0, a1, a2): if not a0.IsFloat: raise Exception("fmadd(): requires floating point operands!") if not a0.IsSpecial: ar, sr = _check3(a0, a1, a2) for i in range(sr): ar[i] = _ek.fmadd(a0[i], a1[i], a2[i]) return ar else: return a0 * a1 + a2
def linspace_(cls, min, max, size=1): result = cls.empty_(size) step = (max - min) / (len(result) - 1) if cls.IsFloat: for i in range(len(result)): result[i] = min + step * i else: for i in range(len(result)): result[i] = _ek.fmadd(step, i, min) return result
def sincos(x): Float = type(x) Int = ek.int_array_t(Float) xa = ek.abs(x) j = Int(xa * 1.2732395447351626862) j = (j + Int(1)) & ~Int(1) y = Float(j) Shift = Float.Type.Size * 8 - 3 sign_sin = ek.reinterpret_array(Float, j << Shift) ^ x sign_cos = ek.reinterpret_array(Float, (~(j - Int(2)) << Shift)) y = xa - y * 0.78515625 \ - y * 2.4187564849853515625e-4 \ - y * 3.77489497744594108e-8 z = y * y z |= ek.eq(xa, ek.Infinity) s = poly2(z, -1.6666654611e-1, 8.3321608736e-3, -1.9515295891e-4) * z c = poly2(z, 4.166664568298827e-2, -1.388731625493765e-3, 2.443315711809948e-5) * z s = ek.fmadd(s, y, y) c = ek.fmadd(c, z, ek.fmadd(z, -0.5, 1)) polymask = ek.eq(j & Int(2), ek.zero(Int)) return ( ek.mulsign(ek.select(polymask, s, c), sign_sin), ek.mulsign(ek.select(polymask, c, s), sign_cos) )
def tabulate_pdf(self): """ Numerically integrate the provided probability density function over each cell to generate an array resembling the histogram computed by ``tabulate_histogram()``. The function uses the trapezoid rule over intervals discretized into ``self.ires`` separate function evaluations. """ from mitsuba.core import Float, Vector2f, ScalarVector2f extents = self.bounds.extents() endpoint = self.bounds.max - extents / ScalarVector2f(self.res) # Compute a set of nodes where the PDF should be evaluated x, y = ek.meshgrid( ek.linspace(Float, self.bounds.min.x, endpoint.x, self.res.x), ek.linspace(Float, self.bounds.min.y, endpoint.y, self.res.y)) endpoint = extents / ScalarVector2f(self.res) eps = 1e-4 nx = ek.linspace(Float, eps, endpoint.x * (1 - eps), self.ires) ny = ek.linspace(Float, eps, endpoint.y * (1 - eps), self.ires) wx = [1 / (self.ires - 1)] * self.ires wy = [1 / (self.ires - 1)] * self.ires wx[0] = wx[-1] = wx[0] * .5 wy[0] = wy[-1] = wy[0] * .5 integral = 0 self.histogram_start = time.time() for yi, dy in enumerate(ny): for xi, dx in enumerate(nx): xy = self.domain.map_forward(Vector2f(x + dx, y + dy)) pdf = self.pdf_func(xy) integral = ek.fmadd(pdf, wx[xi] * wy[yi], integral) self.histogram_end = time.time() self.pdf = integral * (ek.hprod(extents / ScalarVector2f(self.res)) * self.sample_count) # A few sanity checks pdf_min = ek.hmin(self.pdf) / self.sample_count if not pdf_min >= 0: self._log('Failure: Encountered a cell with a ' 'negative PDF value: %f' % pdf_min) self.fail = True self.pdf_sum = ek.hsum(self.pdf) / self.sample_count if self.pdf_sum > 1.1: self._log('Failure: PDF integrates to a value greater ' 'than 1.0: %f' % self.pdf_sum) self.fail = True
def det(m): if not _ek.is_matrix_v(m): raise Exception("Unsupported target type!") if m.Size == 1: return m[0, 0] elif m.Size == 2: return _ek.fmsub(m[0, 0], m[1, 1], m[0, 1] * m[1, 0]) elif m.Size == 3: return _ek.dot(m[0], _ek.cross(m[1], m[2])) elif m.Size == 4: col0, col1, col2, col3 = m col1 = _ek.shuffle((2, 3, 0, 1), col1) col3 = _ek.shuffle((2, 3, 0, 1), col3) temp = _ek.shuffle((1, 0, 3, 2), col2 * col3) row0 = col1 * temp temp = _ek.shuffle((2, 3, 0, 1), temp) row0 = _ek.fmsub(col1, temp, row0) temp = _ek.shuffle((1, 0, 3, 2), col1 * col2) row0 = _ek.fmadd(col3, temp, row0) temp = _ek.shuffle((2, 3, 0, 1), temp) row0 = _ek.fnmadd(col3, temp, row0) col1 = _ek.shuffle((2, 3, 0, 1), col1) col2 = _ek.shuffle((2, 3, 0, 1), col2) temp = _ek.shuffle((1, 0, 3, 2), col1 * col3) row0 = _ek.fmadd(col2, temp, row0) temp = _ek.shuffle((2, 3, 0, 1), temp) row0 = _ek.fnmadd(col2, temp, row0) return _ek.dot(col0, row0) else: raise Exception('Unsupported array size!')
def mul_(a0, a1): if not a0.IsArithmetic: raise Exception("mul(): requires arithmetic operands!") if _ek.array_depth_v(a1) < _ek.array_depth_v(a0): # Avoid type promotion in scalars multiplication, which would # be costly for special types (matrices, quaternions, etc.) sr = len(a0) ar = a0.empty_(sr if a0.Size == Dynamic else 0) for i in range(len(a0)): ar[i] = a0[i] * a1 return ar ar, sr = _check2(a0, a1) if not a0.IsSpecial: for i in range(sr): ar[i] = a0[i] * a1[i] elif a0.IsComplex: ar.real = _ek.fmsub(a0.real, a1.real, a0.imag * a1.imag) ar.imag = _ek.fmadd(a0.real, a1.imag, a0.imag * a1.real) elif a0.IsQuaternion: tbl = (4, 3, -2, 1, -3, 4, 1, 2, 2, -1, 4, 3, -1, -2, -3, 4) for i in range(4): accum = 0 for j in range(4): idx = tbl[i * 4 + j] value = a1[abs(idx) - 1] accum = _ek.fmadd(a0[j], value if idx > 0 else -value, accum) ar[i] = accum elif a0.IsMatrix: raise Exception( "mul(): please use the matrix multiplication operator '@' instead." ) else: raise Exception("mul(): unsupported array type!") return ar
def quat_to_euler(q): q_y_2 = _ek.sqr(q.y) sinr_cosp = 2 * _ek.fmadd(q.w, q.x, q.y * q.z) cosr_cosp = _ek.fnmadd(2, _ek.fmadd(q.x, q.x, q_y_2), 1) roll = _ek.atan2(sinr_cosp, cosr_cosp) # pitch (y-axis rotation) sinp = 2 * _ek.fmsub(q.w, q.y, q.z * q.x) if (_ek.abs(sinp) >= 1.0): pitch = _ek.copysign(0.5 * _ek.Pi, sinp) else: pitch = _ek.asin(sinp) # yaw (z-axis rotation) siny_cosp = 2 * _ek.fmadd(q.w, q.z, q.x * q.y) cosy_cosp = _ek.fnmadd(2, _ek.fmadd(q.z, q.z, q_y_2), 1) yaw = _ek.atan2(siny_cosp, cosy_cosp) name = _ek.detail.array_name('Array', q.Type, [3], q.IsScalar) module = _modules.get(q.__module__) Array3f = getattr(module, name) return Array3f(roll, pitch, yaw)
def dot_(a0, a1): size = len(a0) if size == 0: raise Exception("dot(): zero-sized array!") if size != len(a1): raise Exception("dot(): incompatible array sizes!") if type(a0) is not type(a1): raise Exception("Type mismatch!") value = a0[0] * a1[0] if a0.IsFloat: for i in range(1, size): value = _ek.fmadd(a0[i], a1[i], value) else: for i in range(1, size): value += a0[i] * a1[i] return value
def test04_operators(): I3 = ek.scalar.Array3i F3 = ek.scalar.Array3f assert (I3(1, 2, 3) + I3(0, 1, 0) == I3(1, 3, 3)) assert (I3(1, 2, 3) - I3(0, 1, 0) == I3(1, 1, 3)) assert (I3(1, 2, 3) * I3(0, 1, 0) == I3(0, 2, 0)) assert (I3(1, 2, 3) // I3(2, 2, 2) == I3(0, 1, 1)) assert (I3(1, 2, 3) % I3(2, 2, 2) == I3(1, 0, 1)) assert (I3(1, 2, 3) << 1 == I3(2, 4, 6)) assert (I3(1, 2, 3) >> 1 == I3(0, 1, 1)) assert (I3(1, 2, 3) & 1 == I3(1, 0, 1)) assert (I3(1, 2, 3) | 1 == I3(1, 3, 3)) assert (I3(1, 2, 3) ^ 1 == I3(0, 3, 2)) assert (-I3(1, 2, 3) == I3(-1, -2, -3)) assert (~I3(1, 2, 3) == I3(-2, -3, -4)) assert (ek.abs(I3(1, -2, 3)) == I3(1, 2, 3)) assert (abs(I3(1, -2, 3)) == I3(1, 2, 3)) assert (ek.abs(I3(1, -2, 3)) == I3(1, 2, 3)) assert (ek.fmadd(F3(1, 2, 3), F3(2, 3, 1), F3(1, 1, 1)) == F3(3, 7, 4)) assert (ek.fmsub(F3(1, 2, 3), F3(2, 3, 1), F3(1, 1, 1)) == F3(1, 5, 2)) assert (ek.fnmadd(F3(1, 2, 3), F3(2, 3, 1), F3(1, 1, 1)) == F3(-1, -5, -2)) assert (ek.fnmsub(F3(1, 2, 3), F3(2, 3, 1), F3(1, 1, 1)) == F3(-3, -7, -4))
def poly2(x, c0, c1, c2): x2 = ek.sqr(x) return ek.fmadd(x2, c2, ek.fmadd(x, c1, c0))
def inverse_transpose(m): if not _ek.is_matrix_v(m): raise Exception("Unsupported target type!") t = type(m) if m.Size == 1: return t(_ek.rcp(m[0, 0])) elif m.Size == 2: inv_det = _ek.rcp(_ek.fmsub(m[0, 0], m[1, 1], m[0, 1] * m[1, 0])) return t(m[1, 1] * inv_det, -m[1, 0] * inv_det, -m[0, 1] * inv_det, m[0, 0] * inv_det) elif m.Size == 3: col0, col1, col2 = m row0 = _ek.cross(col1, col2) row1 = _ek.cross(col2, col0) row2 = _ek.cross(col0, col1) inv_det = _ek.rcp(_ek.dot(col0, row0)) return t(row0 * inv_det, row1 * inv_det, row2 * inv_det) elif m.Size == 4: col0, col1, col2, col3 = m col1 = _ek.shuffle((2, 3, 0, 1), col1) col3 = _ek.shuffle((2, 3, 0, 1), col3) temp = _ek.shuffle((1, 0, 3, 2), col2 * col3) row0 = col1 * temp row1 = col0 * temp temp = _ek.shuffle((2, 3, 0, 1), temp) row0 = _ek.fmsub(col1, temp, row0) row1 = _ek.shuffle((2, 3, 0, 1), _ek.fmsub(col0, temp, row1)) temp = _ek.shuffle((1, 0, 3, 2), col1 * col2) row0 = _ek.fmadd(col3, temp, row0) row3 = col0 * temp temp = _ek.shuffle((2, 3, 0, 1), temp) row0 = _ek.fnmadd(col3, temp, row0) row3 = _ek.shuffle((2, 3, 0, 1), _ek.fmsub(col0, temp, row3)) temp = _ek.shuffle((1, 0, 3, 2), _ek.shuffle((2, 3, 0, 1), col1) * col3) col2 = _ek.shuffle((2, 3, 0, 1), col2) row0 = _ek.fmadd(col2, temp, row0) row2 = col0 * temp temp = _ek.shuffle((2, 3, 0, 1), temp) row0 = _ek.fnmadd(col2, temp, row0) row2 = _ek.shuffle((2, 3, 0, 1), _ek.fmsub(col0, temp, row2)) temp = _ek.shuffle((1, 0, 3, 2), col0 * col1) row2 = _ek.fmadd(col3, temp, row2) row3 = _ek.fmsub(col2, temp, row3) temp = _ek.shuffle((2, 3, 0, 1), temp) row2 = _ek.fmsub(col3, temp, row2) row3 = _ek.fnmadd(col2, temp, row3) temp = _ek.shuffle((1, 0, 3, 2), col0 * col3) row1 = _ek.fnmadd(col2, temp, row1) row2 = _ek.fmadd(col1, temp, row2) temp = _ek.shuffle((2, 3, 0, 1), temp) row1 = _ek.fmadd(col2, temp, row1) row2 = _ek.fnmadd(col1, temp, row2) temp = _ek.shuffle((1, 0, 3, 2), col0 * col2) row1 = _ek.fmadd(col3, temp, row1) row3 = _ek.fnmadd(col1, temp, row3) temp = _ek.shuffle((2, 3, 0, 1), temp) row1 = _ek.fnmadd(col3, temp, row1) row3 = _ek.fmadd(col1, temp, row3) inv_det = _ek.rcp(_ek.dot(col0, row0)) return t(row0 * inv_det, row1 * inv_det, row2 * inv_det, row3 * inv_det) else: raise Exception('Unsupported array size!')
def test07_from_builtin(): # Fmadd should fallback to regular multiply (complex)-add assert ek.fmadd(C(2, 2), C(5, 5), C(5, 6)) == C(5, 26)