From 123ea666af1fe7431b28c7030b2e4785e6ccd6d2 Mon Sep 17 00:00:00 2001 From: Tao Wang Date: Fri, 4 Oct 2024 00:31:56 -0400 Subject: [PATCH 01/10] api refactoring wip --- src/integrators.py | 42 +++++------ src/maps.py | 181 ++++++++++++++++++--------------------------- 2 files changed, 95 insertions(+), 128 deletions(-) diff --git a/src/integrators.py b/src/integrators.py index 7084970..522def1 100644 --- a/src/integrators.py +++ b/src/integrators.py @@ -13,24 +13,24 @@ class Integrator: def __init__( self, # bounds: Union[List[Tuple[float, float]], np.ndarray], - map, + map=None, neval: int = 1000, - batch_size: int = None, + nbatch: int = None, device="cpu", # device=torch.device("cuda" if torch.cuda.is_available() else "cpu"), ): - if not isinstance(map, Map): - map = Affine(map) + #if not isinstance(map, Map): + # map = Affine(map) self.dim = map.dim self.map = map self.neval = neval - if batch_size is None: - self.batch_size = neval + if nbatch is None: + self.nbatch = neval self.neval = neval else: - self.batch_size = batch_size - self.neval = -(-neval // batch_size) * batch_size + self.nbatch = nbatch + self.neval = -(-neval // nbatch) * nbatch self.device = device @@ -44,18 +44,18 @@ def __init__( map, nitn: int = 10, neval: int = 1000, - batch_size: int = None, + nbatch: int = None, device="cpu", adapt=False, alpha=0.5, ): - super().__init__(map, neval, batch_size, device) + super().__init__(map, neval, nbatch, device) self.adapt = adapt self.alpha = alpha self.nitn = nitn def __call__(self, f: Callable, **kwargs): - u = torch.rand(self.batch_size, self.dim, device=self.device) + u = torch.rand(self.nbatch, self.dim, device=self.device) x, _ = self.map.forward(u) f_values = f(x) f_size = len(f_values) if isinstance(f_values, (list, tuple)) else 1 @@ -66,14 +66,14 @@ def __call__(self, f: Callable, **kwargs): # var = torch.zeros((f_size, f_size), dtype=type_fval, device=self.device) result = RAvg(weighted=self.adapt) - epoch = self.neval // self.batch_size + epoch = self.neval // self.nbatch for itn in range(self.nitn): mean[:] = 0 var[:] = 0 for _ in range(epoch): y = torch.rand( - self.batch_size, self.dim, dtype=torch.float64, device=self.device + self.nbatch, self.dim, dtype=torch.float64, device=self.device ) x, jac = self.map.forward(y) @@ -108,25 +108,25 @@ def __init__( map: Map, nitn: int = 10, neval=10000, - batch_size=None, + nbatch=None, n_burnin=500, device="cpu", adapt=False, alpha=0.5, ): - super().__init__(map, nitn, neval, batch_size, device, adapt, alpha) + super().__init__(map, nitn, neval, nbatch, device, adapt, alpha) self.n_burnin = n_burnin def __call__( self, f: Callable, - proposal_dist="global_uniform", + proposal_dist="uniform", thinning=1, mix_rate=0.0, **kwargs, ): epsilon = 1e-16 # Small value to ensure numerical stability - vars_shape = (self.batch_size, self.dim) + vars_shape = (self.nbatch, self.dim) current_y = torch.rand(vars_shape, dtype=torch.float64, device=self.device) current_x, current_jac = self.map.forward(current_y) current_fval = f(current_x) @@ -149,7 +149,7 @@ def __call__( result = RAvg(weighted=self.adapt) result_ref = RAvg(weighted=self.adapt) - epoch = self.neval // self.batch_size + epoch = self.neval // self.nbatch n_meas = 0 for itn in range(self.nitn): for i in range(epoch): @@ -162,7 +162,7 @@ def __call__( acceptance_probs = new_weight / current_weight * new_jac / current_jac accept = ( - torch.rand(self.batch_size, dtype=torch.float64, device=self.device) + torch.rand(self.nbatch, dtype=torch.float64, device=self.device) <= acceptance_probs ) @@ -190,7 +190,7 @@ def __call__( ) result.sum_neval += self.neval result.add(gvar.gvar(mean.item(), ((var / n_meas) ** 0.5).item())) - result_ref.sum_neval += self.batch_size + result_ref.sum_neval += self.nbatch result_ref.add( gvar.gvar(mean_ref.item(), ((var_ref / n_meas) ** 0.5).item()) ) @@ -204,7 +204,7 @@ def _propose(self, u, proposal_dist, **kwargs): if proposal_dist == "random_walk": step_size = kwargs.get("step_size", 0.2) return (u + (torch.rand_like(u) - 0.5) * step_size) % 1.0 - elif proposal_dist == "global_uniform": + elif proposal_dist == "uniform": return torch.rand_like(u) # elif proposal_dist == "gaussian": # mean = kwargs.get("mean", torch.zeros_like(u)) diff --git a/src/maps.py b/src/maps.py index 9ab2881..10e9769 100644 --- a/src/maps.py +++ b/src/maps.py @@ -1,63 +1,47 @@ import torch import numpy as np +from torch import nn - -class Map: - def __init__(self, map_spec, device="cpu"): - # if isinstance(map_spec, dict): - # self.map_spec = { - # k: torch.tensor(v, device=device) for k, v in map_spec.items() - # } - if isinstance(map_spec, (list, np.ndarray)): - self.map_spec = torch.tensor(map_spec, dtype=torch.float64, device=device) +class Map(nn.Module): + def __init__(self, bounds, device="cpu"): + super().__init__() + if isinstance(bounds, (list, np.ndarray)): + self.bounds = torch.tensor(bounds, dtype=torch.float64, device=device) else: raise ValueError("Unsupported map specification") - self.dim = self.map_spec.shape[0] + self.dim = self.bounds.shape[0] self.device = device - def forward(self, y): + def forward(self, u): raise NotImplementedError("Subclasses must implement this method") - + def inverse(self, x): raise NotImplementedError("Subclasses must implement this method") - - def log_det_jacobian(self, y): + + def sample(self, nsample): raise NotImplementedError("Subclasses must implement this method") + class Affine(Map): - def __init__(self, map_spec, device="cpu"): - super().__init__(map_spec, device) - self._A = self.map_spec[:, 1] - self.map_spec[:, 0] + def __init__(self, bounds, device="cpu"): + super().__init__(bounds, device) + self._A = self.bounds[:, 1] - self.bounds[:, 0] self._jac1 = torch.prod(self._A) - def forward(self, y): - return y * self._A + self.map_spec[:, 0], self._jac1.repeat(y.shape[0]) + def forward(self, u): + return u * self._A + self.bounds[:, 0], torch.log(self._jac1.repeat(u.shape[0])) def inverse(self, x): - return (x - self.map_spec[:, 0]) / self._A, self._jac1.repeat(x.shape[0]) + return (x - self.bounds[:, 0]) / self._A, torch.log(self._jac1.repeat(x.shape[0])) - def log_det_jacobian(self, y): - return torch.log(self._jac1) * y.shape[0] -class AdaptiveMap(Map): - def __init__(self, map_spec, alpha=0.5, device="cpu"): - super().__init__(map_spec, device) - self.alpha = alpha - - def add_training_data(self, y, f): - pass - - def adapt(self, alpha=0.0): - pass - - -class Vegas(AdaptiveMap): - def __init__(self, map_spec, ninc=1000, alpha=0.5, device="cpu"): - super().__init__(map_spec, alpha, device) +class Vegas(Map): + def __init__(self, bounds, ninc=1000, alpha=0.5, device="cpu"): + super().__init__(bounds, device) # self.nbin = nbin - + self.alpha = alpha if isinstance(ninc, int): self.ninc = torch.ones(self.dim, dtype=torch.int32, device=device) * ninc else: @@ -72,8 +56,8 @@ def __init__(self, map_spec, ninc=1000, alpha=0.5, device="cpu"): for d in range(self.dim): self.grid[d, : self.ninc[d] + 1] = torch.linspace( - self.map_spec[d, 0], - self.map_spec[d, 1], + self.bounds[d, 0], + self.bounds[d, 1], self.ninc[d] + 1, dtype=torch.float64, device=self.device, @@ -83,8 +67,8 @@ def __init__(self, map_spec, ninc=1000, alpha=0.5, device="cpu"): ) self.clear() - def add_training_data(self, y, f): - """Add training data ``f`` for ``y``-space points ``y``. + def add_training_data(self, u, f): + """Add training data ``f`` for ``u``-space points ``u``. Accumulates training data for later use by ``self.adapt()``. Grid increments will be made smaller in regions where @@ -93,19 +77,19 @@ def add_training_data(self, y, f): when ``f`` is constant across the grid. Args: - y (tensor): ``y`` values corresponding to the training data. - ``y`` is a contiguous 2-d tensor, where ``y[j, d]`` + u (tensor): ``u`` values corresponding to the training data. + ``u`` is a contiguous 2-d tensor, where ``u[j, d]`` is for points along direction ``d``. f (tensor): Training function values. ``f[j]`` corresponds to - point ``y[j, d]`` in ``y``-space. + point ``u[j, d]`` in ``u``-space. """ if self.sum_f is None: self.sum_f = torch.zeros_like(self.inc) self.n_f = torch.zeros_like(self.inc) + 1e-10 - iy = torch.floor(y * self.ninc).long() + iu = torch.floor(u * self.ninc).long() for d in range(self.dim): - self.sum_f[d, iy[:, d]] += torch.abs(f) - self.n_f[d, iy[:, d]] += 1 + self.sum_f[d, iu[:, d]] += torch.abs(f) + self.n_f[d, iu[:, d]] += 1 def adapt(self, alpha=0.0): """Adapt grid to accumulated training data. @@ -218,101 +202,84 @@ def clear(self): self.n_f = None @torch.no_grad() - def forward(self, y): - y = y.to(self.device) - y_ninc = y * self.ninc - iy = torch.floor(y_ninc).long() - dy_ninc = y_ninc - iy - - x = torch.empty_like(y) - jac = torch.ones(y.shape[0], device=x.device) + def forward(self, u): + u = u.to(self.device) + u_ninc = u * self.ninc + iu = torch.floor(u_ninc).long() + du_ninc = u_ninc - iu + + x = torch.empty_like(u) + jac = torch.ones(u.shape[0], device=x.device) # self.jac.fill_(1.0) for d in range(self.dim): - # Handle the case where iy < ninc + # Handle the case where iu < ninc ninc = self.ninc[d] - mask = iy[:, d] < ninc + mask = iu[:, d] < ninc if mask.any(): x[mask, d] = ( - self.grid[d, iy[mask, d]] - + self.inc[d, iy[mask, d]] * dy_ninc[mask, d] + self.grid[d, iu[mask, d]] + + self.inc[d, iu[mask, d]] * du_ninc[mask, d] ) - jac[mask] *= self.inc[d, iy[mask, d]] * ninc + jac[mask] *= self.inc[d, iu[mask, d]] * ninc - # Handle the case where iy >= ninc + # Handle the case where iu >= ninc mask_inv = ~mask if mask_inv.any(): x[mask_inv, d] = self.grid[d, ninc] jac[mask_inv] *= self.inc[d, ninc - 1] * ninc - return x, jac + return x, torch.log(jac) @torch.no_grad() def inverse(self, x): # self.jac.fill_(1.0) x = x.to(self.device) - y = torch.empty_like(x) + u = torch.empty_like(x) jac = torch.ones(x.shape[0], device=x.device) for d in range(self.dim): ninc = self.ninc[d] - iy = torch.searchsorted(self.grid[d, :], x[:, d].contiguous(), right=True) + iu = torch.searchsorted(self.grid[d, :], x[:, d].contiguous(), right=True) - mask_valid = (iy > 0) & (iy <= ninc) - mask_lower = iy <= 0 - mask_upper = iy > ninc + mask_valid = (iu > 0) & (iu <= ninc) + mask_lower = iu <= 0 + mask_upper = iu > ninc - # Handle valid range (0 < iy <= ninc) + # Handle valid range (0 < iu <= ninc) if mask_valid.any(): - iyi_valid = iy[mask_valid] - 1 - y[mask_valid, d] = ( - iyi_valid - + (x[mask_valid, d] - self.grid[d, iyi_valid]) - / self.inc[d, iyi_valid] + iui_valid = iu[mask_valid] - 1 + u[mask_valid, d] = ( + iui_valid + + (x[mask_valid, d] - self.grid[d, iui_valid]) + / self.inc[d, iui_valid] ) / ninc - jac[mask_valid] *= self.inc[d, iyi_valid] * ninc + jac[mask_valid] *= self.inc[d, iui_valid] * ninc - # Handle lower bound (iy <= 0)\ + # Handle lower bound (iu <= 0)\ if mask_lower.any(): - y[mask_lower, d] = 0.0 + u[mask_lower, d] = 0.0 jac[mask_lower] *= self.inc[d, 0] * ninc - # Handle upper bound (iy > ninc) + # Handle upper bound (iu > ninc) if mask_upper.any(): - y[mask_upper, d] = 1.0 + u[mask_upper, d] = 1.0 jac[mask_upper] *= self.inc[d, ninc - 1] * ninc - return y, jac - - @torch.no_grad() - def log_det_jacobian(self, y): - y = y.to(self.device) - y_ninc = y * self.ninc - iy = torch.floor(y_ninc).long() - - jac = torch.ones(y.shape[0], device=x.device) - for d in range(self.dim): - # Handle the case where iy < ninc - mask = iy[:, d] < self.ninc - if mask.any(): - jac[mask] *= self.inc[d, iy[mask, d]] * self.ninc - - # Handle the case where iy >= ninc - mask_inv = ~mask - if mask_inv.any(): - jac[mask_inv] *= self.inc[d, self.ninc - 1] * self.ninc + return u, torch.log(jac) - return torch.sum(torch.log(jac), dim=-1) + def sample(self, nsample): + return super() -class NormalizingFlow(AdaptiveMap): - def __init__(self, map_spec, flow_model, alpha=0.5, device="cpu"): - super().__init__(map_spec, alpha, device) +class NormalizingFlow(Map): + def __init__(self, bounds, flow_model, device="cpu"): + super().__init__(bounds, device) self.flow_model = flow_model.to(device) def forward(self, u): - return self.flow_model.forward(u)[0] + return self.flow_model.forward(u) def inverse(self, x): - return self.flow_model.inverse(x)[0] - - def log_det_jacobian(self, u): - return self.flow_model.forward(u)[1] + return self.flow_model.inverse(x) + + def sample(self, nsample): + return self.flow_model.sample(nsample) From cdde5234c82c77da37aa53374fb6e32a8a1d5eae Mon Sep 17 00:00:00 2001 From: Tao Wang Date: Tue, 8 Oct 2024 20:44:10 -0400 Subject: [PATCH 02/10] wip --- src/integrators.py | 52 +++++++++++++++++++++------------------------- src/maps.py | 40 ++++++++++++++++++++++++++++------- 2 files changed, 56 insertions(+), 36 deletions(-) diff --git a/src/integrators.py b/src/integrators.py index 522def1..417e75f 100644 --- a/src/integrators.py +++ b/src/integrators.py @@ -1,7 +1,7 @@ from typing import Callable, Union, List, Tuple, Dict import torch from utils import RAvg -from maps import Map, Affine +from maps import Map, Affine, NormalizingFlow import gvar @@ -13,17 +13,19 @@ class Integrator: def __init__( self, # bounds: Union[List[Tuple[float, float]], np.ndarray], - map=None, + maps: NormalizingFlow, neval: int = 1000, nbatch: int = None, device="cpu", + dtype = torch.float32, # device=torch.device("cuda" if torch.cuda.is_available() else "cpu"), ): #if not isinstance(map, Map): # map = Affine(map) - self.dim = map.dim - self.map = map + self.dim = maps.dim + self.bounds = maps.bounds + self.maps = maps self.neval = neval if nbatch is None: self.nbatch = neval @@ -33,7 +35,7 @@ def __init__( self.neval = -(-neval // nbatch) * nbatch self.device = device - + self.dtype = dtype def __call__(self, f: Callable, **kwargs): raise NotImplementedError("Subclasses must implement this method") @@ -46,17 +48,14 @@ def __init__( neval: int = 1000, nbatch: int = None, device="cpu", - adapt=False, - alpha=0.5, ): super().__init__(map, neval, nbatch, device) - self.adapt = adapt - self.alpha = alpha self.nitn = nitn def __call__(self, f: Callable, **kwargs): - u = torch.rand(self.nbatch, self.dim, device=self.device) - x, _ = self.map.forward(u) + # u = torch.rand(self.nbatch, self.dim, device=self.device) + # x, _ = self.map.forward(u) + x,_ = self.maps.sample(self.nbatch) f_values = f(x) f_size = len(f_values) if isinstance(f_values, (list, tuple)) else 1 type_fval = f_values.dtype if f_size == 1 else type(f_values[0].dtype) @@ -72,24 +71,19 @@ def __call__(self, f: Callable, **kwargs): mean[:] = 0 var[:] = 0 for _ in range(epoch): - y = torch.rand( - self.nbatch, self.dim, dtype=torch.float64, device=self.device - ) - x, jac = self.map.forward(y) - + # y = torch.rand( + # self.nbatch, self.dim, dtype=torch.float64, device=self.device + # ) + # x, jac = self.map.forward(y) + x, log_detJ = self.maps.sample(self.nbatch) f_values = f(x) - batch_results = self._multiply_by_jacobian(f_values, jac) + batch_results = self._multiply_by_jacobian(f_values, torch.exp(log_detJ) ) mean += torch.mean(batch_results, dim=-1) / epoch var += torch.var(batch_results, dim=-1) / (self.neval * epoch) - if self.adapt: - self.map.add_training_data(y, batch_results**2) - result.sum_neval += self.neval result.add(gvar.gvar(mean.item(), (var**0.5).item())) - if self.adapt: - self.map.adapt(alpha=self.alpha) return result @@ -109,13 +103,13 @@ def __init__( nitn: int = 10, neval=10000, nbatch=None, - n_burnin=500, + nburnin=500, device="cpu", adapt=False, alpha=0.5, ): super().__init__(map, nitn, neval, nbatch, device, adapt, alpha) - self.n_burnin = n_burnin + self.nburnin = nburnin def __call__( self, @@ -126,9 +120,11 @@ def __call__( **kwargs, ): epsilon = 1e-16 # Small value to ensure numerical stability - vars_shape = (self.nbatch, self.dim) - current_y = torch.rand(vars_shape, dtype=torch.float64, device=self.device) - current_x, current_jac = self.map.forward(current_y) + #vars_shape = (self.nbatch, self.dim) + # current_y = torch.rand(vars_shape, dtype=torch.float64, device=self.device) + # current_x, current_jac = self.map.forward(current_y) + current_x, current_jac = self.maps.sample(self.nbatch) + current_jac = torch.exp(current_jac) current_fval = f(current_x) current_weight = mix_rate / current_jac + (1 - mix_rate) * current_fval.abs() current_weight.masked_fill_(current_weight < epsilon, epsilon) @@ -171,7 +167,7 @@ def __call__( current_weight = torch.where(accept, new_weight, current_weight) current_jac = torch.where(accept, new_jac, current_jac) - if i < self.n_burnin and (self.adapt or itn == 0): + if i < self.nburnin and (self.adapt or itn == 0): continue elif i % thinning == 0: n_meas += 1 diff --git a/src/maps.py b/src/maps.py index 10e9769..34e28dc 100644 --- a/src/maps.py +++ b/src/maps.py @@ -2,6 +2,37 @@ import numpy as np from torch import nn +class NormalizingFlow(nn.Module): + def __init__(self, q0, maps): + super().__init__() + if not maps: + raise ValueError("Maps can not be empty.") + self.q0 = q0 + self.maps = maps + self.dim = maps[0].dim + self.bounds = maps[0].bounds + def forward(self, u): + log_detJ = torch.zeros(len(u), device=u.device) + for map in self.flows: + u, log_detj = map.forward(u) + log_detJ += log_detj + return u, log_detJ + def inverse(self, x): + log_detJ = torch.zeros(len(x), device=x.device) + for i in range(len(self.maps) - 1, -1, -1): + x, log_detj = self.maps[i].inverse(x) + log_detJ += log_detj + return x, log_detJ + def sample(self, nsample): + u, log_detJ = self.q0.sample(nsample) + for map in self.maps: + u, log_detj = map(u) + log_detJ += log_detj + return u, log_detJ + + + + class Map(nn.Module): def __init__(self, bounds, device="cpu"): super().__init__() @@ -17,10 +48,6 @@ def forward(self, u): def inverse(self, x): raise NotImplementedError("Subclasses must implement this method") - - def sample(self, nsample): - raise NotImplementedError("Subclasses must implement this method") - class Affine(Map): @@ -266,8 +293,7 @@ def inverse(self, x): return u, torch.log(jac) - def sample(self, nsample): - return super() + class NormalizingFlow(Map): @@ -281,5 +307,3 @@ def forward(self, u): def inverse(self, x): return self.flow_model.inverse(x) - def sample(self, nsample): - return self.flow_model.sample(nsample) From 6e1468342233c4d28590061853cf02e615c5be73 Mon Sep 17 00:00:00 2001 From: Tao Wang Date: Thu, 10 Oct 2024 21:23:05 -0400 Subject: [PATCH 03/10] wip --- src/base.py | 29 ++++++++++++++++++++++++++++ src/integrators.py | 44 ++++++++++++++++++++++++++----------------- src/maps.py | 47 ++++++++++++++++++---------------------------- 3 files changed, 74 insertions(+), 46 deletions(-) create mode 100644 src/base.py diff --git a/src/base.py b/src/base.py new file mode 100644 index 0000000..6e9a8fe --- /dev/null +++ b/src/base.py @@ -0,0 +1,29 @@ +import torch +from torch import nn +class BaseDistribution(nn.Module): + """ + Base distribution of a flow-based model + Parameters do not depend of target variable (as is the case for a VAE encoder) + """ + + def __init__(self, bounds): + super().__init__() + self.bounds = bounds + def sample(self, nsamples=1, **kwargs): + """Samples from base distribution + + Args: + num_samples: Number of samples to draw from the distriubtion + + Returns: + Samples drawn from the distribution + """ + raise NotImplementedError + +class Uniform(BaseDistribution): + """ + Multivariate uniform distribution + """ + + def __init__(self, bounds): + super().__init__(bounds) diff --git a/src/integrators.py b/src/integrators.py index 417e75f..c9ec5bc 100644 --- a/src/integrators.py +++ b/src/integrators.py @@ -1,7 +1,8 @@ from typing import Callable, Union, List, Tuple, Dict import torch from utils import RAvg -from maps import Map, Affine, NormalizingFlow +from maps import Map, Affine, CompositeMap +from base import Uniform import gvar @@ -12,19 +13,17 @@ class Integrator: def __init__( self, - # bounds: Union[List[Tuple[float, float]], np.ndarray], - maps: NormalizingFlow, + bounds: List[Tuple[float, float]], #, np.ndarray], + q0 = None, + maps = None, neval: int = 1000, nbatch: int = None, device="cpu", - dtype = torch.float32, # device=torch.device("cuda" if torch.cuda.is_available() else "cpu"), ): - #if not isinstance(map, Map): - # map = Affine(map) - - self.dim = maps.dim - self.bounds = maps.bounds + self.dim = len(bounds) + if not q0: + q0 = Uniform(bounds) self.maps = maps self.neval = neval if nbatch is None: @@ -35,10 +34,15 @@ def __init__( self.neval = -(-neval // nbatch) * nbatch self.device = device - self.dtype = dtype def __call__(self, f: Callable, **kwargs): raise NotImplementedError("Subclasses must implement this method") - + def sample(self, nsample, **kwargs): + u, log_detJ= self.q0.sample(nsample, **kwargs) + if not self.maps: + return u, log_detJ + else: + u, log_detj = self.maps.forward(u) + return u, log_detJ + log_detj class MonteCarlo(Integrator): def __init__( @@ -51,11 +55,11 @@ def __init__( ): super().__init__(map, neval, nbatch, device) self.nitn = nitn - + def __call__(self, f: Callable, **kwargs): # u = torch.rand(self.nbatch, self.dim, device=self.device) # x, _ = self.map.forward(u) - x,_ = self.maps.sample(self.nbatch) + x,_ = self.sample(self.nbatch) f_values = f(x) f_size = len(f_values) if isinstance(f_values, (list, tuple)) else 1 type_fval = f_values.dtype if f_size == 1 else type(f_values[0].dtype) @@ -75,7 +79,7 @@ def __call__(self, f: Callable, **kwargs): # self.nbatch, self.dim, dtype=torch.float64, device=self.device # ) # x, jac = self.map.forward(y) - x, log_detJ = self.maps.sample(self.nbatch) + x, log_detJ = self.sample(self.nbatch) f_values = f(x) batch_results = self._multiply_by_jacobian(f_values, torch.exp(log_detJ) ) @@ -84,9 +88,15 @@ def __call__(self, f: Callable, **kwargs): result.sum_neval += self.neval result.add(gvar.gvar(mean.item(), (var**0.5).item())) - return result - + + # def sample(self, nsample): + # u, log_detJ = self.q0.sample(nsample) + # for map in self.maps: + # u, log_detj = map(u) + # log_detJ += log_detj + # return u, log_detJ + def _multiply_by_jacobian(self, values, jac): # if isinstance(values, dict): # return {k: v * torch.exp(log_det_J) for k, v in values.items()} @@ -123,7 +133,7 @@ def __call__( #vars_shape = (self.nbatch, self.dim) # current_y = torch.rand(vars_shape, dtype=torch.float64, device=self.device) # current_x, current_jac = self.map.forward(current_y) - current_x, current_jac = self.maps.sample(self.nbatch) + current_x, current_jac = self.sample(self.nbatch) current_jac = torch.exp(current_jac) current_fval = f(current_x) current_weight = mix_rate / current_jac + (1 - mix_rate) * current_fval.abs() diff --git a/src/maps.py b/src/maps.py index 34e28dc..e4fdc0f 100644 --- a/src/maps.py +++ b/src/maps.py @@ -2,35 +2,6 @@ import numpy as np from torch import nn -class NormalizingFlow(nn.Module): - def __init__(self, q0, maps): - super().__init__() - if not maps: - raise ValueError("Maps can not be empty.") - self.q0 = q0 - self.maps = maps - self.dim = maps[0].dim - self.bounds = maps[0].bounds - def forward(self, u): - log_detJ = torch.zeros(len(u), device=u.device) - for map in self.flows: - u, log_detj = map.forward(u) - log_detJ += log_detj - return u, log_detJ - def inverse(self, x): - log_detJ = torch.zeros(len(x), device=x.device) - for i in range(len(self.maps) - 1, -1, -1): - x, log_detj = self.maps[i].inverse(x) - log_detJ += log_detj - return x, log_detJ - def sample(self, nsample): - u, log_detJ = self.q0.sample(nsample) - for map in self.maps: - u, log_detj = map(u) - log_detJ += log_detj - return u, log_detJ - - class Map(nn.Module): @@ -49,6 +20,24 @@ def forward(self, u): def inverse(self, x): raise NotImplementedError("Subclasses must implement this method") +class CompositeMap(Map): + def __init__(self, maps, device="cpu"): + if not maps: + raise ValueError("Maps can not be empty.") + super().__init__(maps[-1].bounds, device) + self.maps = maps + def forward(self, u): + log_detJ = torch.zeros(len(u), device=u.device) + for map in self.flows: + u, log_detj = map.forward(u) + log_detJ += log_detj + return u, log_detJ + def inverse(self, x): + log_detJ = torch.zeros(len(x), device=x.device) + for i in range(len(self.maps) - 1, -1, -1): + x, log_detj = self.maps[i].inverse(x) + log_detJ += log_detj + return x, log_detJ class Affine(Map): def __init__(self, bounds, device="cpu"): From 03e38e7dbe67b2a7bfc2ab6e5e152536f1b7a9c8 Mon Sep 17 00:00:00 2001 From: Tao Wang Date: Sat, 12 Oct 2024 01:09:00 -0400 Subject: [PATCH 04/10] mc_test is working --- src/base.py | 22 +++++++++++++--- src/integrators.py | 65 ++++++++++++++++++++++++++-------------------- src/maps.py | 37 +++++++++++++------------- src/mc_test.py | 16 +++++++----- 4 files changed, 83 insertions(+), 57 deletions(-) diff --git a/src/base.py b/src/base.py index 6e9a8fe..98764ad 100644 --- a/src/base.py +++ b/src/base.py @@ -1,14 +1,18 @@ import torch from torch import nn + + class BaseDistribution(nn.Module): """ Base distribution of a flow-based model Parameters do not depend of target variable (as is the case for a VAE encoder) """ - def __init__(self, bounds): + def __init__(self, bounds, device="cpu"): super().__init__() self.bounds = bounds + self.device = device + def sample(self, nsamples=1, **kwargs): """Samples from base distribution @@ -19,11 +23,21 @@ def sample(self, nsamples=1, **kwargs): Samples drawn from the distribution """ raise NotImplementedError - + + class Uniform(BaseDistribution): """ Multivariate uniform distribution """ - def __init__(self, bounds): - super().__init__(bounds) + def __init__(self, bounds, device="cpu"): + super().__init__(bounds, device) + + def sample(self, nsamples=1, **kwargs): + dim = len(self.bounds) + u = torch.rand((nsamples, dim), device=self.device) + log_detJ = torch.zeros(nsamples, device=self.device) + for i, bound in enumerate(self.bounds): + u[:, i] = (bound[1] - bound[0]) * u[:, i] + bound[0] + log_detJ += -torch.log(torch.tensor(bound[1] - bound[0])) + return u, log_detJ diff --git a/src/integrators.py b/src/integrators.py index c9ec5bc..435be5d 100644 --- a/src/integrators.py +++ b/src/integrators.py @@ -13,17 +13,20 @@ class Integrator: def __init__( self, - bounds: List[Tuple[float, float]], #, np.ndarray], - q0 = None, - maps = None, + bounds: List[Tuple[float, float]], # , np.ndarray], + q0=None, + maps=None, neval: int = 1000, nbatch: int = None, device="cpu", + adapt=False, # device=torch.device("cuda" if torch.cuda.is_available() else "cpu"), ): + self.adapt = adapt self.dim = len(bounds) if not q0: q0 = Uniform(bounds) + self.q0 = q0 self.maps = maps self.neval = neval if nbatch is None: @@ -34,32 +37,38 @@ def __init__( self.neval = -(-neval // nbatch) * nbatch self.device = device + def __call__(self, f: Callable, **kwargs): raise NotImplementedError("Subclasses must implement this method") + def sample(self, nsample, **kwargs): - u, log_detJ= self.q0.sample(nsample, **kwargs) + u, log_detJ = self.q0.sample(nsample) if not self.maps: return u, log_detJ else: u, log_detj = self.maps.forward(u) - return u, log_detJ + log_detj + return u, log_detJ + log_detj + class MonteCarlo(Integrator): def __init__( self, - map, + bounds: List[Tuple[float, float]], # , np.ndarray], + q0=None, + maps=None, nitn: int = 10, neval: int = 1000, nbatch: int = None, device="cpu", + adapt=False, ): - super().__init__(map, neval, nbatch, device) + super().__init__(bounds, q0, maps, neval, nbatch, device, adapt) self.nitn = nitn - + def __call__(self, f: Callable, **kwargs): # u = torch.rand(self.nbatch, self.dim, device=self.device) # x, _ = self.map.forward(u) - x,_ = self.sample(self.nbatch) + x, _ = self.sample(self.nbatch) f_values = f(x) f_size = len(f_values) if isinstance(f_values, (list, tuple)) else 1 type_fval = f_values.dtype if f_size == 1 else type(f_values[0].dtype) @@ -81,7 +90,9 @@ def __call__(self, f: Callable, **kwargs): # x, jac = self.map.forward(y) x, log_detJ = self.sample(self.nbatch) f_values = f(x) - batch_results = self._multiply_by_jacobian(f_values, torch.exp(log_detJ) ) + batch_results = self._multiply_by_jacobian( + f_values, torch.exp(log_detJ) + ) mean += torch.mean(batch_results, dim=-1) / epoch var += torch.var(batch_results, dim=-1) / (self.neval * epoch) @@ -89,14 +100,14 @@ def __call__(self, f: Callable, **kwargs): result.sum_neval += self.neval result.add(gvar.gvar(mean.item(), (var**0.5).item())) return result - + # def sample(self, nsample): # u, log_detJ = self.q0.sample(nsample) # for map in self.maps: # u, log_detj = map(u) # log_detJ += log_detj # return u, log_detJ - + def _multiply_by_jacobian(self, values, jac): # if isinstance(values, dict): # return {k: v * torch.exp(log_det_J) for k, v in values.items()} @@ -109,16 +120,17 @@ def _multiply_by_jacobian(self, values, jac): class MCMC(MonteCarlo): def __init__( self, - map: Map, + bounds: List[Tuple[float, float]], # , np.ndarray], + q0=None, + maps=None, nitn: int = 10, neval=10000, nbatch=None, nburnin=500, device="cpu", adapt=False, - alpha=0.5, ): - super().__init__(map, nitn, neval, nbatch, device, adapt, alpha) + super().__init__(bounds, q0, maps, nitn, neval, nbatch, device, adapt) self.nburnin = nburnin def __call__( @@ -130,10 +142,14 @@ def __call__( **kwargs, ): epsilon = 1e-16 # Small value to ensure numerical stability - #vars_shape = (self.nbatch, self.dim) - # current_y = torch.rand(vars_shape, dtype=torch.float64, device=self.device) - # current_x, current_jac = self.map.forward(current_y) - current_x, current_jac = self.sample(self.nbatch) + # vars_shape = (self.nbatch, self.dim) + current_y, current_jac = self.q0.sample(self.nbatch) + if self.maps: + current_x, jac = self.maps.forward(current_y) + current_jac += jac + else: + current_x = current_y + # current_x, current_jac = self.sample(self.nbatch) current_jac = torch.exp(current_jac) current_fval = f(current_x) current_weight = mix_rate / current_jac + (1 - mix_rate) * current_fval.abs() @@ -160,7 +176,7 @@ def __call__( for itn in range(self.nitn): for i in range(epoch): proposed_y[:] = self._propose(current_y, proposal_dist, **kwargs) - proposed_x[:], new_jac = self.map.forward(proposed_y) + proposed_x[:], new_jac = self.maps.forward(proposed_y) new_fval[:] = f(proposed_x) new_weight = mix_rate / new_jac + (1 - mix_rate) * new_fval.abs() @@ -177,7 +193,7 @@ def __call__( current_weight = torch.where(accept, new_weight, current_weight) current_jac = torch.where(accept, new_jac, current_jac) - if i < self.nburnin and (self.adapt or itn == 0): + if i < self.nburnin and itn == 0: continue elif i % thinning == 0: n_meas += 1 @@ -190,10 +206,6 @@ def __call__( mean_ref += torch.mean(batch_results_ref, dim=-1) / epoch var_ref += torch.var(batch_results_ref, dim=-1) / epoch - if self.adapt: - self.map.add_training_data( - current_y, (current_fval * current_jac) ** 2 - ) result.sum_neval += self.neval result.add(gvar.gvar(mean.item(), ((var / n_meas) ** 0.5).item())) result_ref.sum_neval += self.nbatch @@ -201,9 +213,6 @@ def __call__( gvar.gvar(mean_ref.item(), ((var_ref / n_meas) ** 0.5).item()) ) - if self.adapt: - self.map.adapt(alpha=self.alpha) - return result / result_ref def _propose(self, u, proposal_dist, **kwargs): diff --git a/src/maps.py b/src/maps.py index e4fdc0f..ef4a884 100644 --- a/src/maps.py +++ b/src/maps.py @@ -3,7 +3,6 @@ from torch import nn - class Map(nn.Module): def __init__(self, bounds, device="cpu"): super().__init__() @@ -16,22 +15,25 @@ def __init__(self, bounds, device="cpu"): def forward(self, u): raise NotImplementedError("Subclasses must implement this method") - + def inverse(self, x): raise NotImplementedError("Subclasses must implement this method") + class CompositeMap(Map): def __init__(self, maps, device="cpu"): if not maps: raise ValueError("Maps can not be empty.") super().__init__(maps[-1].bounds, device) self.maps = maps + def forward(self, u): log_detJ = torch.zeros(len(u), device=u.device) - for map in self.flows: - u, log_detj = map.forward(u) - log_detJ += log_detj + for map in self.maps: + u, log_detj = map.forward(u) + log_detJ += log_detj return u, log_detJ + def inverse(self, x): log_detJ = torch.zeros(len(x), device=x.device) for i in range(len(self.maps) - 1, -1, -1): @@ -39,6 +41,7 @@ def inverse(self, x): log_detJ += log_detj return x, log_detJ + class Affine(Map): def __init__(self, bounds, device="cpu"): super().__init__(bounds, device) @@ -49,8 +52,9 @@ def forward(self, u): return u * self._A + self.bounds[:, 0], torch.log(self._jac1.repeat(u.shape[0])) def inverse(self, x): - return (x - self.bounds[:, 0]) / self._A, torch.log(self._jac1.repeat(x.shape[0])) - + return (x - self.bounds[:, 0]) / self._A, torch.log( + self._jac1.repeat(x.shape[0]) + ) class Vegas(Map): @@ -282,17 +286,14 @@ def inverse(self, x): return u, torch.log(jac) - - -class NormalizingFlow(Map): - def __init__(self, bounds, flow_model, device="cpu"): - super().__init__(bounds, device) - self.flow_model = flow_model.to(device) +# class NormalizingFlow(Map): +# def __init__(self, bounds, flow_model, device="cpu"): +# super().__init__(bounds, device) +# self.flow_model = flow_model.to(device) - def forward(self, u): - return self.flow_model.forward(u) +# def forward(self, u): +# return self.flow_model.forward(u) - def inverse(self, x): - return self.flow_model.inverse(x) - +# def inverse(self, x): +# return self.flow_model.inverse(x) diff --git a/src/mc_test.py b/src/mc_test.py index 932fe6a..c55df20 100644 --- a/src/mc_test.py +++ b/src/mc_test.py @@ -28,21 +28,23 @@ def half_sphere_integrand(x): dim = 2 -map_spec = [(-1, 1), (-1, 1)] +bounds = [(-1, 1), (-1, 1)] -affine_map = Affine(map_spec, device=device) -# vegas_map = Vegas(map_spec, device=device) +affine_map = Affine(bounds, device=device) +# vegas_map = Vegas(bounds, device=device) # Monte Carlo integration print("Calculate the area of the unit circle using Monte Carlo integration...") -mc_integrator = MonteCarlo(affine_map, neval=400000, batch_size=1000, device=device) +mc_integrator = MonteCarlo( + bounds, maps=affine_map, neval=400000, nbatch=1000, device=device +) res = mc_integrator(unit_circle_integrand) print("Plain MC Integral results:") print(f" Integral: {res.mean}") print(f" Error: {res.sdev}") -res = MonteCarlo(map_spec, neval=400000, batch_size=1000, device=device)( +res = MonteCarlo(bounds, neval=400000, nbatch=1000, device=device)( unit_circle_integrand ) print("Plain MC Integral results:") @@ -50,7 +52,7 @@ def half_sphere_integrand(x): print(f" Error: {res.sdev}") mcmc_integrator = MCMC( - map_spec, neval=400000, batch_size=1000, n_burnin=100, device=device + bounds, maps=affine_map, neval=400000, nbatch=1000, nburnin=100, device=device ) res = mcmc_integrator(unit_circle_integrand, mix_rate=0.5) print("MCMC Integral results:") @@ -66,7 +68,7 @@ def half_sphere_integrand(x): print(f" Error: {res.sdev}") mcmc_integrator = MCMC( - map_spec, neval=400000, batch_size=1000, n_burnin=100, device=device + bounds, maps=affine_map, neval=400000, nbatch=1000, nburnin=100, device=device ) res = mcmc_integrator(half_sphere_integrand, mix_rate=0.5) print("MCMC Integral results:") From 18610331d25f7de71122c918d6a19c863bfe038d Mon Sep 17 00:00:00 2001 From: houpc Date: Tue, 15 Oct 2024 02:21:02 +0800 Subject: [PATCH 05/10] bugfix and GPU compatibility --- src/base.py | 20 +++++++++++------ src/integrators.py | 56 +++++++++++++++++++++++----------------------- src/mc_test.py | 15 +++++++------ 3 files changed, 49 insertions(+), 42 deletions(-) diff --git a/src/base.py b/src/base.py index 98764ad..d6aa62c 100644 --- a/src/base.py +++ b/src/base.py @@ -1,5 +1,6 @@ import torch from torch import nn +import numpy as np class BaseDistribution(nn.Module): @@ -10,7 +11,12 @@ class BaseDistribution(nn.Module): def __init__(self, bounds, device="cpu"): super().__init__() - self.bounds = bounds + # self.bounds = bounds + if isinstance(bounds, (list, np.ndarray)): + self.bounds = torch.tensor(bounds, dtype=torch.float64, device=device) + else: + raise ValueError("Unsupported map specification") + self.dim = self.bounds.shape[0] self.device = device def sample(self, nsamples=1, **kwargs): @@ -32,12 +38,12 @@ class Uniform(BaseDistribution): def __init__(self, bounds, device="cpu"): super().__init__(bounds, device) + self._rangebounds = self.bounds[:, 1] - self.bounds[:, 0] def sample(self, nsamples=1, **kwargs): - dim = len(self.bounds) - u = torch.rand((nsamples, dim), device=self.device) - log_detJ = torch.zeros(nsamples, device=self.device) - for i, bound in enumerate(self.bounds): - u[:, i] = (bound[1] - bound[0]) * u[:, i] + bound[0] - log_detJ += -torch.log(torch.tensor(bound[1] - bound[0])) + u = ( + torch.rand((nsamples, self.dim), device=self.device) * self._rangebounds + + self.bounds[:, 0] + ) + log_detJ = torch.log(self._rangebounds).sum().repeat(nsamples) return u, log_detJ diff --git a/src/integrators.py b/src/integrators.py index 435be5d..65fce77 100644 --- a/src/integrators.py +++ b/src/integrators.py @@ -4,6 +4,7 @@ from maps import Map, Affine, CompositeMap from base import Uniform import gvar +import numpy as np class Integrator: @@ -13,19 +14,19 @@ class Integrator: def __init__( self, - bounds: List[Tuple[float, float]], # , np.ndarray], + bounds: Union[List[Tuple[float, float]], np.ndarray], q0=None, maps=None, neval: int = 1000, nbatch: int = None, device="cpu", adapt=False, - # device=torch.device("cuda" if torch.cuda.is_available() else "cpu"), ): self.adapt = adapt self.dim = len(bounds) if not q0: - q0 = Uniform(bounds) + q0 = Uniform(bounds, device=device) + self.bounds = torch.tensor(bounds, dtype=torch.float64, device=device) self.q0 = q0 self.maps = maps self.neval = neval @@ -53,7 +54,7 @@ def sample(self, nsample, **kwargs): class MonteCarlo(Integrator): def __init__( self, - bounds: List[Tuple[float, float]], # , np.ndarray], + bounds: Union[List[Tuple[float, float]], np.ndarray], q0=None, maps=None, nitn: int = 10, @@ -66,8 +67,6 @@ def __init__( self.nitn = nitn def __call__(self, f: Callable, **kwargs): - # u = torch.rand(self.nbatch, self.dim, device=self.device) - # x, _ = self.map.forward(u) x, _ = self.sample(self.nbatch) f_values = f(x) f_size = len(f_values) if isinstance(f_values, (list, tuple)) else 1 @@ -84,10 +83,6 @@ def __call__(self, f: Callable, **kwargs): mean[:] = 0 var[:] = 0 for _ in range(epoch): - # y = torch.rand( - # self.nbatch, self.dim, dtype=torch.float64, device=self.device - # ) - # x, jac = self.map.forward(y) x, log_detJ = self.sample(self.nbatch) f_values = f(x) batch_results = self._multiply_by_jacobian( @@ -101,13 +96,6 @@ def __call__(self, f: Callable, **kwargs): result.add(gvar.gvar(mean.item(), (var**0.5).item())) return result - # def sample(self, nsample): - # u, log_detJ = self.q0.sample(nsample) - # for map in self.maps: - # u, log_detj = map(u) - # log_detJ += log_detj - # return u, log_detJ - def _multiply_by_jacobian(self, values, jac): # if isinstance(values, dict): # return {k: v * torch.exp(log_det_J) for k, v in values.items()} @@ -120,7 +108,7 @@ def _multiply_by_jacobian(self, values, jac): class MCMC(MonteCarlo): def __init__( self, - bounds: List[Tuple[float, float]], # , np.ndarray], + bounds: Union[List[Tuple[float, float]], np.ndarray], q0=None, maps=None, nitn: int = 10, @@ -132,6 +120,9 @@ def __init__( ): super().__init__(bounds, q0, maps, nitn, neval, nbatch, device, adapt) self.nburnin = nburnin + if maps is None: + self.maps = Affine([(0, 1)] * self.dim, device=device) + self._rangebounds = self.bounds[:, 1] - self.bounds[:, 0] def __call__( self, @@ -144,12 +135,11 @@ def __call__( epsilon = 1e-16 # Small value to ensure numerical stability # vars_shape = (self.nbatch, self.dim) current_y, current_jac = self.q0.sample(self.nbatch) - if self.maps: - current_x, jac = self.maps.forward(current_y) - current_jac += jac - else: - current_x = current_y - # current_x, current_jac = self.sample(self.nbatch) + # if self.maps: + current_x, detJ = self.maps.forward(current_y) + current_jac += detJ + # else: + # current_x = current_y current_jac = torch.exp(current_jac) current_fval = f(current_x) current_weight = mix_rate / current_jac + (1 - mix_rate) * current_fval.abs() @@ -157,7 +147,7 @@ def __call__( # current_fval.masked_fill_(current_fval.abs() < epsilon, epsilon) proposed_y = torch.empty_like(current_y) - proposed_x = torch.empty_like(proposed_y) + proposed_x = torch.empty_like(current_x) new_fval = torch.empty_like(current_fval) new_weight = torch.empty_like(current_weight) @@ -177,6 +167,7 @@ def __call__( for i in range(epoch): proposed_y[:] = self._propose(current_y, proposal_dist, **kwargs) proposed_x[:], new_jac = self.maps.forward(proposed_y) + new_jac = torch.exp(new_jac) new_fval[:] = f(proposed_x) new_weight = mix_rate / new_jac + (1 - mix_rate) * new_fval.abs() @@ -213,14 +204,23 @@ def __call__( gvar.gvar(mean_ref.item(), ((var_ref / n_meas) ** 0.5).item()) ) - return result / result_ref + return result / result_ref * self._rangebounds.prod() def _propose(self, u, proposal_dist, **kwargs): if proposal_dist == "random_walk": step_size = kwargs.get("step_size", 0.2) - return (u + (torch.rand_like(u) - 0.5) * step_size) % 1.0 + step_sizes = self._rangebounds * step_size + step = ( + torch.empty(self.dim, device=self.device).uniform_(-1, 1) * step_sizes + ) + new_u = (u + step - self.bounds[:, 0]) % self._rangebounds + self.bounds[ + :, 0 + ] + return new_u + # return (u + (torch.rand_like(u) - 0.5) * step_size) % 1.0 elif proposal_dist == "uniform": - return torch.rand_like(u) + # return torch.rand_like(u) + return torch.rand_like(u) * self._rangebounds + self.bounds[:, 0] # elif proposal_dist == "gaussian": # mean = kwargs.get("mean", torch.zeros_like(u)) # std = kwargs.get("std", torch.ones_like(u)) diff --git a/src/mc_test.py b/src/mc_test.py index c55df20..700946e 100644 --- a/src/mc_test.py +++ b/src/mc_test.py @@ -5,6 +5,7 @@ set_seed(42) device = get_device() +# device = torch.device("cpu") def unit_circle_integrand(x): @@ -37,7 +38,11 @@ def half_sphere_integrand(x): print("Calculate the area of the unit circle using Monte Carlo integration...") mc_integrator = MonteCarlo( - bounds, maps=affine_map, neval=400000, nbatch=1000, device=device + # bounds, maps=affine_map, neval=400000, nbatch=1000, device=device + bounds, + neval=400000, + nbatch=1000, + device=device, ) res = mc_integrator(unit_circle_integrand) print("Plain MC Integral results:") @@ -51,9 +56,7 @@ def half_sphere_integrand(x): print(f" Integral: {res.mean}") print(f" Error: {res.sdev}") -mcmc_integrator = MCMC( - bounds, maps=affine_map, neval=400000, nbatch=1000, nburnin=100, device=device -) +mcmc_integrator = MCMC(bounds, neval=400000, nbatch=1000, nburnin=100, device=device) res = mcmc_integrator(unit_circle_integrand, mix_rate=0.5) print("MCMC Integral results:") print(f" Integral: {res.mean}") @@ -67,9 +70,7 @@ def half_sphere_integrand(x): print(f" Integral: {res.mean}") print(f" Error: {res.sdev}") -mcmc_integrator = MCMC( - bounds, maps=affine_map, neval=400000, nbatch=1000, nburnin=100, device=device -) +mcmc_integrator = MCMC(bounds, neval=400000, nbatch=1000, nburnin=100, device=device) res = mcmc_integrator(half_sphere_integrand, mix_rate=0.5) print("MCMC Integral results:") print(f" Integral: {res.mean}") From 7123f22fb7d80f5412192d64eb8e05e13d5df9e6 Mon Sep 17 00:00:00 2001 From: Tao Wang Date: Tue, 22 Oct 2024 00:23:51 -0400 Subject: [PATCH 06/10] Refactor api wip --- src/base.py | 12 ++- src/integrators.py | 234 +++++++++++++++++++++++++-------------------- src/maps.py | 29 +++--- src/mc_test.py | 4 +- 4 files changed, 156 insertions(+), 123 deletions(-) diff --git a/src/base.py b/src/base.py index d6aa62c..ab7ffd0 100644 --- a/src/base.py +++ b/src/base.py @@ -9,11 +9,12 @@ class BaseDistribution(nn.Module): Parameters do not depend of target variable (as is the case for a VAE encoder) """ - def __init__(self, bounds, device="cpu"): + def __init__(self, bounds, device="cpu", dtype=torch.float64): super().__init__() + self.dtype = dtype # self.bounds = bounds if isinstance(bounds, (list, np.ndarray)): - self.bounds = torch.tensor(bounds, dtype=torch.float64, device=device) + self.bounds = torch.tensor(bounds, dtype=dtype, device=device) else: raise ValueError("Unsupported map specification") self.dim = self.bounds.shape[0] @@ -36,13 +37,14 @@ class Uniform(BaseDistribution): Multivariate uniform distribution """ - def __init__(self, bounds, device="cpu"): - super().__init__(bounds, device) + def __init__(self, bounds, device="cpu", dtype=torch.float64): + super().__init__(bounds, device, dtype) self._rangebounds = self.bounds[:, 1] - self.bounds[:, 0] def sample(self, nsamples=1, **kwargs): u = ( - torch.rand((nsamples, self.dim), device=self.device) * self._rangebounds + torch.rand((nsamples, self.dim), device=self.device, dtype=self.dtype) + * self._rangebounds + self.bounds[:, 0] ) log_detJ = torch.log(self._rangebounds).sum().repeat(nsamples) diff --git a/src/integrators.py b/src/integrators.py index 65fce77..9e19fc9 100644 --- a/src/integrators.py +++ b/src/integrators.py @@ -1,7 +1,7 @@ from typing import Callable, Union, List, Tuple, Dict import torch from utils import RAvg -from maps import Map, Affine, CompositeMap +from maps import Map, Linear, CompositeMap from base import Uniform import gvar import numpy as np @@ -9,25 +9,32 @@ class Integrator: """ - Base class for all integrators. + Base class for all integrators. This class is designed to handle integration tasks + over a specified domain (bounds) using a sampling method (q0) and optional + transformation maps. """ def __init__( self, - bounds: Union[List[Tuple[float, float]], np.ndarray], + bounds, q0=None, maps=None, neval: int = 1000, nbatch: int = None, device="cpu", - adapt=False, + dtype=torch.float64, ): - self.adapt = adapt + if not isinstance(bounds, (list, np.ndarray)): + raise TypeError("bounds must be a list or a NumPy array.") + self.dtype = dtype self.dim = len(bounds) if not q0: - q0 = Uniform(bounds, device=device) - self.bounds = torch.tensor(bounds, dtype=torch.float64, device=device) + q0 = Uniform(bounds, device=device, dtype=dtype) + self.bounds = torch.tensor(bounds, dtype=dtype, device=device) self.q0 = q0 + if maps: + if not self.dtype == maps.dtype: + raise ValueError("Float type of maps should be same as integrator.") self.maps = maps self.neval = neval if nbatch is None: @@ -54,46 +61,41 @@ def sample(self, nsample, **kwargs): class MonteCarlo(Integrator): def __init__( self, - bounds: Union[List[Tuple[float, float]], np.ndarray], + bounds, q0=None, maps=None, - nitn: int = 10, neval: int = 1000, nbatch: int = None, device="cpu", - adapt=False, + dtype=torch.float64, ): - super().__init__(bounds, q0, maps, neval, nbatch, device, adapt) - self.nitn = nitn + super().__init__(bounds, q0, maps, neval, nbatch, device, dtype) def __call__(self, f: Callable, **kwargs): x, _ = self.sample(self.nbatch) f_values = f(x) f_size = len(f_values) if isinstance(f_values, (list, tuple)) else 1 - type_fval = f_values.dtype if f_size == 1 else type(f_values[0].dtype) - - mean = torch.zeros(f_size, dtype=type_fval, device=self.device) - var = torch.zeros(f_size, dtype=type_fval, device=self.device) + # type_fval = f_values.dtype if f_size == 1 else type(f_values[0].dtype) + # mean = torch.zeros(f_size, dtype=type_fval, device=self.device) + # var = torch.zeros(f_size, dtype=type_fval, device=self.device) # var = torch.zeros((f_size, f_size), dtype=type_fval, device=self.device) - - result = RAvg(weighted=self.adapt) + mean = torch.zeros(f_size, dtype=self.dtype, device=self.device) + var = torch.zeros(f_size, dtype=self.dtype, device=self.device) + result = RAvg() epoch = self.neval // self.nbatch - for itn in range(self.nitn): - mean[:] = 0 - var[:] = 0 - for _ in range(epoch): - x, log_detJ = self.sample(self.nbatch) - f_values = f(x) - batch_results = self._multiply_by_jacobian( - f_values, torch.exp(log_detJ) - ) + mean[:] = 0 + var[:] = 0 + for _ in range(epoch): + x, log_detJ = self.sample(self.nbatch) + f_values = f(x) + batch_results = self._multiply_by_jacobian(f_values, torch.exp(log_detJ)) - mean += torch.mean(batch_results, dim=-1) / epoch - var += torch.var(batch_results, dim=-1) / (self.neval * epoch) + mean += torch.mean(batch_results, dim=-1) / epoch + var += torch.var(batch_results, dim=-1) / (self.neval * epoch) - result.sum_neval += self.neval - result.add(gvar.gvar(mean.item(), (var**0.5).item())) + result.sum_neval += self.neval + result.add(gvar.gvar(mean.item(), (var**0.5).item())) return result def _multiply_by_jacobian(self, values, jac): @@ -105,29 +107,48 @@ def _multiply_by_jacobian(self, values, jac): return values * jac +def random_walk(dim, bounds, device, dtype, u, **kwargs): + rangebounds = bounds[:, 1] - bounds[:, 0] + step_size = kwargs.get("step_size", 0.2) + step_sizes = rangebounds * step_size + step = torch.empty(dim, device=device, dtype=dtype).uniform_(-1, 1) * step_sizes + new_u = (u + step - bounds[:, 0]) % rangebounds + bounds[:, 0] + return new_u + + +def uniform(dim, bounds, device, dtype, u, **kwargs): + rangebounds = bounds[:, 1] - bounds[:, 0] + return torch.rand_like(u) * rangebounds + bounds[:, 0] + + +def gaussian(dim, bounds, device, dtype, u, **kwargs): + mean = kwargs.get("mean", torch.zeros_like(u)) + std = kwargs.get("std", torch.ones_like(u)) + return torch.normal(mean, std) + + class MCMC(MonteCarlo): def __init__( self, - bounds: Union[List[Tuple[float, float]], np.ndarray], + bounds, q0=None, maps=None, - nitn: int = 10, neval=10000, nbatch=None, nburnin=500, device="cpu", - adapt=False, + dtype=torch.float64, ): - super().__init__(bounds, q0, maps, nitn, neval, nbatch, device, adapt) + super().__init__(bounds, q0, maps, neval, nbatch, device, dtype) self.nburnin = nburnin if maps is None: - self.maps = Affine([(0, 1)] * self.dim, device=device) + self.maps = Linear([(0, 1)] * self.dim, device=device) self._rangebounds = self.bounds[:, 1] - self.bounds[:, 0] def __call__( self, f: Callable, - proposal_dist="uniform", + proposal_dist: Callable = uniform, thinning=1, mix_rate=0.0, **kwargs, @@ -146,84 +167,93 @@ def __call__( current_weight.masked_fill_(current_weight < epsilon, epsilon) # current_fval.masked_fill_(current_fval.abs() < epsilon, epsilon) - proposed_y = torch.empty_like(current_y) - proposed_x = torch.empty_like(current_x) - new_fval = torch.empty_like(current_fval) - new_weight = torch.empty_like(current_weight) + # proposed_y = torch.empty_like(current_y) + # proposed_x = torch.empty_like(current_x) + # new_fval = torch.empty_like(current_fval) + # new_weight = torch.empty_like(current_weight) f_size = len(current_fval) if isinstance(current_fval, (list, tuple)) else 1 - type_fval = current_fval.dtype if f_size == 1 else type(current_fval[0].dtype) - mean = torch.zeros(f_size, dtype=type_fval, device=self.device) + # type_fval = current_fval.dtype if f_size == 1 else type(current_fval[0].dtype) + # mean = torch.zeros(f_size, dtype=type_fval, device=self.device) + mean = torch.zeros(f_size, dtype=self.dtype, device=self.device) mean_ref = torch.zeros_like(mean) - var = torch.zeros(f_size, dtype=type_fval, device=self.device) + # var = torch.zeros(f_size, dtype=type_fval, device=self.device) + var = torch.zeros(f_size, dtype=self.dtype, device=self.device) var_ref = torch.zeros_like(mean) - result = RAvg(weighted=self.adapt) - result_ref = RAvg(weighted=self.adapt) + result = RAvg() + result_ref = RAvg() epoch = self.neval // self.nbatch n_meas = 0 - for itn in range(self.nitn): - for i in range(epoch): - proposed_y[:] = self._propose(current_y, proposal_dist, **kwargs) - proposed_x[:], new_jac = self.maps.forward(proposed_y) - new_jac = torch.exp(new_jac) - new_fval[:] = f(proposed_x) - new_weight = mix_rate / new_jac + (1 - mix_rate) * new_fval.abs() + def _propose(current_y, current_fval, current_weight, current_jac): + proposed_y = proposal_dist( + self.dim, self.bounds, self.device, self.dtype, current_y, **kwargs + ) + proposed_x, new_jac = self.maps.forward(proposed_y) + new_jac = torch.exp(new_jac) + + new_fval = f(proposed_x) + new_weight = mix_rate / new_jac + (1 - mix_rate) * new_fval.abs() - acceptance_probs = new_weight / current_weight * new_jac / current_jac + acceptance_probs = new_weight / current_weight * new_jac / current_jac - accept = ( - torch.rand(self.nbatch, dtype=torch.float64, device=self.device) - <= acceptance_probs - ) + accept = ( + torch.rand(self.nbatch, dtype=self.dtype, device=self.device) + <= acceptance_probs + ) - current_y = torch.where(accept.unsqueeze(1), proposed_y, current_y) - current_fval = torch.where(accept, new_fval, current_fval) - current_weight = torch.where(accept, new_weight, current_weight) - current_jac = torch.where(accept, new_jac, current_jac) - - if i < self.nburnin and itn == 0: - continue - elif i % thinning == 0: - n_meas += 1 - batch_results = current_fval / current_weight - - mean += torch.mean(batch_results, dim=-1) / epoch - var += torch.var(batch_results, dim=-1) / epoch - - batch_results_ref = 1 / (current_jac * current_weight) - mean_ref += torch.mean(batch_results_ref, dim=-1) / epoch - var_ref += torch.var(batch_results_ref, dim=-1) / epoch - - result.sum_neval += self.neval - result.add(gvar.gvar(mean.item(), ((var / n_meas) ** 0.5).item())) - result_ref.sum_neval += self.nbatch - result_ref.add( - gvar.gvar(mean_ref.item(), ((var_ref / n_meas) ** 0.5).item()) + current_y = torch.where(accept.unsqueeze(1), proposed_y, current_y) + current_fval = torch.where(accept, new_fval, current_fval) + current_weight = torch.where(accept, new_weight, current_weight) + current_jac = torch.where(accept, new_jac, current_jac) + return current_y, current_fval, current_weight, current_jac + + for i in range(self.nburnin): + current_y, current_fval, current_weight, current_jac = _propose( + current_y, current_fval, current_weight, current_jac ) + for i in range(epoch // thinning): + for j in range(thinning): + current_y, current_fval, current_weight, current_jac = _propose( + current_y, current_fval, current_weight, current_jac + ) + n_meas += 1 + batch_results = current_fval / current_weight + + mean += torch.mean(batch_results, dim=-1) / epoch + var += torch.var(batch_results, dim=-1) / epoch + + batch_results_ref = 1 / (current_jac * current_weight) + mean_ref += torch.mean(batch_results_ref, dim=-1) / epoch + var_ref += torch.var(batch_results_ref, dim=-1) / epoch + + result.sum_neval += self.neval + result.add(gvar.gvar(mean.item(), ((var / n_meas) ** 0.5).item())) + result_ref.sum_neval += self.nbatch + result_ref.add(gvar.gvar(mean_ref.item(), ((var_ref / n_meas) ** 0.5).item())) return result / result_ref * self._rangebounds.prod() - def _propose(self, u, proposal_dist, **kwargs): - if proposal_dist == "random_walk": - step_size = kwargs.get("step_size", 0.2) - step_sizes = self._rangebounds * step_size - step = ( - torch.empty(self.dim, device=self.device).uniform_(-1, 1) * step_sizes - ) - new_u = (u + step - self.bounds[:, 0]) % self._rangebounds + self.bounds[ - :, 0 - ] - return new_u - # return (u + (torch.rand_like(u) - 0.5) * step_size) % 1.0 - elif proposal_dist == "uniform": - # return torch.rand_like(u) - return torch.rand_like(u) * self._rangebounds + self.bounds[:, 0] - # elif proposal_dist == "gaussian": - # mean = kwargs.get("mean", torch.zeros_like(u)) - # std = kwargs.get("std", torch.ones_like(u)) - # return torch.normal(mean, std) - else: - raise ValueError(f"Unknown proposal distribution: {proposal_dist}") + # def _propose(self, u, proposal_dist, **kwargs): + # if proposal_dist == "random_walk": + # step_size = kwargs.get("step_size", 0.2) + # step_sizes = self._rangebounds * step_size + # step = ( + # torch.empty(self.dim, device=self.device).uniform_(-1, 1) * step_sizes + # ) + # new_u = (u + step - self.bounds[:, 0]) % self._rangebounds + self.bounds[ + # :, 0 + # ] + # return new_u + # # return (u + (torch.rand_like(u) - 0.5) * step_size) % 1.0 + # elif proposal_dist == "uniform": + # # return torch.rand_like(u) + # return torch.rand_like(u) * self._rangebounds + self.bounds[:, 0] + # # elif proposal_dist == "gaussian": + # # mean = kwargs.get("mean", torch.zeros_like(u)) + # # std = kwargs.get("std", torch.ones_like(u)) + # # return torch.normal(mean, std) + # else: + # raise ValueError(f"Unknown proposal distribution: {proposal_dist}") diff --git a/src/maps.py b/src/maps.py index ef4a884..2b6a6df 100644 --- a/src/maps.py +++ b/src/maps.py @@ -4,14 +4,15 @@ class Map(nn.Module): - def __init__(self, bounds, device="cpu"): + def __init__(self, bounds, device="cpu", dtype=torch.float64): super().__init__() if isinstance(bounds, (list, np.ndarray)): - self.bounds = torch.tensor(bounds, dtype=torch.float64, device=device) + self.bounds = torch.tensor(bounds, dtype=dtype, device=device) else: raise ValueError("Unsupported map specification") self.dim = self.bounds.shape[0] self.device = device + self.dtype = dtype def forward(self, u): raise NotImplementedError("Subclasses must implement this method") @@ -21,30 +22,30 @@ def inverse(self, x): class CompositeMap(Map): - def __init__(self, maps, device="cpu"): + def __init__(self, maps, device="cpu", dtype=torch.float64): if not maps: raise ValueError("Maps can not be empty.") - super().__init__(maps[-1].bounds, device) + super().__init__(maps[-1].bounds, device, dtype) self.maps = maps def forward(self, u): - log_detJ = torch.zeros(len(u), device=u.device) + log_detJ = torch.zeros(len(u), device=u.device, dtype=self.dtype) for map in self.maps: u, log_detj = map.forward(u) log_detJ += log_detj return u, log_detJ def inverse(self, x): - log_detJ = torch.zeros(len(x), device=x.device) + log_detJ = torch.zeros(len(x), device=x.device, dtype=self.dtype) for i in range(len(self.maps) - 1, -1, -1): x, log_detj = self.maps[i].inverse(x) log_detJ += log_detj return x, log_detJ -class Affine(Map): - def __init__(self, bounds, device="cpu"): - super().__init__(bounds, device) +class Linear(Map): + def __init__(self, bounds, device="cpu", dtype=torch.float64): + super().__init__(bounds, device, dtype) self._A = self.bounds[:, 1] - self.bounds[:, 0] self._jac1 = torch.prod(self._A) @@ -58,8 +59,8 @@ def inverse(self, x): class Vegas(Map): - def __init__(self, bounds, ninc=1000, alpha=0.5, device="cpu"): - super().__init__(bounds, device) + def __init__(self, bounds, ninc=1000, alpha=0.5, device="cpu", dtype=torch.float64): + super().__init__(bounds, device, dtype) # self.nbin = nbin self.alpha = alpha if isinstance(ninc, int): @@ -68,10 +69,10 @@ def __init__(self, bounds, ninc=1000, alpha=0.5, device="cpu"): self.ninc = torch.tensor(ninc, dtype=torch.int32, device=device) self.inc = torch.empty( - self.dim, self.ninc.max(), dtype=torch.float64, device=self.device + self.dim, self.ninc.max(), dtype=self.dtype, device=self.device ) self.grid = torch.empty( - self.dim, self.ninc.max() + 1, dtype=torch.float64, device=self.device + self.dim, self.ninc.max() + 1, dtype=self.dtype, device=self.device ) for d in range(self.dim): @@ -79,7 +80,7 @@ def __init__(self, bounds, ninc=1000, alpha=0.5, device="cpu"): self.bounds[d, 0], self.bounds[d, 1], self.ninc[d] + 1, - dtype=torch.float64, + dtype=self.dtype, device=self.device, ) self.inc[d, : self.ninc[d]] = ( diff --git a/src/mc_test.py b/src/mc_test.py index 700946e..47740c0 100644 --- a/src/mc_test.py +++ b/src/mc_test.py @@ -1,6 +1,6 @@ import torch from integrators import MonteCarlo, MCMC -from maps import Vegas, Affine +from maps import Vegas, Linear from utils import set_seed, get_device set_seed(42) @@ -31,7 +31,7 @@ def half_sphere_integrand(x): dim = 2 bounds = [(-1, 1), (-1, 1)] -affine_map = Affine(bounds, device=device) +affine_map = Linear(bounds, device=device) # vegas_map = Vegas(bounds, device=device) # Monte Carlo integration From 8ea252bd01235a4a619570f773fe7a991578fb84 Mon Sep 17 00:00:00 2001 From: houpc Date: Wed, 23 Oct 2024 10:25:41 +0800 Subject: [PATCH 07/10] support integrand value is list or tuple and update statistics --- src/integrators.py | 165 ++++++++++++++++++++++----------------------- src/mc_test.py | 64 +++++++++++++----- 2 files changed, 128 insertions(+), 101 deletions(-) diff --git a/src/integrators.py b/src/integrators.py index 9e19fc9..a97910b 100644 --- a/src/integrators.py +++ b/src/integrators.py @@ -75,28 +75,33 @@ def __call__(self, f: Callable, **kwargs): x, _ = self.sample(self.nbatch) f_values = f(x) f_size = len(f_values) if isinstance(f_values, (list, tuple)) else 1 - # type_fval = f_values.dtype if f_size == 1 else type(f_values[0].dtype) - # mean = torch.zeros(f_size, dtype=type_fval, device=self.device) - # var = torch.zeros(f_size, dtype=type_fval, device=self.device) - # var = torch.zeros((f_size, f_size), dtype=type_fval, device=self.device) - mean = torch.zeros(f_size, dtype=self.dtype, device=self.device) - var = torch.zeros(f_size, dtype=self.dtype, device=self.device) - result = RAvg() + type_fval = f_values.dtype if f_size == 1 else f_values[0].dtype + epoch = self.neval // self.nbatch + mean_values = torch.zeros((f_size, epoch), dtype=type_fval, device=self.device) + std_values = torch.zeros_like(mean_values) - mean[:] = 0 - var[:] = 0 - for _ in range(epoch): + for iepoch in range(epoch): x, log_detJ = self.sample(self.nbatch) f_values = f(x) batch_results = self._multiply_by_jacobian(f_values, torch.exp(log_detJ)) - mean += torch.mean(batch_results, dim=-1) / epoch - var += torch.var(batch_results, dim=-1) / (self.neval * epoch) - - result.sum_neval += self.neval - result.add(gvar.gvar(mean.item(), (var**0.5).item())) - return result + mean_values[:, iepoch] = torch.mean(batch_results, dim=-1) + std_values[:, iepoch] = torch.std(batch_results, dim=-1) / self.nbatch**0.5 + + results = np.array([RAvg() for _ in range(f_size)]) + for iepoch in range(epoch): + for j in range(f_size): + results[j].sum_neval += self.nbatch + results[j].add( + gvar.gvar( + mean_values[j, iepoch].item(), std_values[j, iepoch].item() + ) + ) + if f_size == 1: + return results[0] + else: + return results def _multiply_by_jacobian(self, values, jac): # if isinstance(values, dict): @@ -150,52 +155,50 @@ def __call__( f: Callable, proposal_dist: Callable = uniform, thinning=1, - mix_rate=0.0, + mix_rate=0.5, **kwargs, ): epsilon = 1e-16 # Small value to ensure numerical stability + epoch = self.neval // self.nbatch # vars_shape = (self.nbatch, self.dim) current_y, current_jac = self.q0.sample(self.nbatch) - # if self.maps: current_x, detJ = self.maps.forward(current_y) current_jac += detJ - # else: - # current_x = current_y current_jac = torch.exp(current_jac) current_fval = f(current_x) - current_weight = mix_rate / current_jac + (1 - mix_rate) * current_fval.abs() - current_weight.masked_fill_(current_weight < epsilon, epsilon) - # current_fval.masked_fill_(current_fval.abs() < epsilon, epsilon) + f_size = len(current_fval) if isinstance(current_fval, (list, tuple)) else 1 - # proposed_y = torch.empty_like(current_y) - # proposed_x = torch.empty_like(current_x) - # new_fval = torch.empty_like(current_fval) - # new_weight = torch.empty_like(current_weight) + if f_size > 1: + current_fval = sum(current_fval) - f_size = len(current_fval) if isinstance(current_fval, (list, tuple)) else 1 - # type_fval = current_fval.dtype if f_size == 1 else type(current_fval[0].dtype) - # mean = torch.zeros(f_size, dtype=type_fval, device=self.device) - mean = torch.zeros(f_size, dtype=self.dtype, device=self.device) - mean_ref = torch.zeros_like(mean) - # var = torch.zeros(f_size, dtype=type_fval, device=self.device) - var = torch.zeros(f_size, dtype=self.dtype, device=self.device) - var_ref = torch.zeros_like(mean) + def _integrand(x): + return sum(f(x)) + else: - result = RAvg() - result_ref = RAvg() + def _integrand(x): + return f(x) - epoch = self.neval // self.nbatch - n_meas = 0 + type_fval = current_fval.dtype - def _propose(current_y, current_fval, current_weight, current_jac): + current_weight = mix_rate / current_jac + (1 - mix_rate) * current_fval.abs() + current_weight.masked_fill_(current_weight < epsilon, epsilon) + + n_meas = epoch // thinning + mean_values = torch.zeros((f_size, n_meas), dtype=type_fval, device=self.device) + std_values = torch.zeros_like(mean_values) + mean_refvalues = torch.zeros(n_meas, dtype=type_fval, device=self.device) + std_refvalues = torch.zeros_like(mean_refvalues) + + def _propose(current_y, current_x, current_weight, current_jac): proposed_y = proposal_dist( self.dim, self.bounds, self.device, self.dtype, current_y, **kwargs ) proposed_x, new_jac = self.maps.forward(proposed_y) new_jac = torch.exp(new_jac) - new_fval = f(proposed_x) + new_fval = _integrand(proposed_x) new_weight = mix_rate / new_jac + (1 - mix_rate) * new_fval.abs() + new_weight.masked_fill_(new_weight < epsilon, epsilon) acceptance_probs = new_weight / current_weight * new_jac / current_jac @@ -205,55 +208,49 @@ def _propose(current_y, current_fval, current_weight, current_jac): ) current_y = torch.where(accept.unsqueeze(1), proposed_y, current_y) - current_fval = torch.where(accept, new_fval, current_fval) + # current_fval = torch.where(accept, new_fval, current_fval) + current_x = torch.where(accept.unsqueeze(1), proposed_x, current_x) current_weight = torch.where(accept, new_weight, current_weight) current_jac = torch.where(accept, new_jac, current_jac) - return current_y, current_fval, current_weight, current_jac + return current_y, current_x, current_weight, current_jac for i in range(self.nburnin): - current_y, current_fval, current_weight, current_jac = _propose( - current_y, current_fval, current_weight, current_jac + current_y, current_x, current_weight, current_jac = _propose( + current_y, current_x, current_weight, current_jac ) - for i in range(epoch // thinning): + + for imeas in range(n_meas): for j in range(thinning): - current_y, current_fval, current_weight, current_jac = _propose( - current_y, current_fval, current_weight, current_jac + current_y, current_x, current_weight, current_jac = _propose( + current_y, current_x, current_weight, current_jac ) - n_meas += 1 - batch_results = current_fval / current_weight - - mean += torch.mean(batch_results, dim=-1) / epoch - var += torch.var(batch_results, dim=-1) / epoch + batch_results = self._multiply_by_jacobian( + f(current_x), 1.0 / current_weight + ) batch_results_ref = 1 / (current_jac * current_weight) - mean_ref += torch.mean(batch_results_ref, dim=-1) / epoch - var_ref += torch.var(batch_results_ref, dim=-1) / epoch - - result.sum_neval += self.neval - result.add(gvar.gvar(mean.item(), ((var / n_meas) ** 0.5).item())) - result_ref.sum_neval += self.nbatch - result_ref.add(gvar.gvar(mean_ref.item(), ((var_ref / n_meas) ** 0.5).item())) - - return result / result_ref * self._rangebounds.prod() - - # def _propose(self, u, proposal_dist, **kwargs): - # if proposal_dist == "random_walk": - # step_size = kwargs.get("step_size", 0.2) - # step_sizes = self._rangebounds * step_size - # step = ( - # torch.empty(self.dim, device=self.device).uniform_(-1, 1) * step_sizes - # ) - # new_u = (u + step - self.bounds[:, 0]) % self._rangebounds + self.bounds[ - # :, 0 - # ] - # return new_u - # # return (u + (torch.rand_like(u) - 0.5) * step_size) % 1.0 - # elif proposal_dist == "uniform": - # # return torch.rand_like(u) - # return torch.rand_like(u) * self._rangebounds + self.bounds[:, 0] - # # elif proposal_dist == "gaussian": - # # mean = kwargs.get("mean", torch.zeros_like(u)) - # # std = kwargs.get("std", torch.ones_like(u)) - # # return torch.normal(mean, std) - # else: - # raise ValueError(f"Unknown proposal distribution: {proposal_dist}") + + mean_values[:, imeas] = torch.mean(batch_results, dim=-1) + std_values[:, imeas] = torch.std(batch_results, dim=-1) / self.nbatch**0.5 + + mean_refvalues[imeas] = torch.mean(batch_results_ref, dim=-1) + std_refvalues[imeas] = ( + torch.std(batch_results_ref, dim=-1) / self.nbatch**0.5 + ) + + results = np.array([RAvg() for _ in range(f_size)]) + results_ref = RAvg() + for imeas in range(n_meas): + results_ref.sum_neval += self.nbatch + results_ref.add( + gvar.gvar(mean_refvalues[imeas].item(), std_refvalues[imeas].item()) + ) + for j in range(f_size): + results[j].sum_neval += self.nbatch + results[j].add( + gvar.gvar(mean_values[j, imeas].item(), std_values[j, imeas].item()) + ) + if f_size == 1: + return results[0] / results_ref * self._rangebounds.prod() + else: + return results / results_ref * self._rangebounds.prod().item() diff --git a/src/mc_test.py b/src/mc_test.py index d5f8c8f..f22d071 100644 --- a/src/mc_test.py +++ b/src/mc_test.py @@ -9,21 +9,11 @@ def test_nothing(): - pass + pass + def unit_circle_integrand(x): inside_circle = (x[:, 0] ** 2 + x[:, 1] ** 2 < 1).double() - # return { - # "scalar": inside_circle, - # "vector": torch.stack([inside_circle, 2 * inside_circle], dim=1), - # "matrix": torch.stack( - # [ - # inside_circle.unsqueeze(1).repeat(1, 3), - # (2 * inside_circle).unsqueeze(1).repeat(1, 3), - # ], - # dim=1, - # ), - # } return inside_circle @@ -31,8 +21,21 @@ def half_sphere_integrand(x): return torch.clamp(1 - (x[:, 0] ** 2 + x[:, 1] ** 2), min=0) * 2 +def integrand_list(x): + return [unit_circle_integrand(x), half_sphere_integrand(x) / 2] + + +def integrand_list1(x): + dx2 = torch.zeros(x.shape[0], dtype=x.dtype, device=x.device) + for d in range(4): + dx2 += (x[:, d] - 0.5) ** 2 + f = torch.exp(-200 * dx2) + return [f, f * x[:, 0], f * x[:, 0] ** 2] + + dim = 2 bounds = [(-1, 1), (-1, 1)] +n_eval = 400000 affine_map = Linear(bounds, device=device) # vegas_map = Vegas(bounds, device=device) @@ -41,9 +44,9 @@ def half_sphere_integrand(x): print("Calculate the area of the unit circle using Monte Carlo integration...") mc_integrator = MonteCarlo( - # bounds, maps=affine_map, neval=400000, nbatch=1000, device=device + # bounds, maps=affine_map, neval=n_eval, nbatch=1000, device=device bounds, - neval=400000, + neval=n_eval, nbatch=1000, device=device, ) @@ -52,14 +55,14 @@ def half_sphere_integrand(x): print(f" Integral: {res.mean}") print(f" Error: {res.sdev}") -res = MonteCarlo(bounds, neval=400000, nbatch=1000, device=device)( +res = MonteCarlo(bounds, neval=n_eval, nbatch=1000, device=device)( unit_circle_integrand ) print("Plain MC Integral results:") print(f" Integral: {res.mean}") print(f" Error: {res.sdev}") -mcmc_integrator = MCMC(bounds, neval=400000, nbatch=1000, nburnin=100, device=device) +mcmc_integrator = MCMC(bounds, neval=n_eval, nbatch=1000, nburnin=100, device=device) res = mcmc_integrator(unit_circle_integrand, mix_rate=0.5) print("MCMC Integral results:") print(f" Integral: {res.mean}") @@ -73,8 +76,35 @@ def half_sphere_integrand(x): print(f" Integral: {res.mean}") print(f" Error: {res.sdev}") -mcmc_integrator = MCMC(bounds, neval=400000, nbatch=1000, nburnin=100, device=device) +mcmc_integrator = MCMC(bounds, neval=n_eval, nbatch=1000, nburnin=100, device=device) res = mcmc_integrator(half_sphere_integrand, mix_rate=0.5) print("MCMC Integral results:") print(f" Integral: {res.mean}") print(f" Error: {res.sdev}") + +# Two integrands +res = mc_integrator(integrand_list) +print("Plain MC Integral results:") +print(f" Integral 1: {res[0].mean} +- {res[0].sdev}") +print(f" Integral 2: {res[1].mean} +- {res[1].sdev}") + +res = mcmc_integrator(integrand_list, mix_rate=0.5) +print("MCMC Integral results:") +print(f" Integral 1: {res[0].mean} +- {res[0].sdev}") +print(f" Integral 2: {res[1].mean} +- {res[1].sdev}") + +bounds = [(0, 1)] * 4 +mc_integrator = MonteCarlo( + bounds, + neval=n_eval, + nbatch=1000, + device=device, +) +mcmc_integrator = MCMC(bounds, neval=n_eval, nbatch=1000, nburnin=100, device=device) +res = mc_integrator(integrand_list1) +print("I[0] =", res[0], " I[1] =", res[1], " I[2] =", res[2]) +print(" =", res[1] / res[0]) + +res = mcmc_integrator(integrand_list1, mix_rate=0.5) +print("I[0] =", res[0], " I[1] =", res[1], " I[2] =", res[2]) +print(" =", res[1] / res[0]) From 3b39fa9d87e20752fa42c2c8fe56e3515067b3b4 Mon Sep 17 00:00:00 2001 From: houpc Date: Wed, 23 Oct 2024 17:21:16 +0800 Subject: [PATCH 08/10] udpate vegas map and support vegas-mc (simple test pass) --- examples/benchmark_vegas.py | 93 +++++++++++++++++ src/base.py | 3 + src/integrators.py | 8 +- src/maps.py | 56 +++++++--- src/mc_test.py | 202 ++++++++++++++++++++++++++++-------- src/vegas_test.py | 81 +++++++++++++++ 6 files changed, 382 insertions(+), 61 deletions(-) create mode 100644 examples/benchmark_vegas.py create mode 100644 src/vegas_test.py diff --git a/examples/benchmark_vegas.py b/examples/benchmark_vegas.py new file mode 100644 index 0000000..d830c5d --- /dev/null +++ b/examples/benchmark_vegas.py @@ -0,0 +1,93 @@ +import vegas +import numpy as np +import gvar +import torch + +dim = 4 +nitn = 10 +ninc = 1000 + + +@vegas.lbatchintegrand +def f_batch(x): + dx2 = 0.0 + for d in range(dim): + dx2 += (x[:, d] - 0.5) ** 2 + return np.exp(-200 * dx2) + # ans = np.empty((x.shape[0], 3), float) + # dx2 = 0.0 + # for d in range(dim): + # dx2 += (x[:, d] - 0.5) ** 2 + # ans[:, 0] = np.exp(-200 * dx2) + # ans[:, 1] = x[:, 0] * ans[:, 0] + # ans[:, 2] = x[:, 0] ** 2 * ans[:, 0] + # return ans + + +def smc(f, map, neval, dim): + "integrates f(y) over dim-dimensional unit hypercube" + y = np.random.uniform(0, 1, (neval, dim)) + jac = np.empty(y.shape[0], float) + x = np.empty(y.shape, float) + map.map(y, x, jac) + fy = jac * f(x) + return (np.average(fy), np.std(fy) / neval**0.5) + + +def mc(f, neval, dim): + "integrates f(y) over dim-dimensional unit hypercube" + y = np.random.uniform(0, 1, (neval, dim)) + fy = f(y) + return (np.average(fy), np.std(fy) / neval**0.5) + + +m = vegas.AdaptiveMap(dim * [[0, 1]], ninc=ninc) +ny = 20000 +# torch.manual_seed(0) +# y = torch.rand((ny, dim), dtype=torch.float64).numpy() +y = np.random.uniform(0.0, 1.0, (ny, dim)) # 1000 random y's + +x = np.empty(y.shape, float) # work space +jac = np.empty(y.shape[0], float) +f2 = np.empty(y.shape[0], float) + +for itn in range(10): # 5 iterations to adapt + m.map(y, x, jac) # compute x's and jac + + f2 = (jac * f_batch(x)) ** 2 + m.add_training_data(y, f2) # adapt + # if itn == 0: + # print(np.array(memoryview(m.sum_f))) + # print(np.array(memoryview(m.n_f))) + m.adapt(alpha=0.5) + + +# with map +r = smc(f_batch, m, 50_000, dim) +print(" SMC + map:", f"{r[0]} +- {r[1]}") + +# without map +r = mc(f_batch, 50_000, dim) +print("SMC (no map):", f"{r[0]} +- {r[1]}") + + +# vegas with adaptive stratified sampling +print("VEGAS using adaptive stratified sampling") +integ = vegas.Integrator(dim * [[0, 1]]) +training = integ(f_batch, nitn=10, neval=20000) # adapt grid + +# final analysis +result = integ(f_batch, nitn=1, neval=50_000, adapt=False) +print(result) +result = integ(f_batch, nitn=5, neval=10_000, adapt=False) +print(result) +result = integ(f_batch, nitn=5, neval=10_000) +print(result) +# print("I[0] =", result[0], " I[1] =", result[1], " I[2] =", result[2]) +# print("Q = %.2f\n" % result.Q) +# print(" =", result[1] / result[0]) +# print( +# "sigma_x**2 = - **2 =", +# result[2] / result[0] - (result[1] / result[0]) ** 2, +# ) +# print("\ncorrelation matrix:\n", gv.evalcorr(result)) diff --git a/src/base.py b/src/base.py index ab7ffd0..07055b1 100644 --- a/src/base.py +++ b/src/base.py @@ -15,6 +15,8 @@ def __init__(self, bounds, device="cpu", dtype=torch.float64): # self.bounds = bounds if isinstance(bounds, (list, np.ndarray)): self.bounds = torch.tensor(bounds, dtype=dtype, device=device) + elif isinstance(bounds, torch.Tensor): + self.bounds = bounds else: raise ValueError("Unsupported map specification") self.dim = self.bounds.shape[0] @@ -42,6 +44,7 @@ def __init__(self, bounds, device="cpu", dtype=torch.float64): self._rangebounds = self.bounds[:, 1] - self.bounds[:, 0] def sample(self, nsamples=1, **kwargs): + # torch.manual_seed(0) # test seed u = ( torch.rand((nsamples, self.dim), device=self.device, dtype=self.dtype) * self._rangebounds diff --git a/src/integrators.py b/src/integrators.py index a97910b..0f3e638 100644 --- a/src/integrators.py +++ b/src/integrators.py @@ -72,7 +72,7 @@ def __init__( super().__init__(bounds, q0, maps, neval, nbatch, device, dtype) def __call__(self, f: Callable, **kwargs): - x, _ = self.sample(self.nbatch) + x, _ = self.sample(1) f_values = f(x) f_size = len(f_values) if isinstance(f_values, (list, tuple)) else 1 type_fval = f_values.dtype if f_size == 1 else f_values[0].dtype @@ -189,7 +189,7 @@ def _integrand(x): mean_refvalues = torch.zeros(n_meas, dtype=type_fval, device=self.device) std_refvalues = torch.zeros_like(mean_refvalues) - def _propose(current_y, current_x, current_weight, current_jac): + def one_step(current_y, current_x, current_weight, current_jac): proposed_y = proposal_dist( self.dim, self.bounds, self.device, self.dtype, current_y, **kwargs ) @@ -215,13 +215,13 @@ def _propose(current_y, current_x, current_weight, current_jac): return current_y, current_x, current_weight, current_jac for i in range(self.nburnin): - current_y, current_x, current_weight, current_jac = _propose( + current_y, current_x, current_weight, current_jac = one_step( current_y, current_x, current_weight, current_jac ) for imeas in range(n_meas): for j in range(thinning): - current_y, current_x, current_weight, current_jac = _propose( + current_y, current_x, current_weight, current_jac = one_step( current_y, current_x, current_weight, current_jac ) diff --git a/src/maps.py b/src/maps.py index 2b6a6df..93081b0 100644 --- a/src/maps.py +++ b/src/maps.py @@ -1,6 +1,10 @@ -import torch import numpy as np +import torch from torch import nn +from base import Uniform +import sys + +TINY = 10 ** (sys.float_info.min_10_exp + 50) class Map(nn.Module): @@ -75,6 +79,9 @@ def __init__(self, bounds, ninc=1000, alpha=0.5, device="cpu", dtype=torch.float self.dim, self.ninc.max() + 1, dtype=self.dtype, device=self.device ) + self._A = self.bounds[:, 1] - self.bounds[:, 0] + self._jaclinear = torch.prod(self._A) + for d in range(self.dim): self.grid[d, : self.ninc[d] + 1] = torch.linspace( self.bounds[d, 0], @@ -88,7 +95,28 @@ def __init__(self, bounds, ninc=1000, alpha=0.5, device="cpu", dtype=torch.float ) self.clear() - def add_training_data(self, u, f): + def train(self, nsamples, f, nitn=5, alpha=0.5): + q0 = Uniform(self.bounds, device=self.device, dtype=self.dtype) + u, log_detJ0 = q0.sample(nsamples) + + fval = f(u) + f_size = len(fval) if isinstance(fval, (list, tuple)) else 1 + if f_size > 1: + + def _integrand(x): + return sum(f(x)) + else: + + def _integrand(x): + return f(x) + + for _ in range(nitn): + x, log_detJ = self.forward(u) + f2 = torch.exp(2 * (log_detJ + log_detJ0)) * _integrand(x) ** 2 + self.add_training_data(u, f2) + self.adapt(alpha) + + def add_training_data(self, u, fval): """Add training data ``f`` for ``u``-space points ``u``. Accumulates training data for later use by ``self.adapt()``. @@ -106,11 +134,13 @@ def add_training_data(self, u, f): """ if self.sum_f is None: self.sum_f = torch.zeros_like(self.inc) - self.n_f = torch.zeros_like(self.inc) + 1e-10 - iu = torch.floor(u * self.ninc).long() + self.n_f = torch.zeros_like(self.inc) + TINY + iu = (u - self.bounds[:, 0]) / self._A * self.ninc + iu = torch.floor(iu).long() for d in range(self.dim): - self.sum_f[d, iu[:, d]] += torch.abs(f) - self.n_f[d, iu[:, d]] += 1 + indices = iu[:, d] + self.sum_f[d].scatter_add_(0, indices, fval.abs()) + self.n_f[d].scatter_add_(0, indices, torch.ones_like(fval)) def adapt(self, alpha=0.0): """Adapt grid to accumulated training data. @@ -172,9 +202,9 @@ def adapt(self, alpha=0.0): ) sum_f = torch.sum(tmp_f[:ninc]) if sum_f > 0: - avg_f[:ninc] = tmp_f[:ninc] / sum_f + 1e-10 + avg_f[:ninc] = tmp_f[:ninc] / sum_f + TINY else: - avg_f[:ninc] = 1e-10 + avg_f[:ninc] = TINY avg_f[:ninc] = ( -(1 - avg_f[:ninc]) / torch.log(avg_f[:ninc]) ) ** alpha @@ -226,8 +256,10 @@ def clear(self): def forward(self, u): u = u.to(self.device) u_ninc = u * self.ninc - iu = torch.floor(u_ninc).long() - du_ninc = u_ninc - iu + # iu = torch.floor(u_ninc).long() + iu = (u - self.bounds[:, 0]) / self._A * self.ninc + iu = torch.floor(iu).long() + du_ninc = u_ninc - torch.floor(u_ninc).long() x = torch.empty_like(u) jac = torch.ones(u.shape[0], device=x.device) @@ -249,7 +281,7 @@ def forward(self, u): x[mask_inv, d] = self.grid[d, ninc] jac[mask_inv] *= self.inc[d, ninc - 1] * ninc - return x, torch.log(jac) + return x, torch.log(jac / self._jaclinear) @torch.no_grad() def inverse(self, x): @@ -285,7 +317,7 @@ def inverse(self, x): u[mask_upper, d] = 1.0 jac[mask_upper] *= self.inc[d, ninc - 1] * ninc - return u, torch.log(jac) + return u, torch.log(jac / self._jaclinear) # class NormalizingFlow(Map): diff --git a/src/mc_test.py b/src/mc_test.py index f22d071..ba74199 100644 --- a/src/mc_test.py +++ b/src/mc_test.py @@ -4,12 +4,8 @@ from utils import set_seed, get_device set_seed(42) -device = get_device() -# device = torch.device("cpu") - - -def test_nothing(): - pass +# device = get_device() +device = torch.device("cpu") def unit_circle_integrand(x): @@ -33,78 +29,194 @@ def integrand_list1(x): return [f, f * x[:, 0], f * x[:, 0] ** 2] +def sharp_peak(x): + dx2 = torch.zeros(x.shape[0], dtype=x.dtype, device=x.device) + for d in range(4): + dx2 += (x[:, d] - 0.5) ** 2 + return torch.exp(-200 * dx2) + + dim = 2 bounds = [(-1, 1), (-1, 1)] n_eval = 400000 +n_batch = 10000 +n_therm = 10 -affine_map = Linear(bounds, device=device) -# vegas_map = Vegas(bounds, device=device) +vegas_map = Vegas(bounds, device=device, ninc=10) -# Monte Carlo integration -print("Calculate the area of the unit circle using Monte Carlo integration...") +# True value of pi +print(f"pi = {torch.pi} \n") + +# Start Monte Carlo integration, including plain-MC, MCMC, vegas, and vegas-MCMC mc_integrator = MonteCarlo( - # bounds, maps=affine_map, neval=n_eval, nbatch=1000, device=device bounds, neval=n_eval, - nbatch=1000, + nbatch=n_batch, + # nbatch=n_eval, device=device, ) +mcmc_integrator = MCMC( + bounds, neval=n_eval, nbatch=n_batch, nburnin=n_therm, device=device +) + +print("Calculate the area of the unit circle f(x1, x2) in the bounds [-1, 1]^2...") res = mc_integrator(unit_circle_integrand) -print("Plain MC Integral results:") -print(f" Integral: {res.mean}") -print(f" Error: {res.sdev}") +print("Plain MC Integral results: ", res) + +res = mcmc_integrator(unit_circle_integrand, mix_rate=0.5) +print("MCMC Integral results: ", res) -res = MonteCarlo(bounds, neval=n_eval, nbatch=1000, device=device)( - unit_circle_integrand +vegas_map.train(20000, unit_circle_integrand, alpha=0.5) +vegas_integrator = MonteCarlo( + bounds, + maps=vegas_map, + neval=n_eval, + nbatch=n_batch, + # nbatch=n_eval, + device=device, ) -print("Plain MC Integral results:") -print(f" Integral: {res.mean}") -print(f" Error: {res.sdev}") +res = vegas_integrator(unit_circle_integrand) +print("VEGAS Integral results: ", res) -mcmc_integrator = MCMC(bounds, neval=n_eval, nbatch=1000, nburnin=100, device=device) -res = mcmc_integrator(unit_circle_integrand, mix_rate=0.5) -print("MCMC Integral results:") -print(f" Integral: {res.mean}") -print(f" Error: {res.sdev}") +vegasmcmc_integrator = MCMC( + bounds, + maps=vegas_map, + neval=n_eval, + nbatch=n_batch, + nburnin=n_therm, + device=device, +) +res = vegasmcmc_integrator(unit_circle_integrand, mix_rate=0.5) +print("VEGAS-MCMC Integral results: ", res, "\n") -# True value of pi -print(f"True value of pi: {torch.pi:.6f}") + +print( + r"Calculate the integral g(x1, x2) = $2 \max(1-(x_1^2+x_2^2), 0)$ in the bounds [-1, 1]^2..." +) res = mc_integrator(half_sphere_integrand) -print("Plain MC Integral results:") -print(f" Integral: {res.mean}") -print(f" Error: {res.sdev}") +print("Plain MC Integral results: ", res) -mcmc_integrator = MCMC(bounds, neval=n_eval, nbatch=1000, nburnin=100, device=device) res = mcmc_integrator(half_sphere_integrand, mix_rate=0.5) -print("MCMC Integral results:") -print(f" Integral: {res.mean}") -print(f" Error: {res.sdev}") +print("MCMC Integral results:", res) + +# train the vegas map +vegas_map.train(20000, half_sphere_integrand, nitn=10, alpha=0.5) + +res = vegas_integrator(half_sphere_integrand) +print("VEGAS Integral results: ", res) + +res = vegasmcmc_integrator(half_sphere_integrand, mix_rate=0.5) +print("VEGAS-MCMC Integral results: ", res) + +print("\nCalculate the integral [f(x1, x2), g(x1, x2)/2] in the bounds [-1, 1]^2") # Two integrands res = mc_integrator(integrand_list) print("Plain MC Integral results:") -print(f" Integral 1: {res[0].mean} +- {res[0].sdev}") -print(f" Integral 2: {res[1].mean} +- {res[1].sdev}") +print(" Integral 1: ", res[0]) +print(" Integral 2: ", res[1]) res = mcmc_integrator(integrand_list, mix_rate=0.5) print("MCMC Integral results:") -print(f" Integral 1: {res[0].mean} +- {res[0].sdev}") -print(f" Integral 2: {res[1].mean} +- {res[1].sdev}") +print(f" Integral 1: ", res[0]) +print(f" Integral 2: ", res[1]) + +# print("VEAGS map is trained for g(x1, x2)") +vegas_map.train(20000, integrand_list, nitn=10, alpha=0.5) +res = vegas_integrator(integrand_list) +print("VEGAS Integral results:") +print(" Integral 1: ", res[0]) +print(" Integral 2: ", res[1]) + +res = vegasmcmc_integrator(integrand_list, mix_rate=0.5) +print("VEGAS-MCMC Integral results:") +print(" Integral 1: ", res[0]) +print(" Integral 2: ", res[1]) + +print("\nCalculate the integral [h(X), x1 * h(X), x1^2 * h(X)] in the bounds [0, 1]^4") +print("h(X) = exp(-200 * (x1^2 + x2^2 + x3^2 + x4^2))") bounds = [(0, 1)] * 4 mc_integrator = MonteCarlo( bounds, neval=n_eval, - nbatch=1000, + nbatch=n_batch, + # nbatch=n_eval, device=device, ) -mcmc_integrator = MCMC(bounds, neval=n_eval, nbatch=1000, nburnin=100, device=device) +mcmc_integrator = MCMC( + bounds, neval=n_eval, nbatch=n_batch, nburnin=n_therm, device=device +) +print("Plain MC Integral results:") res = mc_integrator(integrand_list1) -print("I[0] =", res[0], " I[1] =", res[1], " I[2] =", res[2]) -print(" =", res[1] / res[0]) - +print( + " I[0] =", + res[0], + " I[1] =", + res[1], + " I[2] =", + res[2], + " I[1]/I[0] =", + res[1] / res[0], +) +print("MCMC Integral results:") res = mcmc_integrator(integrand_list1, mix_rate=0.5) -print("I[0] =", res[0], " I[1] =", res[1], " I[2] =", res[2]) -print(" =", res[1] / res[0]) +print( + " I[0] =", + res[0], + " I[1] =", + res[1], + " I[2] =", + res[2], + " I[1]/I[0] =", + res[1] / res[0], +) + +vegas_map = Vegas(bounds, device=device) +print("train VEGAS map for h(X)...") +# vegas_map.train(20000, sharp_peak, nitn=10, alpha=0.5) +vegas_map.train(20000, integrand_list1, nitn=10, alpha=0.5) + +print("VEGAS Integral results:") +vegas_integrator = MonteCarlo( + bounds, + maps=vegas_map, + neval=n_eval, + nbatch=n_batch, + # nbatch=n_eval, + device=device, +) +res = vegas_integrator(integrand_list1) +print( + " I[0] =", + res[0], + " I[1] =", + res[1], + " I[2] =", + res[2], + " I[1]/I[0] =", + res[1] / res[0], +) + +print("VEGAS-MCMC Integral results:") +vegasmcmc_integrator = MCMC( + bounds, + maps=vegas_map, + neval=n_eval, + nbatch=n_batch, + nburnin=n_therm, + device=device, +) +res = vegasmcmc_integrator(integrand_list1, mix_rate=0.5) +print( + " I[0] =", + res[0], + " I[1] =", + res[1], + " I[2] =", + res[2], + " I[1]/I[0] =", + res[1] / res[0], +) diff --git a/src/vegas_test.py b/src/vegas_test.py new file mode 100644 index 0000000..a0b7c5e --- /dev/null +++ b/src/vegas_test.py @@ -0,0 +1,81 @@ +import torch +from integrators import MonteCarlo, MCMC +from maps import Vegas, Linear +from utils import set_seed, get_device + +set_seed(42) +# device = get_device() +device = torch.device("cpu") + + +def integrand_list1(x): + dx2 = torch.zeros(x.shape[0], dtype=x.dtype, device=x.device) + for d in range(4): + dx2 += (x[:, d] - 0.5) ** 2 + f = torch.exp(-200 * dx2) + return [f, f * x[:, 0], f * x[:, 0] ** 2] + + +def sharp_peak(x): + dx2 = torch.zeros(x.shape[0], dtype=x.dtype, device=x.device) + for d in range(4): + dx2 += (x[:, d] - 0.5) ** 2 + return torch.exp(-200 * dx2) + + +ninc = 1000 +# n_eval = 100000 +n_eval = 50000 +n_batch = 10000 +n_therm = 10 +# Start Monte Carlo integration, including plain-MC, MCMC, vegas, and vegas-MCMC +print("\nCalculate the integral [h(X), x1 * h(X), x1^2 * h(X)] in the bounds [0, 1]^4") +print("h(X) = exp(-200 * (x1^2 + x2^2 + x3^2 + x4^2))") + +bounds = [(0, 1)] * 4 +vegas_map = Vegas(bounds, device=device, ninc=ninc) +print("train VEGAS map for h(X)...") +vegas_map.train(20000, sharp_peak, nitn=10, alpha=0.5) +# print(vegas_map.extract_grid()) + +print("VEGAS Integral results:") +vegas_integrator = MonteCarlo( + bounds, + maps=vegas_map, + neval=n_eval, + # nbatch=n_batch, + nbatch=n_eval, + device=device, +) +res = vegas_integrator(integrand_list1) +print( + " I[0] =", + res[0], + " I[1] =", + res[1], + " I[2] =", + res[2], + " I[1]/I[0] =", + res[1] / res[0], +) + +print("VEGAS-MCMC Integral results:") +vegasmcmc_integrator = MCMC( + bounds, + maps=vegas_map, + neval=n_eval, + nbatch=n_batch, + nburnin=n_therm, + device=device, +) +res = vegasmcmc_integrator(integrand_list1, mix_rate=0.5) +print( + " I[0] =", + res[0], + " I[1] =", + res[1], + " I[2] =", + res[2], + " I[1]/I[0] =", + res[1] / res[0], +) From 686c3b762bc8b432845f4216a9b074b905970e64 Mon Sep 17 00:00:00 2001 From: houpc Date: Wed, 23 Oct 2024 22:34:11 +0800 Subject: [PATCH 09/10] remove bounds dep in Integrator and avoid average for each MC step --- src/integrators.py | 121 +++++++++++++++++++++++---------------------- src/maps.py | 4 +- src/mc_test.py | 21 +++----- src/vegas_test.py | 37 ++++++++++++-- 4 files changed, 105 insertions(+), 78 deletions(-) diff --git a/src/integrators.py b/src/integrators.py index 0f3e638..4b93ad9 100644 --- a/src/integrators.py +++ b/src/integrators.py @@ -16,25 +16,28 @@ class Integrator: def __init__( self, - bounds, - q0=None, maps=None, + bounds=None, + q0=None, neval: int = 1000, nbatch: int = None, device="cpu", dtype=torch.float64, ): - if not isinstance(bounds, (list, np.ndarray)): - raise TypeError("bounds must be a list or a NumPy array.") self.dtype = dtype - self.dim = len(bounds) - if not q0: - q0 = Uniform(bounds, device=device, dtype=dtype) - self.bounds = torch.tensor(bounds, dtype=dtype, device=device) - self.q0 = q0 if maps: if not self.dtype == maps.dtype: raise ValueError("Float type of maps should be same as integrator.") + self.bounds = maps.bounds + else: + if not isinstance(bounds, (list, np.ndarray)): + raise TypeError("bounds must be a list or a NumPy array.") + self.bounds = torch.tensor(bounds, dtype=dtype, device=device) + + self.dim = len(self.bounds) + if not q0: + q0 = Uniform(self.bounds, device=device, dtype=dtype) + self.q0 = q0 self.maps = maps self.neval = neval if nbatch is None: @@ -61,43 +64,47 @@ def sample(self, nsample, **kwargs): class MonteCarlo(Integrator): def __init__( self, - bounds, - q0=None, maps=None, + bounds=None, + q0=None, neval: int = 1000, nbatch: int = None, device="cpu", dtype=torch.float64, ): - super().__init__(bounds, q0, maps, neval, nbatch, device, dtype) + super().__init__(maps, bounds, q0, neval, nbatch, device, dtype) def __call__(self, f: Callable, **kwargs): x, _ = self.sample(1) f_values = f(x) - f_size = len(f_values) if isinstance(f_values, (list, tuple)) else 1 - type_fval = f_values.dtype if f_size == 1 else f_values[0].dtype + if isinstance(f_values, (list, tuple)) and isinstance( + f_values[0], torch.Tensor + ): + f_size = len(f_values) + type_fval = f_values[0].dtype + elif isinstance(f_values, torch.Tensor): + f_size = 1 + type_fval = f_values.dtype + else: + raise TypeError( + "f must return a torch.Tensor or a list/tuple of torch.Tensor." + ) epoch = self.neval // self.nbatch - mean_values = torch.zeros((f_size, epoch), dtype=type_fval, device=self.device) - std_values = torch.zeros_like(mean_values) + values = torch.zeros((self.nbatch, f_size), dtype=type_fval, device=self.device) for iepoch in range(epoch): x, log_detJ = self.sample(self.nbatch) f_values = f(x) batch_results = self._multiply_by_jacobian(f_values, torch.exp(log_detJ)) - mean_values[:, iepoch] = torch.mean(batch_results, dim=-1) - std_values[:, iepoch] = torch.std(batch_results, dim=-1) / self.nbatch**0.5 + values += batch_results / epoch results = np.array([RAvg() for _ in range(f_size)]) - for iepoch in range(epoch): - for j in range(f_size): - results[j].sum_neval += self.nbatch - results[j].add( - gvar.gvar( - mean_values[j, iepoch].item(), std_values[j, iepoch].item() - ) - ) + for i in range(f_size): + _mean = values[:, i].mean().item() + _std = values[:, i].std().item() / self.nbatch**0.5 + results[i].add(gvar.gvar(_mean, _std)) if f_size == 1: return results[0] else: @@ -107,9 +114,9 @@ def _multiply_by_jacobian(self, values, jac): # if isinstance(values, dict): # return {k: v * torch.exp(log_det_J) for k, v in values.items()} if isinstance(values, (list, tuple)): - return torch.stack([v * jac for v in values]) + return torch.stack([v * jac for v in values], dim=-1) else: - return values * jac + return torch.stack([values * jac], dim=-1) def random_walk(dim, bounds, device, dtype, u, **kwargs): @@ -135,16 +142,16 @@ def gaussian(dim, bounds, device, dtype, u, **kwargs): class MCMC(MonteCarlo): def __init__( self, - bounds, - q0=None, maps=None, + bounds=None, + q0=None, neval=10000, nbatch=None, nburnin=500, device="cpu", dtype=torch.float64, ): - super().__init__(bounds, q0, maps, neval, nbatch, device, dtype) + super().__init__(maps, bounds, q0, neval, nbatch, device, dtype) self.nburnin = nburnin if maps is None: self.maps = Linear([(0, 1)] * self.dim, device=device) @@ -160,34 +167,34 @@ def __call__( ): epsilon = 1e-16 # Small value to ensure numerical stability epoch = self.neval // self.nbatch - # vars_shape = (self.nbatch, self.dim) current_y, current_jac = self.q0.sample(self.nbatch) current_x, detJ = self.maps.forward(current_y) current_jac += detJ current_jac = torch.exp(current_jac) current_fval = f(current_x) - f_size = len(current_fval) if isinstance(current_fval, (list, tuple)) else 1 - - if f_size > 1: + if isinstance(current_fval, (list, tuple)) and isinstance( + current_fval[0], torch.Tensor + ): + f_size = len(current_fval) current_fval = sum(current_fval) def _integrand(x): return sum(f(x)) - else: + elif isinstance(current_fval, torch.Tensor): + f_size = 1 def _integrand(x): return f(x) - + else: + raise TypeError( + "f must return a torch.Tensor or a list/tuple of torch.Tensor." + ) type_fval = current_fval.dtype current_weight = mix_rate / current_jac + (1 - mix_rate) * current_fval.abs() current_weight.masked_fill_(current_weight < epsilon, epsilon) n_meas = epoch // thinning - mean_values = torch.zeros((f_size, n_meas), dtype=type_fval, device=self.device) - std_values = torch.zeros_like(mean_values) - mean_refvalues = torch.zeros(n_meas, dtype=type_fval, device=self.device) - std_refvalues = torch.zeros_like(mean_refvalues) def one_step(current_y, current_x, current_weight, current_jac): proposed_y = proposal_dist( @@ -219,6 +226,9 @@ def one_step(current_y, current_x, current_weight, current_jac): current_y, current_x, current_weight, current_jac ) + values = torch.zeros((self.nbatch, f_size), dtype=type_fval, device=self.device) + refvalues = torch.zeros(self.nbatch, dtype=type_fval, device=self.device) + for imeas in range(n_meas): for j in range(thinning): current_y, current_x, current_weight, current_jac = one_step( @@ -230,26 +240,21 @@ def one_step(current_y, current_x, current_weight, current_jac): ) batch_results_ref = 1 / (current_jac * current_weight) - mean_values[:, imeas] = torch.mean(batch_results, dim=-1) - std_values[:, imeas] = torch.std(batch_results, dim=-1) / self.nbatch**0.5 - - mean_refvalues[imeas] = torch.mean(batch_results_ref, dim=-1) - std_refvalues[imeas] = ( - torch.std(batch_results_ref, dim=-1) / self.nbatch**0.5 - ) + values += batch_results / n_meas + refvalues += batch_results_ref / n_meas results = np.array([RAvg() for _ in range(f_size)]) results_ref = RAvg() - for imeas in range(n_meas): - results_ref.sum_neval += self.nbatch - results_ref.add( - gvar.gvar(mean_refvalues[imeas].item(), std_refvalues[imeas].item()) - ) - for j in range(f_size): - results[j].sum_neval += self.nbatch - results[j].add( - gvar.gvar(mean_values[j, imeas].item(), std_values[j, imeas].item()) - ) + + mean_ref = refvalues.mean().item() + std_ref = refvalues.std().item() / self.nbatch**0.5 + + results_ref.add(gvar.gvar(mean_ref, std_ref)) + for i in range(f_size): + _mean = values[:, i].mean().item() + _std = values[:, i].std().item() / self.nbatch**0.5 + results[i].add(gvar.gvar(_mean, _std)) + if f_size == 1: return results[0] / results_ref * self._rangebounds.prod() else: diff --git a/src/maps.py b/src/maps.py index 93081b0..02f5fb1 100644 --- a/src/maps.py +++ b/src/maps.py @@ -95,7 +95,7 @@ def __init__(self, bounds, ninc=1000, alpha=0.5, device="cpu", dtype=torch.float ) self.clear() - def train(self, nsamples, f, nitn=5, alpha=0.5): + def train(self, nsamples, f, epoch=5, alpha=0.5): q0 = Uniform(self.bounds, device=self.device, dtype=self.dtype) u, log_detJ0 = q0.sample(nsamples) @@ -110,7 +110,7 @@ def _integrand(x): def _integrand(x): return f(x) - for _ in range(nitn): + for _ in range(epoch): x, log_detJ = self.forward(u) f2 = torch.exp(2 * (log_detJ + log_detJ0)) * _integrand(x) ** 2 self.add_training_data(u, f2) diff --git a/src/mc_test.py b/src/mc_test.py index ba74199..9c0f60e 100644 --- a/src/mc_test.py +++ b/src/mc_test.py @@ -50,14 +50,13 @@ def sharp_peak(x): # Start Monte Carlo integration, including plain-MC, MCMC, vegas, and vegas-MCMC mc_integrator = MonteCarlo( - bounds, + bounds=bounds, neval=n_eval, nbatch=n_batch, - # nbatch=n_eval, device=device, ) mcmc_integrator = MCMC( - bounds, neval=n_eval, nbatch=n_batch, nburnin=n_therm, device=device + bounds=bounds, neval=n_eval, nbatch=n_batch, nburnin=n_therm, device=device ) print("Calculate the area of the unit circle f(x1, x2) in the bounds [-1, 1]^2...") @@ -69,7 +68,6 @@ def sharp_peak(x): vegas_map.train(20000, unit_circle_integrand, alpha=0.5) vegas_integrator = MonteCarlo( - bounds, maps=vegas_map, neval=n_eval, nbatch=n_batch, @@ -80,7 +78,6 @@ def sharp_peak(x): print("VEGAS Integral results: ", res) vegasmcmc_integrator = MCMC( - bounds, maps=vegas_map, neval=n_eval, nbatch=n_batch, @@ -102,7 +99,7 @@ def sharp_peak(x): print("MCMC Integral results:", res) # train the vegas map -vegas_map.train(20000, half_sphere_integrand, nitn=10, alpha=0.5) +vegas_map.train(20000, half_sphere_integrand, epoch=10, alpha=0.5) res = vegas_integrator(half_sphere_integrand) print("VEGAS Integral results: ", res) @@ -124,7 +121,7 @@ def sharp_peak(x): print(f" Integral 2: ", res[1]) # print("VEAGS map is trained for g(x1, x2)") -vegas_map.train(20000, integrand_list, nitn=10, alpha=0.5) +vegas_map.train(20000, integrand_list, epoch=10, alpha=0.5) res = vegas_integrator(integrand_list) print("VEGAS Integral results:") print(" Integral 1: ", res[0]) @@ -140,14 +137,14 @@ def sharp_peak(x): bounds = [(0, 1)] * 4 mc_integrator = MonteCarlo( - bounds, + bounds=bounds, neval=n_eval, nbatch=n_batch, # nbatch=n_eval, device=device, ) mcmc_integrator = MCMC( - bounds, neval=n_eval, nbatch=n_batch, nburnin=n_therm, device=device + bounds=bounds, neval=n_eval, nbatch=n_batch, nburnin=n_therm, device=device ) print("Plain MC Integral results:") res = mc_integrator(integrand_list1) @@ -176,12 +173,11 @@ def sharp_peak(x): vegas_map = Vegas(bounds, device=device) print("train VEGAS map for h(X)...") -# vegas_map.train(20000, sharp_peak, nitn=10, alpha=0.5) -vegas_map.train(20000, integrand_list1, nitn=10, alpha=0.5) +# vegas_map.train(20000, sharp_peak, epoch=10, alpha=0.5) +vegas_map.train(20000, integrand_list1, epoch=10, alpha=0.5) print("VEGAS Integral results:") vegas_integrator = MonteCarlo( - bounds, maps=vegas_map, neval=n_eval, nbatch=n_batch, @@ -202,7 +198,6 @@ def sharp_peak(x): print("VEGAS-MCMC Integral results:") vegasmcmc_integrator = MCMC( - bounds, maps=vegas_map, neval=n_eval, nbatch=n_batch, diff --git a/src/vegas_test.py b/src/vegas_test.py index a0b7c5e..4ca34a5 100644 --- a/src/vegas_test.py +++ b/src/vegas_test.py @@ -3,7 +3,7 @@ from maps import Vegas, Linear from utils import set_seed, get_device -set_seed(42) +# set_seed(42) # device = get_device() device = torch.device("cpu") @@ -23,11 +23,40 @@ def sharp_peak(x): return torch.exp(-200 * dx2) +def func(x): + return torch.log(x[:, 0]) / torch.sqrt(x[:, 0]) + + ninc = 1000 -# n_eval = 100000 n_eval = 50000 n_batch = 10000 n_therm = 10 + +print("\nCalculate the integral log(x)/x^0.5 in the bounds [0, 1]") + +print("train VEGAS map") +vegas_map = Vegas([(0, 1)], device=device, ninc=ninc) +vegas_map.train(20000, func, epoch=10, alpha=0.5) + +vegas_integrator = MonteCarlo( + maps=vegas_map, + neval=1000000, + nbatch=n_batch, + device=device, +) +res = vegas_integrator(func) +print("VEGAS Integral results: ", res) + +vegasmcmc_integrator = MCMC( + maps=vegas_map, + neval=1000000, + nbatch=n_batch, + nburnin=n_therm, + device=device, +) +res = vegasmcmc_integrator(func, mix_rate=0.5) +print("VEGAS-MCMC Integral results: ", res) + # Start Monte Carlo integration, including plain-MC, MCMC, vegas, and vegas-MCMC print("\nCalculate the integral [h(X), x1 * h(X), x1^2 * h(X)] in the bounds [0, 1]^4") print("h(X) = exp(-200 * (x1^2 + x2^2 + x3^2 + x4^2))") @@ -35,12 +64,11 @@ def sharp_peak(x): bounds = [(0, 1)] * 4 vegas_map = Vegas(bounds, device=device, ninc=ninc) print("train VEGAS map for h(X)...") -vegas_map.train(20000, sharp_peak, nitn=10, alpha=0.5) +vegas_map.train(20000, sharp_peak, epoch=10, alpha=0.5) # print(vegas_map.extract_grid()) print("VEGAS Integral results:") vegas_integrator = MonteCarlo( - bounds, maps=vegas_map, neval=n_eval, # nbatch=n_batch, @@ -61,7 +89,6 @@ def sharp_peak(x): print("VEGAS-MCMC Integral results:") vegasmcmc_integrator = MCMC( - bounds, maps=vegas_map, neval=n_eval, nbatch=n_batch, From e4b16f7def586f152f7c87c468439c159208cafa Mon Sep 17 00:00:00 2001 From: Tao Wang Date: Wed, 23 Oct 2024 22:33:49 -0400 Subject: [PATCH 10/10] add empty test --- src/mc_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/mc_test.py b/src/mc_test.py index 9c0f60e..9172c8b 100644 --- a/src/mc_test.py +++ b/src/mc_test.py @@ -8,6 +8,10 @@ device = torch.device("cpu") +def test_nothing(): + pass + + def unit_circle_integrand(x): inside_circle = (x[:, 0] ** 2 + x[:, 1] ** 2 < 1).double() return inside_circle