def step(self):
""" Take a gradient step """
self.t += 1
from mitsuba.core import Float
lr_t = ek.detach(Float(self.lr * ek.sqrt(1 - self.beta_2**self.t) /
(1 - self.beta_1**self.t), literal=False))
for k, p in self.params.items():
g_p = ek.gradient(p)
size = ek.slices(g_p)
if size == 0:
continue
elif size != ek.slices(self.state[k][0]):
# Reset state if data size has changed
self._reset(k)
m_tp, v_tp = self.state[k]
m_t = self.beta_1 * m_tp + (1 - self.beta_1) * g_p
v_t = self.beta_2 * v_tp + (1 - self.beta_2) * ek.sqr(g_p)
self.state[k] = (m_t, v_t)
u = ek.detach(p) - lr_t * m_t / (ek.sqrt(v_t) + self.epsilon)
u = type(p)(u)
ek.set_requires_gradient(u)
self.params[k] = u
def test01_coordinate_system(variant_scalar_rgb):
from mitsuba.core import coordinate_system, warp
def branchless_onb(n):
"""
Building an Orthonormal Basis, Revisited
Tom Duff, James Burgess, Per Christensen, Christophe Hery,
Andrew Kensler, Max Liani, and Ryusuke Villemin
"""
sign = ek.copysign(1.0, n[2])
a = -1.0 / (sign + n[2])
b = n[0] * n[1] * a
return (
[1.0 + sign * n[0] * n[0] * a, sign * b, -sign * n[0]],
[b, sign + n[1] * n[1] * a, -n[1]]
)
a = [0.70710678, -0. , -0.70710678]
b = [-0., 1., 0.]
assert ek.allclose(
branchless_onb([ek.sqrt(0.5), 0, ek.sqrt(0.5)]), (a, b), atol=1e-6)
assert ek.allclose(
coordinate_system([ek.sqrt(0.5), 0, ek.sqrt(0.5)]), (a, b), atol=1e-6)
for u in ek.linspace(Float, 0, 1, 10):
for v in ek.linspace(Float, 0, 1, 10):
n = warp.square_to_uniform_sphere([u, v])
s1, t1 = branchless_onb(n)
s2, t2 = coordinate_system(n)
assert ek.allclose(s1, s2, atol=1e-6)
assert ek.allclose(t1, t2, atol=1e-6)
def test01_fresnel(variant_scalar_rgb):
from mitsuba.render import fresnel
ct_crit = -ek.sqrt(1 - 1 / 1.5**2)
assert ek.allclose(fresnel(1, 1.5), (0.04, -1, 1.5, 1 / 1.5))
assert ek.allclose(fresnel(-1, 1.5), (0.04, 1, 1 / 1.5, 1.5))
assert ek.allclose(fresnel(1, 1 / 1.5), (0.04, -1, 1 / 1.5, 1.5))
assert ek.allclose(fresnel(-1, 1 / 1.5), (0.04, 1, 1.5, 1 / 1.5))
assert ek.allclose(fresnel(0, 1.5), (1, ct_crit, 1.5, 1 / 1.5))
assert ek.allclose(fresnel(0, 1 / 1.5), (1, 0, 1 / 1.5, 1.5))
# Spot check at 45 deg (1 -> 1.5)
# Source: http://hyperphysics.phy-astr.gsu.edu/hbase/phyopt/freseq.html
F, cos_theta_t, _, scale = fresnel(ek.cos(45 * ek.pi / 180), 1.5)
cos_theta_t_ref = -ek.cos(28.1255057020557 * ek.pi / 180)
F_ref = 0.5 * (0.09201336304552442**2 + 0.3033370452904235**2)
L = (scale * ek.sqrt(1 - ek.cos(45 * ek.pi / 180)**2))**2 + cos_theta_t**2
assert ek.allclose(L, 1)
assert ek.allclose(cos_theta_t, cos_theta_t_ref)
assert ek.allclose(F, F_ref)
# 1.5 -> 1
F, cos_theta_t, _, _ = fresnel(ek.cos(45 * ek.pi / 180), 1 / 1.5)
assert ek.allclose(F, 1)
assert ek.allclose(cos_theta_t, 0)
F, cos_theta_t, _, scale = fresnel(ek.cos(10 * ek.pi / 180), 1 / 1.5)
cos_theta_t_ref = -ek.cos(15.098086605159006 * ek.pi / 180)
F_ref = 0.5 * (0.19046797197779405**2 + 0.20949431963852014**2)
L = (scale * ek.sqrt(1 - ek.cos(10 * ek.pi / 180)**2))**2 + cos_theta_t**2
assert ek.allclose(L, 1)
assert ek.allclose(cos_theta_t, cos_theta_t_ref)
assert ek.allclose(F, F_ref)
def sph_convert(v):
x2 = ek.pow(v.x, 2)
y2 = ek.pow(v.y, 2)
z2 = ek.pow(v.z, 2)
r = ek.sqrt(x2 + y2 + z2)
phi = ek.atan2(v.y, v.x)
theta = ek.atan2(ek.sqrt(x2 + y2), v.z)
return r, theta, phi
def acosh_(a0):
if not a0.IsFloat:
raise Exception("acosh(): requires floating point operands!")
ar, sr = _check1(a0)
if not a0.IsSpecial:
for i in range(sr):
ar[i] = _ek.acosh(a0[i])
elif a0.IsComplex:
return 2 * _ek.log(_ek.sqrt(.5 * (a0 + 1)) + _ek.sqrt(.5 * (a0 - 1)))
else:
raise Exception("acosh(): unsupported array type!")
return ar
def test01_gauss_lobatto(variant_scalar_rgb):
from mitsuba.core.quad import gauss_lobatto
assert ek.allclose(gauss_lobatto(2), [[-1, 1], [1.0, 1.0]])
assert ek.allclose(gauss_lobatto(3),
[[-1, 0, 1], [1.0 / 3.0, 4.0 / 3.0, 1.0 / 3.0]])
assert ek.allclose(gauss_lobatto(4),
[[-1, -ek.sqrt(1.0 / 5.0),
ek.sqrt(1.0 / 5.0), 1],
[1.0 / 6.0, 5.0 / 6.0, 5.0 / 6.0, 1.0 / 6.0]])
assert ek.allclose(
gauss_lobatto(5),
[[-1, -ek.sqrt(3.0 / 7.0), 0,
ek.sqrt(3.0 / 7.0), 1],
[1.0 / 10.0, 49.0 / 90.0, 32.0 / 45.0, 49.0 / 90.0, 1.0 / 10.0]])
def test05_specular_reflection(variant_scalar_rgb):
from mitsuba.core import Matrix4f
from mitsuba.render.mueller import specular_reflection
import numpy as np
# Identity matrix (and no phase shift) at perpendicular incidence
assert ek.allclose(specular_reflection(1, 1.5), Matrix4f.identity()*0.04)
assert ek.allclose(specular_reflection(1, 1/1.5), Matrix4f.identity()*0.04)
# 180 deg phase shift at grazing incidence
assert ek.allclose(specular_reflection(0, 1.5), [[1,0,0,0],[0,1,0,0],[0,0,-1,0],[0,0,0,-1]], atol=1e-6)
assert ek.allclose(specular_reflection(0, 1/1.5), [[1,0,0,0],[0,1,0,0],[0,0,-1,0],[0,0,0,-1]], atol=1e-6)
# Perfect linear polarization at Brewster's angle
M = np.zeros((4, 4))
M[0:2, 0:2] = 0.0739645
assert ek.allclose(specular_reflection(ek.cos(ek.atan(1/1.5)), 1/1.5), M, atol=1e-6)
assert ek.allclose(specular_reflection(ek.cos(ek.atan(1.5)), 1.5), M, atol=1e-6)
# 180 deg phase shift just below Brewster's angle (air -> glass)
M = specular_reflection(ek.cos(ek.atan(1.5)+1e-4), 1.5)
assert M[0, 0] > 0 and M[1, 1] > 0 and M[2, 2] < 0 and M[3, 3] < 0
M = specular_reflection(ek.cos(ek.atan(1/1.5)+1e-4), 1/1.5)
assert M[0, 0] > 0 and M[1, 1] > 0 and M[2, 2] < 0 and M[3, 3] < 0
# Complex phase shift in the total internal reflection case
eta = 1/1.5
cos_theta_min = ek.sqrt((1 - eta**2) / (1 + eta**2))
M = specular_reflection(cos_theta_min, eta).numpy()
assert np.all(M[0:2, 2:4] == 0) and np.all(M[2:4, 0:2] == 0) and ek.allclose(M[0:2, 0:2], np.eye(2), atol=1e-6)
phi_delta = 4*ek.atan(eta)
assert ek.allclose(M[2:4, 2:4], [[ek.cos(phi_delta), ek.sin(phi_delta)], [-ek.sin(phi_delta), ek.cos(phi_delta)]])
def test15_test_avx512_approx():
Float = get_class('enoki.llvm.Float')
x = ek.linspace(Float, 0, 10, 1000)
o = ek.full(Float, 1, 1000)
assert ek.allclose(ek.rsqrt(x), o / ek.sqrt(x), rtol=2e-7, atol=0)
assert ek.allclose(ek.rcp(x), o / x, rtol=2e-7, atol=0)
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 test03_type_promotion():
import enoki.packet as s
assert type(s.Array1b() | s.Array1b()) is s.Array1b
assert type(s.Array1b() ^ s.Array1b()) is s.Array1b
assert type(s.Array1b() & s.Array1b()) is s.Array1b
with pytest.raises(ek.Exception) as ei:
_ = s.Array1b() + s.Array1b()
assert "add(): requires arithmetic operands!" in str(ei.value)
assert type(s.Array1i() + s.Array1i()) is s.Array1i
assert type(s.Array1i() + s.Array1u()) is s.Array1u
assert type(s.Array1i() + s.Array1i64()) is s.Array1i64
assert type(s.Array1i() + s.Array1u64()) is s.Array1u64
assert type(s.Array1i() + s.Array1f()) is s.Array1f
assert type(s.Array1i() + s.Array1f64()) is s.Array1f64
with pytest.raises(ek.Exception) as ei:
_ = ek.sqrt(s.Array1i())
assert "sqrt(): requires floating point operands!" in str(ei.value)
assert type(s.Array1f() + s.Array1f()) is s.Array1f
assert type(s.Array1f() + s.Array1u()) is s.Array1f
assert type(s.Array1f() + s.Array1i64()) is s.Array1f
assert type(s.Array1f() + s.Array1u64()) is s.Array1f
assert type(s.Array1f() + s.Array1f()) is s.Array1f
assert type(s.Array1f() + s.Array1f64()) is s.Array1f64
assert type(s.Array1f() | s.Array1b()) is s.Array1f
assert type(s.Array1f() ^ s.Array1b()) is s.Array1f
assert type(s.Array1f() & s.Array1b()) is s.Array1f
def test02_unit_frame(variant_scalar_rgb):
from mitsuba.core import Frame3f, Vector2f, Vector3f
for theta in [30 * mitsuba.core.math.Pi / 180, 95 * mitsuba.core.math.Pi / 180]:
phi = 73 * mitsuba.core.math.Pi / 180
sin_theta, cos_theta = ek.sin(theta), ek.cos(theta)
sin_phi, cos_phi = ek.sin(phi), ek.cos(phi)
v = Vector3f(
cos_phi * sin_theta,
sin_phi * sin_theta,
cos_theta
)
f = Frame3f(Vector3f(1.0, 2.0, 3.0) / ek.sqrt(14))
v2 = f.to_local(v)
v3 = f.to_world(v2)
assert ek.allclose(v3, v)
assert ek.allclose(Frame3f.cos_theta(v), cos_theta)
assert ek.allclose(Frame3f.sin_theta(v), sin_theta)
assert ek.allclose(Frame3f.cos_phi(v), cos_phi)
assert ek.allclose(Frame3f.sin_phi(v), sin_phi)
assert ek.allclose(Frame3f.cos_theta_2(v), cos_theta * cos_theta)
assert ek.allclose(Frame3f.sin_theta_2(v), sin_theta * sin_theta)
assert ek.allclose(Frame3f.cos_phi_2(v), cos_phi * cos_phi)
assert ek.allclose(Frame3f.sin_phi_2(v), sin_phi * sin_phi)
assert ek.allclose(Vector2f(Frame3f.sincos_phi(v)), [sin_phi, cos_phi])
assert ek.allclose(Vector2f(Frame3f.sincos_phi_2(v)), [sin_phi * sin_phi, cos_phi * cos_phi])
def test04_surface_area(variant_scalar_rgb):
from mitsuba.core import xml, Transform4f
# Unifomly-scaled rectangle
rect = xml.load_dict({
"type":
"rectangle",
"to_world":
Transform4f([[2, 0, 0, 0], [0, 2, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
})
assert ek.allclose(rect.surface_area(), 2.0 * 2.0 * 2.0 * 2.0)
# Rectangle sheared along the Z-axis
rect = xml.load_dict({
"type":
"rectangle",
"to_world":
Transform4f([[1, 0, 0, 0], [0, 1, 0, 0], [1, 0, 1, 0], [0, 0, 0, 1]])
})
assert ek.allclose(rect.surface_area(), 2.0 * 2.0 * ek.sqrt(2.0))
# Rectangle sheared along the X-axis (shouldn't affect surface_area)
rect = xml.load_dict({
"type":
"rectangle",
"to_world":
Transform4f([[1, 1, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
})
assert ek.allclose(rect.surface_area(), 2.0 * 2.0)
def test04_sample_direct(variant_scalar_rgb):
from mitsuba.core import xml, Ray3f
from mitsuba.render import Interaction3f
sphere = xml.load_dict({"type": "sphere"})
def sample_cone(sample, cos_theta_max):
cos_theta = (1 - sample[1]) + sample[1] * cos_theta_max
sin_theta = ek.sqrt(1 - cos_theta * cos_theta)
phi = 2 * ek.pi * sample[0]
s, c = ek.sin(phi), ek.cos(phi)
return [c * sin_theta, s * sin_theta, cos_theta]
it = Interaction3f.zero()
it.p = [0, 0, -3]
it.t = 0
sin_cone_angle = 1.0 / it.p[2]
cos_cone_angle = ek.sqrt(1 - sin_cone_angle**2)
for xi_1 in ek.linspace(Float, 0, 1, 10):
for xi_2 in ek.linspace(Float, 1e-3, 1 - 1e-3, 10):
sample = sphere.sample_direction(it, [xi_2, 1 - xi_1])
d = sample_cone([xi_1, xi_2], cos_cone_angle)
its = sphere.ray_intersect(Ray3f(it.p, d, 0, []))
assert ek.allclose(d, sample.d, atol=1e-5, rtol=1e-5)
assert ek.allclose(its.t, sample.dist, atol=1e-5, rtol=1e-5)
assert ek.allclose(its.p, sample.p, atol=1e-5, rtol=1e-5)
def test04_sample_direct(variant_scalar_rgb):
from mitsuba.core.xml import load_string
from mitsuba.core import Ray3f
from mitsuba.render import Interaction3f
if mitsuba.core.MTS_ENABLE_EMBREE:
pytest.skip("EMBREE enabled")
sphere = load_string('<shape type="sphere" version="2.0.0"/>')
def sample_cone(sample, cos_theta_max):
cos_theta = (1 - sample[1]) + sample[1] * cos_theta_max
sin_theta = ek.sqrt(1 - cos_theta * cos_theta)
phi = 2 * ek.pi * sample[0]
s, c = ek.sin(phi), ek.cos(phi)
return [c * sin_theta, s * sin_theta, cos_theta]
it = Interaction3f.zero()
it.p = [0, 0, -3]
it.t = 0
sin_cone_angle = 1.0 / it.p[2]
cos_cone_angle = ek.sqrt(1 - sin_cone_angle**2)
for xi_1 in ek.linspace(Float, 0, 1, 10):
for xi_2 in ek.linspace(Float, 1e-3, 1 - 1e-3, 10):
sample = sphere.sample_direction(it, [xi_2, 1 - xi_1])
d = sample_cone([xi_1, xi_2], cos_cone_angle)
its = sphere.ray_intersect(Ray3f(it.p, d, 0, []))
assert ek.allclose(d, sample.d, atol=1e-5, rtol=1e-5)
assert ek.allclose(its.t, sample.dist, atol=1e-5, rtol=1e-5)
assert ek.allclose(its.p, sample.p, atol=1e-5, rtol=1e-5)
def test13_sqrt(m):
x = m.Float(1, 4, 16)
ek.enable_grad(x)
y = ek.sqrt(x)
ek.backward(y)
assert ek.allclose(ek.detach(y), [1, 2, 4])
assert ek.allclose(ek.grad(x), [.5, .25, .125])
def test02_gauss_legendre(variant_scalar_rgb):
from mitsuba.core.quad import gauss_legendre
assert ek.allclose(gauss_legendre(1), [[0], [2]])
assert ek.allclose(
gauss_legendre(2),
[[-ek.sqrt(1.0 / 3.0), ek.sqrt(1.0 / 3.0)], [1, 1]])
assert ek.allclose(
gauss_legendre(3),
[[-ek.sqrt(3.0 / 5.0), 0, ek.sqrt(3.0 / 5.0)],
[5.0 / 9.0, 8.0 / 9.0, 5.0 / 9.0]])
assert ek.allclose(gauss_legendre(4), [[
-0.861136,
-0.339981,
0.339981,
0.861136,
], [0.347855, 0.652145, 0.652145, 0.347855]])
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 test02_sample_spectrum(variant_scalar_spectral, obj):
from mitsuba.render import SurfaceInteraction3f
si = SurfaceInteraction3f()
assert ek.allclose(obj.sample_spectrum(si, 0), [500, 150])
assert ek.allclose(obj.sample_spectrum(si, 1), [600, 150])
assert ek.allclose(
obj.sample_spectrum(si, .5),
[500 + 100 * (ek.sqrt(10) / 2 - 1), 150]
)
def rsqrt_(a0):
if not a0.IsFloat:
raise Exception("rsqrt(): requires floating point operands!")
if not a0.IsSpecial:
ar, sr = _check1(a0)
for i in range(sr):
ar[i] = _ek.rsqrt(a0[i])
return ar
else:
return _ek.rcp(_ek.sqrt(a0))
def test13_cont_func(variant_packet_rgb):
# Test continuous 1D distribution integral against analytic result
from mitsuba.core import ContinuousDistribution, Float
x = ek.linspace(Float, -2, 2, 513)
y = ek.exp(-ek.sqr(x))
d = ContinuousDistribution([-2, 2], y)
assert ek.allclose(d.integral(), ek.sqrt(ek.pi) * ek.erf(2.0))
assert ek.allclose(d.eval_pdf([1]), [ek.exp(-1)])
assert ek.allclose(d.sample([0, 0.5, 1]), [-2, 0, 2])
def asin_(a0):
if not a0.IsFloat:
raise Exception("asin(): requires floating point operands!")
ar, sr = _check1(a0)
if not a0.IsSpecial:
for i in range(sr):
ar[i] = _ek.asin(a0[i])
elif a0.IsSpecial:
tmp = _ek.log(type(a0)(-a0.imag, a0.real) + _ek.sqrt(1 - _ek.sqr(a0)))
ar.real = tmp.imag
ar.imag = -tmp.real
else:
raise Exception("asin(): unsupported array type!")
return ar
def test08_misc():
for i in range(-5, 5):
for j in range(-5, 5):
a = ek.sqrt(C(i, j))
b = cmath.sqrt(complex(i, j))
assert ek.allclose(a, C(b))
assert ek.allclose(ek.conj(a), C(b.conjugate()))
assert ek.allclose(ek.abs(a), C(abs(b)))
if i != 0 and j != 0:
a = ek.rsqrt(C(i, j))
b = C(1 / cmath.sqrt(complex(i, j)))
assert ek.allclose(a, b)
def step(self):
""" Take a gradient step """
self.t += 1
from enoki.cuda_autodiff import Float32 as Float
lr_t = ek.detach(
Float(self.lr * ek.sqrt(1 - self.beta_2**self.t) /
(1 - self.beta_1**self.t),
literal=False))
for k in self.bsdf_ad_keys:
g_p = ek.gradient(self.bsdf_map[k].reflectance.data)
size = ek.slices(g_p)
assert (size == ek.slices(self.state[k][0]))
m_tp, v_tp = self.state[k]
m_t = self.beta_1 * m_tp + (1 - self.beta_1) * g_p
v_t = self.beta_2 * v_tp + (1 - self.beta_2) * ek.sqr(g_p)
self.state[k] = (m_t, v_t)
u = ek.detach(self.bsdf_map[k].reflectance.data) - lr_t * m_t / (
ek.sqrt(v_t) + self.epsilon)
u = type(self.bsdf_map[k].reflectance.data)(u)
ek.set_requires_gradient(u)
self.bsdf_map[k].reflectance.data = u
for k in self.mesh_ad_keys:
g_p = ek.gradient(self.mesh_map[k].vertex_positions)
size = ek.slices(g_p)
assert (size == ek.slices(self.state[k][0]))
m_tp, v_tp = self.state[k]
m_t = self.beta_1 * m_tp + (1 - self.beta_1) * g_p
v_t = self.beta_2 * v_tp + (1 - self.beta_2) * ek.sqr(g_p)
self.state[k] = (m_t, v_t)
u = ek.detach(self.mesh_map[k].vertex_positions) - lr_t * m_t / (
ek.sqrt(v_t) + self.epsilon)
u = type(self.mesh_map[k].vertex_positions)(u)
ek.set_requires_gradient(u)
self.mesh_map[k].vertex_positions = u
def sqrt_(a0):
if not a0.IsFloat:
raise Exception("sqrt(): requires floating point operands!")
ar, sr = _check1(a0)
if not a0.IsSpecial:
for i in range(sr):
ar[i] = _ek.sqrt(a0[i])
elif a0.IsComplex:
n = abs(a0)
m = a0.real >= 0
zero = _ek.eq(n, 0)
t1 = _ek.sqrt(.5 * (n + abs(a0.real)))
t2 = .5 * a0.imag / t1
im = _ek.select(m, t2, _ek.copysign(t1, a0.imag))
ar.real = _ek.select(m, t1, abs(t2))
ar.imag = _ek.select(zero, 0, im)
elif a0.IsQuaternion:
ri = _ek.norm(a0.imag)
cs = _ek.sqrt(a0.Complex(a0.real, ri))
ar.imag = a0.imag * (_ek.rcp(ri) * cs.imag)
ar.real = cs.real
else:
raise Exception("sqrt(): unsupported array type!")
return ar
def test03_distance(variant_scalar_rgb):
from mitsuba.core import BoundingBox3f as BBox
assert BBox([1, 2, 3], [2, 3, 5]).distance(BBox([4, 2, 3], [5, 3, 5])) == 2
assert ek.abs(
BBox([1, 2, 3], [2, 3, 5]).distance(BBox([3, 4, 6], [7, 7, 7])) -
ek.sqrt(3)) < 1e-6
assert BBox([1, 2, 3],
[2, 3, 5]).distance(BBox([1.1, 2.2, 3.3],
[1.8, 2.8, 3.8])) == 0
assert BBox([1, 2, 3], [2, 3, 5]).distance([1.5, 2.5, 3.5]) == 0
assert BBox([1, 2, 3], [2, 3, 5]).distance([3, 2.5, 3.5]) == 1
def test40_safe_functions(m):
x = ek.linspace(m.Float, 0, 1, 10)
y = ek.linspace(m.Float, -1, 1, 10)
z = ek.linspace(m.Float, -1, 1, 10)
ek.enable_grad(x, y, z)
x2 = ek.safe_sqrt(x)
y2 = ek.safe_acos(y)
z2 = ek.safe_asin(z)
ek.backward(x2)
ek.backward(y2)
ek.backward(z2)
assert ek.grad(x)[0] == 0
assert ek.allclose(ek.grad(x)[1], .5 / ek.sqrt(1 / 9))
assert x[0] == 0
assert ek.all(ek.isfinite(ek.grad(x)))
assert ek.all(ek.isfinite(ek.grad(y)))
assert ek.all(ek.isfinite(ek.grad(z)))
def test12_cont_eval(variant_packet_rgb):
# Test continuous 1D distribution pdf/cdf against hand-computed reference
from mitsuba.core import ContinuousDistribution
d = ContinuousDistribution([2, 3], [1, 2])
eps = 1e-6
assert ek.allclose(d.integral(), 3.0 / 2.0)
assert ek.allclose(d.normalization(), 2.0 / 3.0)
assert ek.allclose(
d.eval_pdf_normalized([1, 2 - eps, 2, 2.5, 3, 3 + eps, 4]),
[0, 0, 2.0 / 3.0, 1.0, 4.0 / 3.0, 0, 0])
assert ek.allclose(d.eval_cdf_normalized([1, 2, 2.5, 3, 4]),
[0, 0, 5.0 / 12.0, 1, 1])
assert d.sample([0, 1]) == [2, 3]
x, pdf = d.sample_pdf([0, 0.5, 1])
dx = (ek.sqrt(10) - 2) / 2
assert x == [2, 2 + dx, 3]
assert ek.allclose(pdf,
[2.0 / 3.0, (4 * dx + 2 * (1 - dx)) / 3.0, 4.0 / 3.0])
def test06_phase_tir(variant_scalar_rgb):
import numpy as np
from mitsuba.render import fresnel_polarized
eta = 1 / 1.5
# Check phase shift at critical angle (total internal reflection case)
crit = ek.asin(eta)
a_s, a_p, _, _, _ = fresnel_polarized(ek.cos(crit), eta)
assert ek.allclose(ek.arg(a_s), 0.0)
assert ek.allclose(ek.arg(a_p), ek.pi) or ek.allclose(ek.arg(a_p), -ek.pi)
# Check phase shift at grazing angle (total internal reflection case)
a_s, a_p, _, _, _ = fresnel_polarized(0.0, eta)
assert ek.allclose(ek.arg(a_s), ek.pi) or ek.allclose(ek.arg(a_s), -ek.pi)
assert ek.allclose(ek.arg(a_p), 0.0)
# Location of minimum phase difference
cos_theta_min = ek.sqrt((1 - eta**2) / (1 + eta**2))
phi_delta = 4 * ek.atan(eta)
a_s, a_p, _, _, _ = fresnel_polarized(cos_theta_min, eta)
assert ek.allclose(ek.arg(a_s) - ek.arg(a_p), phi_delta)
def sample_cone(sample, cos_theta_max):
cos_theta = (1 - sample[1]) + sample[1] * cos_theta_max
sin_theta = ek.sqrt(1 - cos_theta * cos_theta)
phi = 2 * ek.pi * sample[0]
s, c = ek.sin(phi), ek.cos(phi)
return [c * sin_theta, s * sin_theta, cos_theta]
def stdnormal_cdf(x):
shape = x.shape
cdf = (1.0 - ek.erf(-x.flatten() / ek.sqrt(2.0))) * 0.5
return np.array(cdf).reshape(shape)
