def sample_overlap_worker(wfs, configs, pgrad, nsteps, tstep=0.5): r"""Run nstep Metropolis steps to sample a distribution proportional to :math:`\sum_i |\Psi_i|^2`, where :math:`\Psi_i` = wfs[i] """ nconf, nelec, _ = configs.configs.shape for wf in wfs: wf.recompute(configs) block_avg = {} block_avg["acceptance"] = np.zeros(nsteps) for n in range(nsteps): for e in range(nelec): # a sweep # Propose move grads = [ np.real(wf.gradient(e, configs.electron(e)).T) for wf in wfs ] grad = mc.limdrift(np.mean(grads, axis=0)) gauss = np.random.normal(scale=np.sqrt(tstep), size=(nconf, 3)) newcoorde = configs.configs[:, e, :] + gauss + grad * tstep newcoorde = configs.make_irreducible(e, newcoorde) # Compute reverse move grads, vals = list( zip(*[wf.gradient_value(e, newcoorde) for wf in wfs])) grads = [np.real(g.T) for g in grads] new_grad = mc.limdrift(np.mean(grads, axis=0)) forward = np.sum(gauss**2, axis=1) backward = np.sum((gauss + tstep * (grad + new_grad))**2, axis=1) # Acceptance t_prob = np.exp(1 / (2 * tstep) * (forward - backward)) wf_ratios = np.abs(vals)**2 log_values = np.real(np.array([wf.value()[1] for wf in wfs])) weights = np.exp(2 * (log_values - log_values[0])) ratio = t_prob * np.sum(wf_ratios * weights, axis=0) / weights.sum(axis=0) accept = ratio > np.random.rand(nconf) block_avg["acceptance"][n] += accept.mean() / nelec # Update wave function configs.move(e, newcoorde, accept) for wf in wfs: wf.updateinternals(e, newcoorde, configs, mask=accept) # Collect rolling average save_dat = collect_overlap_data(wfs, configs, pgrad) for k, it in save_dat.items(): if k not in block_avg: block_avg[k] = np.zeros((*it.shape, ), dtype=it.dtype) if k in ["overlap", "overlap_gradient", "weight_final"]: block_avg[k] += save_dat[k] / nsteps else: block_avg[k] += save_dat[k] * save_dat["weight_final"] / nsteps for k, it in block_avg.items(): if k not in [ "overlap", "overlap_gradient", "weight_final", "acceptance" ]: it /= block_avg["weight_final"] return block_avg, configs
def limdrift_cutoff(g, tau, cutoff=1): """ Limit a vector to have a maximum magnitude of cutoff while maintaining direction Args: g: a [nconf,ndim] vector cutoff: the maximum magnitude Returns: The vector with the cut off applied and multiplied by tau. """ return mc.limdrift(g, cutoff) * tau
def sample_overlap(wfs, configs, pgrad, nsteps=100, tstep=0.5): r""" Sample .. math:: \rho(R) = \sum_i |\Psi_i(R)|^2 `pgrad` is expected to be a gradient generator. returns data as follows: `overlap` : .. math:: \left\langle \frac{\Psi_i^* \Psi_j}{\rho} \right\rangle `overlap_gradient`: .. math:: \left\langle \frac{\partial_{im} \Psi_i^* \Psi_j}{\rho} \right\rangle Note that overlap_gradient is only saved for i = f, where f is the final wave function. In addition, any key returned by `pgrad` will be saved for the final wave function. """ nconf, nelec, ndim = configs.configs.shape for wf in wfs: wf.recompute(configs) return_data = {} for step in range(nsteps): # print("step", step) for e in range(nelec): # Propose move grads = [ np.real(wf.gradient(e, configs.electron(e)).T) for wf in wfs ] grad = limdrift(np.mean(grads, axis=0)) gauss = np.random.normal(scale=np.sqrt(tstep), size=(nconf, 3)) newcoorde = configs.configs[:, e, :] + gauss + grad * tstep newcoorde = configs.make_irreducible(e, newcoorde) # Compute reverse move grads = [np.real(wf.gradient(e, newcoorde).T) for wf in wfs] new_grad = limdrift(np.mean(grads, axis=0)) forward = np.sum(gauss**2, axis=1) backward = np.sum((gauss + tstep * (grad + new_grad))**2, axis=1) # Acceptance t_prob = np.exp(1 / (2 * tstep) * (forward - backward)) wf_ratios = np.array([wf.testvalue(e, newcoorde)**2 for wf in wfs]) log_values = np.array([wf.value()[1] for wf in wfs]) ref = log_values[0] weights = np.exp(2 * (log_values - ref)) ratio = (t_prob * np.sum(wf_ratios * weights, axis=0) / np.sum(weights, axis=0)) accept = ratio > np.random.rand(nconf) # Update wave function configs.move(e, newcoorde, accept) for wf in wfs: wf.updateinternals(e, newcoorde, mask=accept) # print("accept", np.mean(accept)) log_values = np.array([wf.value() for wf in wfs]) # print(log_values.shape) ref = np.max(log_values[:, 1, :], axis=0) save_dat = {} denominator = np.sum(np.exp(2 * (log_values[:, 1, :] - ref)), axis=0) normalized_values = log_values[:, 0, :] * np.exp(log_values[:, 1, :] - ref) save_dat["overlap"] = np.mean( np.einsum("ik,jk->ijk", normalized_values, normalized_values) / denominator, axis=-1, ) weight = np.array([ np.exp(-2 * (log_values[i, 1, :] - log_values[:, 1, :])) for i in range(len(wfs)) ]) weight = 1.0 / np.sum(weight, axis=1) dat = pgrad(configs, wfs[-1]) save_dat["overlap_gradient"] = np.mean( np.einsum("km,k,jk->jmk", dat['dppsi'], normalized_values[-1], normalized_values) / denominator, axis=-1, ) for k in dat.keys(): save_dat[k] = np.average(dat[k], axis=0, weights=weight[-1]) save_dat["weight"] = np.mean(weight, axis=1) for k, it in save_dat.items(): if k not in return_data: return_data[k] = np.zeros((nsteps, *it.shape)) return_data[k][step, ...] = it.copy() return return_data, configs
def sample_overlap(wfs, configs, pgrad, nblocks=10, nsteps_per_block=10, nsteps=None, tstep=0.5): r""" Sample .. math:: \rho(R) = \sum_i |\Psi_i(R)|^2 `pgrad` is expected to be a gradient generator. returns data as follows: `overlap` : .. math:: \left\langle \frac{\Psi_i^* \Psi_j}{\rho} \right\rangle `overlap_gradient`: .. math:: \left\langle \frac{\partial_{im} \Psi_i^* \Psi_j}{\rho} \right\rangle Note that overlap_gradient is only saved for i = f, where f is the final wave function. In addition, any key returned by `pgrad` will be saved for the final wave function. """ nconf, nelec, _ = configs.configs.shape if nsteps is not None: nblocks = nsteps nsteps_per_block = 1 for wf in wfs: wf.recompute(configs) return_data = {} for block in range(nblocks): block_avg = {} for step in range(nsteps_per_block): for e in range(nelec): # Propose move grads = [ np.real(wf.gradient(e, configs.electron(e)).T) for wf in wfs ] grad = limdrift(np.mean(grads, axis=0)) gauss = np.random.normal(scale=np.sqrt(tstep), size=(nconf, 3)) newcoorde = configs.configs[:, e, :] + gauss + grad * tstep newcoorde = configs.make_irreducible(e, newcoorde) # Compute reverse move grads = [np.real(wf.gradient(e, newcoorde).T) for wf in wfs] new_grad = limdrift(np.mean(grads, axis=0)) forward = np.sum(gauss**2, axis=1) backward = np.sum((gauss + tstep * (grad + new_grad))**2, axis=1) # Acceptance t_prob = np.exp(1 / (2 * tstep) * (forward - backward)) wf_ratios = np.array( [wf.testvalue(e, newcoorde)**2 for wf in wfs]) log_values = np.array([wf.value()[1] for wf in wfs]) ref = log_values[0] weights = np.exp(2 * (log_values - ref)) ratio = (t_prob * np.sum(wf_ratios * weights, axis=0) / np.sum(weights, axis=0)) accept = ratio > np.random.rand(nconf) # Update wave function configs.move(e, newcoorde, accept) for wf in wfs: wf.updateinternals(e, newcoorde, mask=accept) # print("accept", np.mean(accept)) log_values = np.array([wf.value() for wf in wfs]) # print(log_values.shape) ref = np.max(log_values[:, 1, :], axis=0) save_dat = {} denominator = np.sum(np.exp(2 * (log_values[:, 1, :] - ref)), axis=0) normalized_values = log_values[:, 0, :] * np.exp(log_values[:, 1, :] - ref) save_dat["overlap"] = np.mean( np.einsum("ik,jk->ijk", normalized_values, normalized_values) / denominator, axis=-1, ) weight = np.array([ np.exp(-2 * (log_values[i, 1, :] - log_values[:, 1, :])) for i in range(len(wfs)) ]) weight = 1.0 / np.sum(weight, axis=1) # Fast evaluation of dppsi_reg dppsi = pgrad.transform.serialize_gradients(wfs[-1].pgradient()) node_cut, f = pgrad._node_regr(configs, wfs[-1]) dppsi_regularized = dppsi * f[:, np.newaxis] save_dat["overlap_gradient"] = np.mean( np.einsum( "km,k,jk->jmk", dppsi_regularized, normalized_values[-1], normalized_values, ) / denominator, axis=-1, ) # Weighted average on rest dat = pgrad.avg(configs, wf, weight[-1]) for k in dat.keys(): save_dat[k] = dat[k] save_dat["weight"] = np.mean(weight, axis=1) # Rolling average within block for k, it in save_dat.items(): if k not in block_avg: block_avg[k] = np.zeros((*it.shape, )) block_avg[k] += save_dat[k] / nsteps_per_block # Blocks stored for k, it in block_avg.items(): if k not in return_data: return_data[k] = np.zeros((nblocks, *it.shape)) return_data[k][block, ...] = it.copy() return return_data, configs
def sample_overlap_worker(wfs, configs, energy, transforms, nsteps=10, nblocks=10, tstep=0.5): r"""Run nstep Metropolis steps to sample a distribution proportional to :math:`\sum_i |\Psi_i|^2`, where :math:`\Psi_i` = wfs[i] """ nconf, nelec, _ = configs.configs.shape for wf in wfs: wf.recompute(configs) weighted = [] unweighted = [] for block in range(nblocks): weighted_block = {} unweighted_block = {} for n in range(nsteps): for e in range(nelec): # a sweep # Propose move grads = [ np.real(wf.gradient(e, configs.electron(e)).T) for wf in wfs ] grad = mc.limdrift(np.mean(grads, axis=0)) gauss = np.random.normal(scale=np.sqrt(tstep), size=(nconf, 3)) newcoorde = configs.configs[:, e, :] + gauss + grad * tstep newcoorde = configs.make_irreducible(e, newcoorde) # Compute reverse move grads, vals = list( zip(*[wf.gradient_value(e, newcoorde) for wf in wfs])) grads = [np.real(g.T) for g in grads] new_grad = mc.limdrift(np.mean(grads, axis=0)) forward = np.sum(gauss**2, axis=1) backward = np.sum((gauss + tstep * (grad + new_grad))**2, axis=1) # Acceptance t_prob = np.exp(1 / (2 * tstep) * (forward - backward)) wf_ratios = np.abs(vals)**2 log_values = np.real(np.array([wf.value()[1] for wf in wfs])) weights = np.exp(2 * (log_values - log_values[0])) ratio = t_prob * np.sum(wf_ratios * weights, axis=0) / weights.sum(axis=0) accept = ratio > np.random.rand(nconf) #block_avg["acceptance"][n] += accept.mean() / nelec # Update wave function configs.move(e, newcoorde, accept) for wf in wfs: wf.updateinternals(e, newcoorde, configs, mask=accept) # Collect rolling average weighted_dat, unweighted_dat = collect_overlap_data( wfs, configs, energy, transforms) for k, it in unweighted_dat.items(): if k not in unweighted_block: unweighted_block[k] = np.zeros((*it.shape, ), dtype=it.dtype) unweighted_block[k] += unweighted_dat[k] / nsteps for k, it in weighted_dat.items(): if k not in weighted_block: weighted_block[k] = [ np.zeros((*x.shape, ), dtype=x.dtype) for x in it ] for b, v in zip(weighted_block[k], it): b += v / nsteps weighted.append(weighted_block) unweighted.append(unweighted_block) # here we modify the data so that it's a dictionary of lists of arrays for weighted # and a dictionary of arrays for unweighted # Access weighted as weighted[quantity][block, ...] # Access unweighted as unweighted[quantity][block,...] weighted = invert_list_of_dicts(weighted) unweighted = invert_list_of_dicts(unweighted) for k in weighted.keys(): weighted[k] = np.asarray(weighted[k]) for k in unweighted.keys(): unweighted[k] = np.asarray(unweighted[k]) return weighted, unweighted, configs