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 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 test24_log(m): x = ek.linspace(m.Float, 0.01, 1, 10) ek.enable_grad(x) y = ek.log(x * x) ek.backward(y) log_x = ek.log(ek.sqr(ek.detach(x))) assert ek.allclose(y, log_x) assert ek.allclose(ek.grad(x), 2 / ek.detach(x))
def test23_exp(m): x = ek.linspace(m.Float, 0, 1, 10) ek.enable_grad(x) y = ek.exp(x * x) ek.backward(y) exp_x = ek.exp(ek.sqr(ek.detach(x))) assert ek.allclose(y, exp_x) assert ek.allclose(ek.grad(x), 2 * ek.detach(x) * exp_x)
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 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 test29_asin(m): x = ek.linspace(m.Float, -.8, .8, 10) ek.enable_grad(x) y = ek.asin(x * x) ek.backward(y) asin_x = ek.asin(ek.sqr(ek.detach(x))) assert ek.allclose(y, asin_x) assert ek.allclose( ek.grad(x), m.Float(-2.08232, -1.3497, -0.906755, -0.534687, -0.177783, 0.177783, 0.534687, 0.906755, 1.3497, 2.08232))
def test30_acos(m): x = ek.linspace(m.Float, -.8, .8, 10) ek.enable_grad(x) y = ek.acos(x * x) ek.backward(y) acos_x = ek.acos(ek.sqr(ek.detach(x))) assert ek.allclose(y, acos_x) assert ek.allclose( ek.grad(x), m.Float(2.08232, 1.3497, 0.906755, 0.534687, 0.177783, -0.177783, -0.534687, -0.906755, -1.3497, -2.08232))
def test31_atan(m): x = ek.linspace(m.Float, -.8, .8, 10) ek.enable_grad(x) y = ek.atan(x * x) ek.backward(y) atan_x = ek.atan(ek.sqr(ek.detach(x))) assert ek.allclose(y, atan_x) assert ek.allclose( ek.grad(x), m.Float(-1.13507, -1.08223, -0.855508, -0.53065, -0.177767, 0.177767, 0.53065, 0.855508, 1.08223, 1.13507))
def test28_tan(m): x = ek.linspace(m.Float, 0, 1, 10) ek.enable_grad(x) y = ek.tan(x * x) ek.backward(y) tan_x = ek.tan(ek.sqr(ek.detach(x))) assert ek.allclose(y, tan_x) assert ek.allclose( ek.grad(x), m.Float(0., 0.222256, 0.44553, 0.674965, 0.924494, 1.22406, 1.63572, 2.29919, 3.58948, 6.85104))
def test27_sec(m): x = ek.linspace(m.Float, 1, 2, 10) ek.enable_grad(x) y = ek.sec(x * x) ek.backward(y) sec_x = ek.sec(ek.sqr(ek.detach(x))) assert ek.allclose(y, sec_x) assert ek.allclose( ek.grad(x), m.Float(5.76495, 19.2717, 412.208, 61.794, 10.3374, 3.64885, 1.35811, -0.0672242, -1.88437, -7.08534))
def asinh_(a0): if not a0.IsFloat: raise Exception("asinh(): requires floating point operands!") ar, sr = _check1(a0) if not a0.IsSpecial: for i in range(sr): ar[i] = _ek.asinh(a0[i]) elif a0.IsComplex: return _ek.log(a0 + _ek.sqrt(_ek.sqr(a0) + 1)) else: raise Exception("asinh(): unsupported array type!") return ar
def test26_csc(m): x = ek.linspace(m.Float, 1, 2, 10) ek.enable_grad(x) y = ek.csc(x * x) ek.backward(y) csc_x = ek.csc(ek.sqr(ek.detach(x))) assert ek.allclose(y, csc_x) assert ek.allclose(ek.grad(x), m.Float(-1.52612, -0.822733, -0.189079, 0.572183, 1.88201, 5.34839, 24.6017, 9951.86, 20.1158, 4.56495), rtol=5e-5)
def test28_cot(m): x = ek.linspace(m.Float, 1, 2, 10) ek.enable_grad(x) y = ek.cot(x * x) ek.backward(y) cot_x = ek.cot(ek.sqr(ek.detach(x))) assert ek.allclose(y, cot_x) assert ek.allclose(ek.grad(x), m.Float(-2.82457, -2.49367, -2.45898, -2.78425, -3.81687, -7.12557, -26.3248, -9953.63, -22.0932, -6.98385), rtol=5e-5)
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 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 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)
# Construct an Adam optimizer that will adjust the parameters 'params' opt = Adam(params, lr=.02) time_a = time.time() iterations = 100 for it in range(iterations): # Perform a differentiable rendering of the scene image = render(scene, optimizer=opt, unbiased=True, spp=1) write_bitmap('out_%03i.png' % it, image, crop_size) write_bitmap('envmap_%03i.png' % it, params['my_envmap.data'], (param_res[0], param_res[1])) # Objective: MSE between 'image' and 'image_ref' ob_val = ek.hsum(ek.sqr(image - image_ref)) / len(image) # Back-propagate errors to input parameters ek.backward(ob_val) # Optimizer: take a gradient step opt.step() # Compare iterate against ground-truth value err_ref = ek.hsum(ek.sqr(param_ref - params['my_envmap.data'])) print('Iteration %03i: error=%g' % (it, err_ref[0]), end='\r') time_b = time.time() print() print('%f ms per iteration' % (((time_b - time_a) * 1000) / iterations))
def K(self, x, m_): return ek.rsqrt(1 - m_ * ek.sqr(ek.sin(x)))
def backward(self): grad_out = self.grad_out() grad_in = grad_out * self.inv_norm grad_in -= self.value * (ek.dot(self.value, grad_in) * ek.sqr(self.inv_norm)) self.set_grad_in('value', grad_in)
def forward(self): grad_in = self.grad_in('value') grad_out = grad_in * self.inv_norm grad_out -= self.value * (ek.dot(self.value, grad_out) * ek.sqr(self.inv_norm)) self.set_grad_out(grad_out)
params['red.reflectance.value'] = [.9, .9, .9] params.update() # Construct an Adam optimizer that will adjust the parameters 'params' opt = Adam(params, lr=.2) time_a = time.time() iterations = 100 for it in range(iterations): # Perform a differentiable rendering of the scene image = render(scene, optimizer=opt, unbiased=True, spp=1) write_bitmap('out_%03i.png' % it, image, crop_size) # Objective: MSE between 'image' and 'image_ref' ob_val = ek.hsum(ek.sqr(image - image_ref)) / len(image) # Back-propagate errors to input parameters ek.backward(ob_val) # Optimizer: take a gradient step opt.step() # Compare iterate against ground-truth value err_ref = ek.hsum(ek.sqr(param_ref - params['red.reflectance.value'])) print('Iteration %03i: error=%g' % (it, err_ref[0]), end='\r') time_b = time.time() print() print('%f ms per iteration' % (((time_b - time_a) * 1000) / iterations))
def render_sample(scene, sampler, rays, bdata, heightmap_pybind, bssrdf=None): """ Sample RTE TODO: Support multi channel sampling Args: scene: Target scene object sampler: Sampler object for random number rays: Given rays for sampling bdata: BSSRDF Data object heightmap_pybind: Object for getting height map around incident position. Refer src/librender/python/heightmap.cpp Returns: result: Sampling RTE result valid_rays: Mask data whether rays are valid or not scatter: Scatter components of Sampling RTE result non_scatter: Non scatter components of Sampling RTE result invalid_sample: Sampling RTE result with invalid sampled data by VAEBSSRDF """ eta = Float(1.0) emission_weight = Float(1.0) throughput = Spectrum(1.0) result = Spectrum(0.0) scatter = Spectrum(0.0) non_scatter = Spectrum(0.0) invalid_sample = Spectrum(0.0) active = True is_bssrdf = False ##### First interaction ##### si = scene.ray_intersect(rays, active) active = si.is_valid() & active valid_rays = si.is_valid() emitter = si.emitter(scene, active) depth = 0 # Set channel # At and after evaluating BSSRDF, a ray consider only this one channel n_channels = 3 channel = UInt32( ek.min(sampler.next_1d(active) * n_channels, n_channels - 1)) d_out_local = Vector3f().zero() d_out_pdf = Float(0) sss = Mask(False) while (True): depth += 1 if config.aovs and depth == 2: sss = is_bssrdf ##### Interaction with emitters ##### emission_val = emission_weight * throughput * Emitter.eval_vec( emitter, si, active) result += ek.select(active, emission_val, Spectrum(0.0)) invalid_sample += ek.select(active, emission_val, Spectrum(0.0)) scatter += ek.select(active & sss, emission_val, Spectrum(0.0)) non_scatter += ek.select(active & ~sss, emission_val, Spectrum(0.0)) active = active & si.is_valid() # Process russian roulette if depth > config.rr_depth: q = ek.min(ek.hmax(throughput) * ek.sqr(eta), 0.95) active = active & (sampler.next_1d(active) < q) throughput *= ek.rcp(q) # Stop if the number of bouces exceeds the given limit bounce, or # all rays are invalid. latter check is done only when the limit # bounce is infinite if depth >= config.max_depth: break ##### Emitter sampling ##### bsdf = si.bsdf(rays) ctx = BSDFContext() active_e = active & has_flag(BSDF.flags_vec(bsdf), BSDFFlags.Smooth) ds, emitter_val = scene.sample_emitter_direction( si, sampler.next_2d(active_e), True, active_e) active_e &= ek.neq(ds.pdf, 0.0) # Query the BSDF for that emitter-sampled direction wo = si.to_local(ds.d) bsdf_val = BSDF.eval_vec(bsdf, ctx, si, wo, active_e) # Determine density of sampling that same direction using BSDF sampling bsdf_pdf = BSDF.pdf_vec(bsdf, ctx, si, wo, active_e) mis = ek.select(ds.delta, Float(1), mis_weight(ds.pdf, bsdf_pdf)) emission_val = mis * throughput * bsdf_val * emitter_val result += ek.select(active, emission_val, Spectrum(0.0)) invalid_sample += ek.select(active, emission_val, Spectrum(0.0)) scatter += ek.select(active & sss, emission_val, Spectrum(0.0)) non_scatter += ek.select(active & ~sss, emission_val, Spectrum(0.0)) ##### BSDF sampling ##### bs, bsdf_val = BSDF.sample_vec(bsdf, ctx, si, sampler.next_1d(active), sampler.next_2d(active), active) ##### BSSRDF replacing ##### if (config.enable_bssrdf): # Replace bsdf samples by ones of BSSRDF bs.wo = ek.select(is_bssrdf, d_out_local, bs.wo) bs.pdf = ek.select(is_bssrdf, d_out_pdf, bs.pdf) bs.sampled_component = ek.select(is_bssrdf, UInt32(1), bs.sampled_component) bs.sampled_type = ek.select(is_bssrdf, UInt32(+BSDFFlags.DeltaTransmission), bs.sampled_type) ############################ throughput *= ek.select(is_bssrdf, Float(1.0), bsdf_val) active &= ek.any(ek.neq(throughput, 0)) eta *= bs.eta # Intersect the BSDF ray against the scene geometry rays = RayDifferential3f(si.spawn_ray(si.to_world(bs.wo))) si_bsdf = scene.ray_intersect(rays, active) ##### Checking BSSRDF ##### if (config.enable_bssrdf): # Whether the BSDF is BSS RDF or not? is_bssrdf = (active & has_flag(BSDF.flags_vec(bsdf), BSDFFlags.BSSRDF) & (Frame3f.cos_theta(bs.wo) < Float(0.0)) & (Frame3f.cos_theta(si.wi) > Float(0.0))) # Decide whether we should use 0-scattering or multiple scattering is_zero_scatter = utils_render.check_zero_scatter( sampler, si_bsdf, bs, channel, is_bssrdf) is_bssrdf = is_bssrdf & ~is_zero_scatter throughput *= ek.select(is_bssrdf, ek.sqr(bs.eta), Float(1.0)) ########################### ###### Process for BSSRDF ###### if (config.enable_bssrdf and not ek.none(is_bssrdf)): # Get projected samples from BSSRDF projected_si, project_suc, abs_prob = bssrdf.sample_bssrdf( scene, bsdf, bs, si, bdata, heightmap_pybind, channel, is_bssrdf) if config.visualize_invalid_sample and (depth <= 1): active = active & (~is_bssrdf | project_suc) invalid_sample += ek.select((is_bssrdf & (~project_suc)), Spectrum([100, 0, 0]), Spectrum(0.0)) # Sample outgoing direction from projected position d_out_local, d_out_pdf = utils_render.resample_wo( sampler, is_bssrdf) # Apply absorption probability throughput *= ek.select(is_bssrdf, Spectrum(1) - abs_prob, Spectrum(1)) # Replace interactions by sampled ones from BSSRDF si_bsdf = SurfaceInteraction3f().masked_si(si_bsdf, projected_si, is_bssrdf) ################################ # Determine probability of having sampled that same # direction using emitter sampling emitter = si_bsdf.emitter(scene, active) ds = DirectionSample3f(si_bsdf, si) ds.object = emitter delta = has_flag(bs.sampled_type, BSDFFlags.Delta) emitter_pdf = ek.select(delta, Float(0.0), scene.pdf_emitter_direction(si, ds)) emission_weight = mis_weight(bs.pdf, emitter_pdf) si = si_bsdf return result, valid_rays, scatter, non_scatter, invalid_sample
def poly2(x, c0, c1, c2): x2 = ek.sqr(x) return ek.fmadd(x2, c2, ek.fmadd(x, c1, c0))
def elevation(d): dist = ek.sqrt(ek.sqr(d.x) + ek.sqr(d.y) + ek.sqr(d.z - 1.)) return 2. * ek.safe_asin(.5 * dist)
if ek.any(ek.isnan(image)): print("[WARNING] NaNs in the image.") # Write a gamma encoded PNG image_np = image.numpy().reshape(height, width, 3) output_file = output_path + 'out_%03i.png' % i print("Writing image %s" % (output_file)) Bitmap(image_np).convert(pixel_format=Bitmap.PixelFormat.RGB, component_format=Struct.Type.UInt8, srgb_gamma=True).write(output_file) # Objective function if task == 'plain2bumpy' or task == 'bumpy2bumpy': loss = ek.hsum(ek.hsum( ek.sqr(image - image_ref))) / (height * width * 3) if task == 'bumpy2plain': loss = ek.hsum(ek.hsum( ek.sqr(image - image_plain))) / (height * width * 3) print("Iteration %i: loss=%f" % (i, loss[0])) loss_list.append(loss[0]) diff_image_init = ek.hsum(ek.hsum( ek.sqr(image - image_init))) / (height * width * 3) #print("difference with initial image = %f" % (diff_image_init[0])) diff_render_init.append(diff_image_init[0]) diff_init = ek.hsum( ek.sqr(params_opt['displacements'] - disp_tex_data_init * amp)) / ek.slices(params_opt['displacements'])
film = scene.sensors()[0].film() crop_size = film.crop_size() iterations = render_config['num_iteration'] for it in range(iterations): # Perform a differentiable rendering of the scene image = render(scene, optimizer=opt, unbiased=True, spp=4) if it % 10 == 0: write_bitmap(os.path.join(outpath, 'out_%04i.png' % it), image, crop_size) write_bitmap(os.path.join(outpath, 'texture_%04i.png' % it), params[opt_param_name], (param_res[1], param_res[0])) # Objective : MSE between 'image' and 'image_ref' ob_val = ek.hsum(ek.sqr(image - image_ref)) / len(image) # Back-propropagate errors to input parameters ek.backward(ob_val) # Optimizer : take a gradient step opt.step() # Compare iterate against ground-truth value err_ref = ek.hsum(ek.sqr(param_ref - params[opt_param_name])) losses = np.append(losses, err_ref) print('Iteration %04i: error=%g' % (it, err_ref[0]), end='\r') time_b = time.time() np.savetxt(os.path.join(outpath, 'losses.csv'), losses, delimiter=",")
mask_len = ek.hsum(masks[diff_param_owner])[0] errors = list() converged = False it = 0 while converged != True and it <= 100: # Perform a differentiable rendering of the scene image = render(scene, optimizer=opt, unbiased=True, spp=1) write_bitmap(dump_path + 'out_%03i.png' % it, image, crop_size) # Objective: MSE between 'image' and 'image_ref' #ob_val = ek.hsum(ek.sqr(image - image_ref)) / len(image) if use_masks: ob_val = ek.hsum( masks[diff_param_owner] * ek.sqr(image - image_ref)) / mask_len else: ob_val = ek.hsum(ek.sqr(image - image_ref)) / len(image) # Back-propagate errors to input parameters ek.backward(ob_val) # Optimizer: take a gradient step opt.step() err_ref = ek.hsum(ek.sqr(diff_param_ref - params[diff_param_key])) #if err_ref[0] < 0.0001: #converged = True errors.append(err_ref[0]) print('Iteration %03i : error= %g' % (it, err_ref[0])) it += 1
if ek.any(ek.any(ek.isnan(params[vertex_pos_key]))): print("[WARNING] NaNs in the vertex positions.") if ek.any(ek.isnan(image)): print("[WARNING] NaNs in the image.") # Write a gamma encoded PNG image_np = image.numpy().reshape(height, width, 3) output_file = output_path + 'out_%03i_%03i.png' % (j, i) print("Writing image %s" % (output_file)) Bitmap(image_np).convert(pixel_format=Bitmap.PixelFormat.RGB, component_format=Struct.Type.UInt8, srgb_gamma=True).write(output_file) # Objective function if task == 'plain2bumpy' or task == 'bumpy2bumpy': loss = ek.hsum(ek.hsum(ek.sqr(image - image_ref))) / (height*width*3) if task == 'bumpy2plain': loss = ek.hsum(ek.hsum(ek.sqr(image - image_plain))) / (height*width*3) print("Iteration %i-%i: loss=%f" % (j, i, loss[0])) loss_list.append(loss[0]) diff_image_init = ek.hsum(ek.hsum(ek.sqr(image - image_init))) / (height*width*3) #print("difference with initial image = %f" % (diff_image_init[0])) diff_render_init.append(diff_image_init[0]) diff_init = ek.hsum(ek.sqr(params_opt['displacements'] - disp_tex_data_init * amp)) / ek.slices(params_opt['displacements']) diff_ref = ek.hsum(ek.sqr(params_opt['displacements'] - disp_tex_data_ref * amp)) / ek.slices(params_opt['displacements']) #print("diff_init:", diff_init[0]) #print("diff_ref:", diff_ref[0])