From 23258a4fb90d7728e579acfbaf08e3988fa7d79d Mon Sep 17 00:00:00 2001 From: FHof Date: Wed, 5 Jan 2022 14:40:16 +0100 Subject: [PATCH] VEGAS+: Replace torch. with autoray.numpy. operations After this commit it still only works with torch and should indirectly perform the same torch operations as before. --- torchquad/integration/vegas.py | 73 +++++++++------- torchquad/integration/vegas_map.py | 86 +++++++++++-------- torchquad/integration/vegas_stratification.py | 73 +++++++++------- 3 files changed, 133 insertions(+), 99 deletions(-) diff --git a/torchquad/integration/vegas.py b/torchquad/integration/vegas.py index 268cb1a2..c3686753 100644 --- a/torchquad/integration/vegas.py +++ b/torchquad/integration/vegas.py @@ -1,15 +1,16 @@ -import torch +from autoray import numpy as anp +from autoray import infer_backend from loguru import logger from .base_integrator import BaseIntegrator -from .utils import _setup_integration_domain +from .utils import _setup_integration_domain, RNG from .vegas_map import VEGASMap from .vegas_stratification import VEGASStratification class VEGAS(BaseIntegrator): - """VEGAS Enhanced in torch. Refer to https://arxiv.org/abs/2009.05112 . + """VEGAS Enhanced. Refer to https://arxiv.org/abs/2009.05112 . Implementation inspired by https://github.com/ycwu1030/CIGAR/ . EQ refers to equation in the above paper. """ @@ -29,6 +30,7 @@ def integrate( eps_abs=0, max_iterations=20, use_warmup=True, + backend="torch", ): """Integrates the passed function on the passed domain using VEGAS. @@ -43,6 +45,7 @@ def integrate( eps_abs (float, optional): Absolute error to abort at. Defaults to 0. max_iterations (int, optional): Maximum number of vegas iterations to perform. Defaults to 32. use_warmup (bool, optional): If a warmup should be used to initialize the map. Defaults to True. + backend (string, optional): Numerical backend. This argument is ignored if the backend can be inferred from integration_domain. Defaults to "torch". Raises: ValueError: If len(integration_domain) != dim @@ -71,11 +74,12 @@ def integrate( self._starting_N = N // self._max_iterations self._N_increment = N // self._max_iterations self._fn = fn - self._integration_domain = _setup_integration_domain( - dim, integration_domain, "torch" + self._integration_domain = integration_domain = _setup_integration_domain( + dim, integration_domain, backend ) - if seed is not None: - torch.random.manual_seed(seed) + self.backend = infer_backend(integration_domain) + self.dtype = integration_domain.dtype + self.rng = RNG(backend=self.backend, seed=seed) # Initialize the adaptive VEGAS map, # Note that a larger number of intervals may lead to problems if only few evals are allowed @@ -87,7 +91,9 @@ def integrate( # Initialize VEGAS' stratification # Paper section III - self.strat = VEGASStratification(self._N_increment, dim=self._dim) + self.strat = VEGASStratification( + self._N_increment, dim=self._dim, rng=self.rng, backend=self.backend + ) logger.debug("Starting VEGAS") @@ -125,8 +131,8 @@ def integrate( chi2 = self._get_chisq() acc = err / res - if torch.isnan(acc): # capture 0 error - acc = torch.tensor(0.0) + if anp.isnan(acc): # capture 0 error + acc = anp.array(0.0, like=acc) # Abort if errors acceptable logger.debug(f"Iteration {self.it},Chi2={chi2:.4e}") @@ -136,9 +142,12 @@ def integrate( # Adjust number of evals if Chi square indicates instability # EQ 32 if chi2 / 5.0 < 1.0: - self._starting_N = torch.minimum( - torch.tensor(self._starting_N + self._N_increment), - self._starting_N * torch.sqrt(acc / (eps_rel + 1e-8)), + self._starting_N = anp.minimum( + anp.array( + self._starting_N + self._N_increment, + like=acc, + ), + self._starting_N * anp.sqrt(acc / (eps_rel + 1e-8)), ) self.results = [] # reset sample results self.sigma2 = [] # reset sample results @@ -165,8 +174,6 @@ def _warmup_grid(self, warmup_N_it=5, N_samples=1000): f"Running Map Warmup with warmup_N_it={warmup_N_it}, N_samples={N_samples}..." ) - yrnd = torch.zeros(self._dim) # sample points - x = torch.zeros(self._dim) # transformed sample points alpha_start = 0.5 # initial alpha value # TODO in the original paper this is adjusted over time self.alpha = alpha_start @@ -181,8 +188,12 @@ def _warmup_grid(self, warmup_N_it=5, N_samples=1000): jf = 0 # jacobians * function jf2 = 0 + # Sample points yrnd and transformed sample points x # Multiplying by 0.99999999 as the edge case of y=1 leads to an error - yrnd = torch.rand(size=[N_samples, self._dim]) * 0.999999 + yrnd = ( + self.rng.uniform(size=[N_samples, self._dim], dtype=self.dtype) + * 0.999999 + ) x = self.map.get_X(yrnd) f_eval = self._eval(x).squeeze() jac = self.map.get_Jac(yrnd) @@ -198,7 +209,7 @@ def _warmup_grid(self, warmup_N_it=5, N_samples=1000): self.sigma2[-1] += sig2 / N_samples # store results self.map.update_map() # adapt the map # TODO fix for integrals close to 0 - acc = torch.sqrt( + acc = anp.sqrt( self.sigma2[-1] / self.results[-1] ) # compute estimated accuracy, logger.debug( @@ -212,15 +223,13 @@ def _run_iteration(self): """Runs one iteration of VEGAS including stratification and updates the VEGAS map if use_grid_improve is set. Returns: - float: Estimated accuracy. + backend tensor float: Estimated accuracy. """ - y = torch.zeros(self._dim) # stratified sampling points - x = torch.zeros(self._dim) # transformed sample points - neval = self.strat.get_NH(self._starting_N) # Evals per strat cube - self.starting_N = torch.sum(neval) / self.strat.N_cubes # update real neval + self.starting_N = anp.sum(neval) / self.strat.N_cubes # update real neval self._nr_of_fevals += neval.sum() # Locally track function evals + # Stratified sampling points y and transformed sample points x y = self.strat.get_Y(neval) x = self.map.get_X(y) # transform, EQ 8+9 f_eval = self._eval(x).squeeze() # eval integrand @@ -233,19 +242,19 @@ def _run_iteration(self): self.map.accumulate_weight(y, jf_vec2) # EQ 25 jf, jf2 = self.strat.accumulate_weight(neval, jf_vec) # update strat - ih = torch.divide(jf, neval) * self.strat.V_cubes # Compute integral per cube + ih = (jf / neval) * self.strat.V_cubes # Compute integral per cube # Collect results - sig2 = torch.divide(jf2, neval) * (self.strat.V_cubes ** 2) - pow(ih, 2) + sig2 = (jf2 / neval) * (self.strat.V_cubes ** 2) - pow(ih, 2) self.results[-1] = ih.sum() # store results - self.sigma2[-1] = torch.divide(sig2, neval).sum() + self.sigma2[-1] = (sig2 / neval).sum() if self.use_grid_improve: # if on, update adaptive map logger.debug("Running grid improvement") self.map.update_map() self.strat.update_DH() # update stratification - acc = torch.sqrt(self.sigma2[-1] / (self.results[-1])) # estimate accuracy + acc = anp.sqrt(self.sigma2[-1] / (self.results[-1])) # estimate accuracy return acc @@ -254,7 +263,7 @@ def _get_result(self): """Computes mean of results to estimate integral, EQ 30. Returns: - float: Estimated integral. + backend tensor float: Estimated integral. """ res_num = 0 res_den = 0 @@ -262,8 +271,8 @@ def _get_result(self): res_num += res / self.sigma2[idx] res_den += 1.0 / self.sigma2[idx] - if torch.isnan(res_num / res_den): # if variance is 0 just return mean result - return torch.mean(torch.tensor(self.results)) + if anp.isnan(res_num / res_den): # if variance is 0 just return mean result + return anp.mean(anp.array(self.results, like=res_num)) else: return res_num / res_den @@ -271,20 +280,20 @@ def _get_error(self): """Estimates error from variance , EQ 31. Returns: - float: Estimated error. + backend tensor float: Estimated error. """ res = 0 for sig in self.sigma2: res += 1.0 / sig - return 1.0 / torch.sqrt(res) + return 1.0 / anp.sqrt(res) def _get_chisq(self): """Computes chi square from estimated integral and variance, EQ 32. Returns: - float: Chi squared. + backend tensor float: Chi squared. """ I_final = self._get_result() chi2 = 0 diff --git a/torchquad/integration/vegas_map.py b/torchquad/integration/vegas_map.py index e59d207a..a3f08cb9 100644 --- a/torchquad/integration/vegas_map.py +++ b/torchquad/integration/vegas_map.py @@ -1,4 +1,5 @@ -import torch +from autoray import numpy as anp +from autoray import infer_backend class VEGASMap: @@ -12,7 +13,7 @@ def __init__(self, dim, integration_domain, N_intervals=100, alpha=0.5) -> None: Args: dim (int): Dimensionality of the integrand. - integration_domain (list, optional): Integration domain, e.g. [[-1,1],[0,1]]. Defaults to [-1,1]^dim. + integration_domain (backend tensor): Integration domain. N_intervals (int, optional): Number of intervals to split the domain in. Defaults to 100. alpha (float, optional): Alpha from the paper, EQ 19. Defaults to 0.5. """ @@ -20,8 +21,11 @@ def __init__(self, dim, integration_domain, N_intervals=100, alpha=0.5) -> None: self.N_intervals = N_intervals # # of subdivisions self.N_edges = self.N_intervals + 1 # # of subdivsion boundaries self.alpha = alpha # Weight smoothing - self.x_edges = torch.zeros((self.dim, self.N_edges)) # boundary locations - self.dx_edges = torch.zeros((self.dim, self.N_intervals)) # subdomain stepsizes + self.backend = infer_backend(integration_domain) + + # boundary locations and subdomain stepsizes + self.x_edges = anp.zeros((self.dim, self.N_edges), like=self.backend) + self.dx_edges = anp.zeros((self.dim, self.N_intervals), like=self.backend) # Subdivide initial domain equally spaced in N-d, EQ 8 for dim in range(self.dim): @@ -35,28 +39,35 @@ def __init__(self, dim, integration_domain, N_intervals=100, alpha=0.5) -> None: self.dx_edges[dim][i - 1] = stepsize # weights in each intervall - self.weights = torch.zeros((self.dim, self.N_intervals)) - self.smoothed_weights = torch.zeros((self.dim, self.N_intervals)) - self.summed_weights = torch.zeros(self.dim) # sum of weights per dim - self.delta_weights = torch.zeros(self.dim) # EQ 11 - self.std_weight = torch.zeros(self.dim) # EQ 13 - self.avg_weight = torch.zeros(self.dim) # EQ 14 + self.weights = anp.zeros((self.dim, self.N_intervals), like=self.backend) + self.smoothed_weights = anp.zeros( + (self.dim, self.N_intervals), like=self.backend + ) + self.summed_weights = anp.zeros( + [self.dim], like=self.backend + ) # sum of weights per dim + self.delta_weights = anp.zeros([self.dim], like=self.backend) # EQ 11 + self.std_weight = anp.zeros([self.dim], like=self.backend) # EQ 13 + self.avg_weight = anp.zeros([self.dim], like=self.backend) # EQ 14 # numbers of random samples in specific interval - self.counts = torch.zeros((self.dim, self.N_intervals)).long() + self.counts = anp.zeros( + (self.dim, self.N_intervals), + like=self.backend, + ).long() def get_X(self, y): """Get mapped sampling points, EQ 9. Args: - y (torch.tensor): Randomly sampled location(s) + y (backend tensor): Randomly sampled location(s) Returns: - torch.tensor: Mapped points. + backend tensor: Mapped points. """ ID, offset = self._get_interval_ID(y), self._get_interval_offset(y) - res = torch.zeros_like(y) + res = anp.zeros_like(y) for i in range(self.dim): - ID_i = torch.floor(ID[:, i]).long() + ID_i = anp.floor(ID[:, i]).long() res[:, i] = self.x_edges[i, ID_i] + self.dx_edges[i, ID_i] * offset[:, i] return res @@ -67,12 +78,12 @@ def get_Jac(self, y): y ([type]): Sampled locations. Returns: - torch.tensor: Jacobian + backend tensor: Jacobian """ ID = self._get_interval_ID(y) - jac = torch.ones(y.shape[0]) + jac = anp.ones([y.shape[0]], like=y) for i in range(self.dim): - ID_i = torch.floor(ID[:, i]).long() + ID_i = anp.floor(ID[:, i]).long() jac *= self.N_intervals * self.dx_edges[i][ID_i] return jac @@ -80,21 +91,21 @@ def _get_interval_ID(self, y): """Get the integer part of the desired mapping , EQ 10. Args: - y (float): Sampled point + y (backend tensor): Sampled points Returns: - int: Integer part of mapped point. + backend tensor: Integer part of mapped points. """ - return torch.floor(y * float(self.N_intervals)) + return anp.floor(y * float(self.N_intervals)) def _get_interval_offset(self, y): """Get the fractional part of the desired mapping , EQ 11. Args: - y (float): Sampled point. + y (backend tensor): Sampled points. Returns: - float: Fractional part of mapped point. + backend tensor: Fractional part of mapped points. """ return (y * self.N_intervals) - self._get_interval_ID(y) @@ -102,13 +113,13 @@ def accumulate_weight(self, y, jf_vec2): """Accumulate weights and counts of the map. Args: - y (float): Sampled point. - jf_vec2 (float): Square of the product of function value and jacobian + y (backend tensor): Sampled points. + jf_vec2 (backend tensor): Square of the product of function values and jacobians """ ID = self._get_interval_ID(y) for i in range(self.dim): - ID_i = torch.floor(ID[:, i]).long() - unique_vals, unique_counts = torch.unique(ID_i, return_counts=True) + ID_i = anp.floor(ID[:, i]).long() + unique_vals, unique_counts = anp.unique(ID_i, return_counts=True) weights_vals = jf_vec2 for val in unique_vals: self.weights[i][val] += weights_vals[ID_i == val].sum() @@ -130,9 +141,7 @@ def _smooth_map(self): # i == 0 d_tmp = (7.0 * self.weights[dim][0] + self.weights[dim][1]) / (8.0 * d_sum) - d_tmp = ( - pow((d_tmp - 1.0) / torch.log(d_tmp), self.alpha) if d_tmp != 0 else 0 - ) + d_tmp = pow((d_tmp - 1.0) / anp.log(d_tmp), self.alpha) if d_tmp != 0 else 0 self.smoothed_weights[dim][0] = d_tmp # i == last @@ -140,9 +149,7 @@ def _smooth_map(self): self.weights[dim][self.N_intervals - 2] + 7.0 * self.weights[dim][self.N_intervals - 1] ) / (8.0 * d_sum) - d_tmp = ( - pow((d_tmp - 1.0) / torch.log(d_tmp), self.alpha) if d_tmp != 0 else 0 - ) + d_tmp = pow((d_tmp - 1.0) / anp.log(d_tmp), self.alpha) if d_tmp != 0 else 0 self.smoothed_weights[dim][-1] = d_tmp # rest @@ -152,7 +159,7 @@ def _smooth_map(self): + self.weights[dim][2:] ) / (8.0 * d_sum) d_tmp[d_tmp != 0] = pow( - (d_tmp[d_tmp != 0] - 1.0) / torch.log(d_tmp[d_tmp != 0]), self.alpha + (d_tmp[d_tmp != 0] - 1.0) / anp.log(d_tmp[d_tmp != 0]), self.alpha ) self.smoothed_weights[dim][1:-1] = d_tmp @@ -166,8 +173,11 @@ def _reset_weight( self, ): """Resets weights.""" - self.weights = torch.zeros((self.dim, self.N_intervals)) - self.counts = torch.zeros((self.dim, self.N_intervals)).long() + self.weights = anp.zeros((self.dim, self.N_intervals), like=self.backend) + self.counts = anp.zeros( + (self.dim, self.N_intervals), + like=self.backend, + ).long() def update_map( self, @@ -179,8 +189,8 @@ def update_map( old_i = 0 d_accu = 0 - indices = torch.zeros(self.N_intervals - 1).long() - d_accu_i = torch.zeros(self.N_intervals - 1) + indices = anp.zeros([self.N_intervals - 1], like=self.backend).long() + d_accu_i = anp.zeros([self.N_intervals - 1], like=self.backend) for new_i in range(1, self.N_intervals): d_accu += self.delta_weights[i] diff --git a/torchquad/integration/vegas_stratification.py b/torchquad/integration/vegas_stratification.py index a4549ebe..4445dc52 100644 --- a/torchquad/integration/vegas_stratification.py +++ b/torchquad/integration/vegas_stratification.py @@ -1,5 +1,4 @@ -import torch -import math +from autoray import numpy as anp class VEGASStratification: @@ -8,41 +7,53 @@ class VEGASStratification: EQ refers to equation in the above paper. """ - def __init__(self, N_increment, dim=1, beta=0.75): + def __init__(self, N_increment, dim, rng, beta=0.75, backend="torch"): """Initialize the VEGAS stratification. Args: - N_increment (int, optional): Number of evaluations per iteration. - dim (int, optional): Dimensionality. Defaults to 1. + N_increment (int): Number of evaluations per iteration. + dim (int): Dimensionality + rng (RNG): Random number generator beta (float, optional): Beta parameter from VEGAS Enhanced. Defaults to 0.75. + backend (string, optional): Numerical backend. Defaults to "torch" """ + self.rng = rng self.dim = dim # stratification steps per dim, EQ 33 - self.N_strat = math.floor((N_increment / 2.0) ** (1.0 / dim)) + self.N_strat = int((N_increment / 2.0) ** (1.0 / dim)) self.N_strat = 1000 if self.N_strat > 1000 else self.N_strat self.beta = beta # variable controlling adaptiveness in stratification 0 to 1 self.N_cubes = self.N_strat ** self.dim # total number of subdomains self.V_cubes = (1.0 / self.N_strat) ** self.dim # volume of hypercubes - self.JF = torch.zeros(self.N_cubes) # jacobian times f eval - self.JF2 = torch.zeros(self.N_cubes) # jacobian^2 times f - self.dh = torch.ones(self.N_cubes) * 1.0 / self.N_cubes # dampened counts - self.strat_counts = torch.zeros(self.N_cubes) # current index counts + + # Use the backend-default dtype; doesn't work with numpy or tensorflow + self.dtype = anp.zeros([1], like=backend).dtype + + # jacobian times f eval and jacobian^2 times f + self.JF = anp.zeros([self.N_cubes], like=backend) + self.JF2 = anp.zeros([self.N_cubes], like=backend) + + # dampened counts + self.dh = anp.ones([self.N_cubes], like=backend) * 1.0 / self.N_cubes + + # current index counts + self.strat_counts = anp.zeros([self.N_cubes], like=backend) def accumulate_weight(self, nevals, weight_all_cubes): """Accumulate weights for the cubes. Args: - nevals (torch.tensor): Number of evals belonging to each cube (sorted). - weight_all_cubes (torch.tensor): Function values. + nevals (backend tensor): Number of evals belonging to each cube (sorted). + weight_all_cubes (backend tensor): Function values. Returns: - torch.tensor,torch.tensor: Computed JF and JF2 + backend tensor, backend tensor: Computed JF and JF2 """ # Build indices for weights - indices = torch.cumsum(nevals, dim=0) + indices = anp.cumsum(nevals, axis=0) # Compute squares of all weights ahead - square_weights = torch.pow(weight_all_cubes, 2.0) + square_weights = weight_all_cubes ** 2.0 # Used to store previous indexstarting point prev = 0 @@ -76,7 +87,7 @@ def update_DH(self): self.dh = d_tmp ** self.beta # for very small d_tmp d_tmp ** self.beta becomes NaN - self.dh[torch.isnan(self.dh)] = 0 + self.dh[anp.isnan(self.dh)] = 0 # Normalize dampening d_sum = sum(self.dh) @@ -91,26 +102,25 @@ def get_NH(self, nevals_exp): nevals_exp (int): Expected number of evaluations. Returns: - torch.tensor: Stratified sample counts per cube. + backend tensor: Stratified sample counts per cube. """ - nh = torch.multiply(self.dh, nevals_exp) - nh = torch.floor(nh) - nh = torch.clip(nh, 2, None).int() + nh = anp.floor(self.dh * nevals_exp) + nh = anp.clip(nh, 2, None).int() return nh def _get_indices(self, idx): """Maps point to stratified point. Args: - idx (int): Target points index. + idx (backend tensor): Target points indices. Returns: - torch.tensor: Mapped point. + backend tensor: Mapped points. """ - res = torch.zeros([len(idx), self.dim]) + res = anp.zeros([len(idx), self.dim], like=idx) tmp = idx for i in range(self.dim): - q = torch.div(tmp, self.N_strat, rounding_mode="floor") + q = anp.div(tmp, self.N_strat, rounding_mode="floor") r = tmp - q * self.N_strat res[:, i] = r tmp = q @@ -120,22 +130,27 @@ def get_Y(self, nevals): """Compute randomly sampled points. Args: - nevals (torch.tensor): Number of samples to draw per stratification cube. + nevals (backend tensor): Number of samples to draw per stratification cube. Returns: - torch.tensor: Sampled points. + backend tensor: Sampled points. """ dy = 1.0 / self.N_strat res_in_all_cubes = [] # Get indices - indices = torch.arange(len(nevals)) + indices = anp.arange(len(nevals), like=nevals) indices = self._get_indices(indices) # Get random numbers (we get a few more just to vectorize properly) # This might increase the memory requirements slightly but is probably # worth it. - random_uni = torch.rand(size=[len(nevals), nevals.max(), self.dim]) * 0.999999 + random_uni = ( + self.rng.uniform( + size=[len(nevals), nevals.max(), self.dim], dtype=self.dtype + ) + * 0.999999 + ) # Sum the random numbers onto the index locations and scale with dy # Note that the resulting tensor is still slightly too large @@ -147,4 +162,4 @@ def get_Y(self, nevals): for idx, N in enumerate(nevals): res_in_all_cubes.append(res[idx, 0:N, :]) - return torch.cat(res_in_all_cubes, dim=0) + return anp.concatenate(res_in_all_cubes, axis=0, like=res)