diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml new file mode 100644 index 0000000..7e4d320 --- /dev/null +++ b/.github/workflows/CI.yml @@ -0,0 +1,79 @@ +name: CI + +on: + push: + branches: + - master + pull_request: + branches: + - master + +permissions: + contents: write + +jobs: + test: + name: Python ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }} + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + version: + - "3.12" + os: + - ubuntu-latest + arch: + - x64 + - x86 + + steps: + - uses: actions/checkout@v2 + + - name: Set up Python ${{ matrix.version }} + uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.version }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install coverage + pip install pytest pytest-cov + pip install -r requirements.txt + + - name: Run tests + run: | + export PYTHONPATH=src + pytest --cov --cov-report=xml + + - name: Upload coverage to Codecov + uses: codecov/codecov-action@v4 + with: + token: ${{ secrets.CODECOV_TOKEN }} + + docs: + name: Documentation + runs-on: ubuntu-latest + needs: test + + steps: + - uses: actions/checkout@v2 + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install sphinx + pip install -r requirements.txt + + - name: Build documentation + run: | + cd docs + sphinx-apidoc -o source ../src + python ../remove_prefix.py + make html + + - name: Deploy documentation to GitHub Pages + uses: peaceiris/actions-gh-pages@v3 + with: + github_token: ${{ secrets.GITHUB_TOKEN }} + publish_dir: docs/build/html diff --git a/.github/workflows/generate_docs b/.github/workflows/generate_docs deleted file mode 100644 index 43a976e..0000000 --- a/.github/workflows/generate_docs +++ /dev/null @@ -1,55 +0,0 @@ -name: CI - -on: - push: - branches: [ "dev" ] - - # Allows you to run this workflow manually from the Actions tab - workflow_dispatch: - -jobs: - build: - runs-on: ubuntu-latest - - steps: - - name: Checkout code - uses: actions/checkout@v3 - with: - ref: dev - - - name: Set up Python - uses: actions/setup-python@v4 - with: - python-version: '3.11' - - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install sphinx - # other dependencies... - - - name: Generate documentation - run: | - cd doc - make html - - - name: Checkout dev branch again - uses: actions/checkout@v3 - with: - ref: dev - path: docs-output - - - name: Copy generated docs to dev branch - run: | - cp -r doc/build/html/* docs-output/docs/ - - - name: Commit and push changes - run: | - cd docs-output - git config --global user.name "github-actions[bot]" - git config --global user.email "github-actions[bot]@users.noreply.github.com" - git add doc - git commit -m "Auto-generated documentation" - git push origin dev - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/README.md b/README.md index 7a3a26a..546e331 100644 --- a/README.md +++ b/README.md @@ -1 +1,3 @@ # MCintegration +[![alpha](https://img.shields.io/badge/docs-alpha-blue.svg)](https://numericaleft.github.io/MCintegration.py/) +[![codecov](https://codecov.io/gh/numericalEFT/MCintegration.py/graph/badge.svg?token=851N2CNOTN)](https://codecov.io/gh/numericalEFT/MCintegration.py) diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 0000000..d0c3cbf --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = source +BUILDDIR = build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/make.bat b/docs/make.bat new file mode 100644 index 0000000..dc1312a --- /dev/null +++ b/docs/make.bat @@ -0,0 +1,35 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=source +set BUILDDIR=build + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.https://www.sphinx-doc.org/ + exit /b 1 +) + +if "%1" == "" goto help + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + +:end +popd diff --git a/docs/source/conf.py b/docs/source/conf.py new file mode 100644 index 0000000..fd16701 --- /dev/null +++ b/docs/source/conf.py @@ -0,0 +1,62 @@ +# Configuration file for the Sphinx documentation builder. +# +# For the full list of built-in configuration values, see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +# -- Project information ----------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information +import os +import sys + +sys.path.insert(0, os.path.abspath('../../src')) + +project = 'MCintegration' +copyright = '2024, Authors' +author = 'Authors' +release = '1.0.0' + +# -- General configuration --------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration + +extensions = [ + 'sphinx.ext.autodoc', + 'sphinx.ext.viewcode', + 'sphinx.ext.napoleon', + 'sphinx.ext.autosummary' +] + +templates_path = ['_templates'] +exclude_patterns = [] + + + +# -- Options for HTML output ------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output + +templates_path = ['_templates'] + +exclude_trees = ['_build'] + +exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] + +source_suffix = '.rst' + +source_encoding = 'utf-8' + +master_doc = 'index' + +napoleon_google_docstring = True +napoleon_numpy_docstring = True +napoleon_include_private_with_doc = False +napoleon_include_special_with_doc = True +napoleon_use_admonition_for_examples = False +napoleon_use_admonition_for_notes = False +napoleon_use_admonition_for_references = False +napoleon_use_ivar = False +napoleon_use_param = True +napoleon_use_rtype = True + +html_theme = 'sphinxdoc' +html_theme = 'nature' +html_theme = 'pyramid' +html_static_path = ['_static'] diff --git a/docs/source/index.rst b/docs/source/index.rst new file mode 100644 index 0000000..4c55073 --- /dev/null +++ b/docs/source/index.rst @@ -0,0 +1,17 @@ +.. MCintegration documentation master file, created by + sphinx-quickstart on Tue Oct 15 10:32:44 2024. + You can adapt this file completely to your liking, but it should at least + contain the root `toctree` directive. + +MCintegration documentation +=============== + +Add your content using ``reStructuredText`` syntax. See the +`reStructuredText `_ +documentation for details. + + +.. toctree:: + :maxdepth: 2 + :caption: Contents: + 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/remove_prefix.py b/remove_prefix.py new file mode 100644 index 0000000..5fa1773 --- /dev/null +++ b/remove_prefix.py @@ -0,0 +1,14 @@ +import os + +docs_source_dir = './source' + +target_file = 'src.rst' + +target_file_path = os.path.join(docs_source_dir, target_file) + +if os.path.exists(target_file_path): + with open(target_file_path, 'r') as f: + content = f.read() + content = content.replace('src.', '') + with open(target_file_path, 'w') as f: + f.write(content) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..a10924c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,8 @@ +torch +numpy +normflows +vegas +mpmath +matplotlib +tqdm +absl-py diff --git a/src/__init__.py b/src/__init__.py index 21701fe..ca21233 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -1,3 +1,2 @@ from .integrators import MonteCarlo, MCMC -from .maps import VegasMap from .utils import RAvg, set_seed, get_device 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 9e19fc9..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,50 +64,59 @@ 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(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 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() + 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 + values = torch.zeros((self.nbatch, f_size), dtype=type_fval, device=self.device) - 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) + values += batch_results / epoch - result.sum_neval += self.neval - result.add(gvar.gvar(mean.item(), (var**0.5).item())) - return result + results = np.array([RAvg() for _ in range(f_size)]) + 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: + return results 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): @@ -130,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) @@ -150,52 +162,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 - # vars_shape = (self.nbatch, self.dim) + epoch = self.neval // self.nbatch 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) + 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)) + 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) - # 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) - - 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) - result = RAvg() - result_ref = RAvg() + n_meas = epoch // thinning - epoch = self.neval // self.nbatch - n_meas = 0 - - def _propose(current_y, current_fval, 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 ) 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 +215,47 @@ 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 = one_step( + current_y, current_x, current_weight, current_jac ) - for i in range(epoch // thinning): + + 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_fval, current_weight, current_jac = _propose( - current_y, current_fval, current_weight, current_jac + current_y, current_x, current_weight, current_jac = one_step( + 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}") + + values += batch_results / n_meas + refvalues += batch_results_ref / n_meas + + results = np.array([RAvg() for _ in range(f_size)]) + results_ref = RAvg() + + 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: + return results / results_ref * self._rangebounds.prod().item() diff --git a/src/maps.py b/src/maps.py index 2b6a6df..02f5fb1 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, epoch=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(epoch): + 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 47740c0..9c0f60e 100644 --- a/src/mc_test.py +++ b/src/mc_test.py @@ -4,23 +4,12 @@ from utils import set_seed, get_device set_seed(42) -device = get_device() -# device = torch.device("cpu") +# device = get_device() +device = torch.device("cpu") 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 @@ -28,50 +17,201 @@ 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] + + +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=400000, nbatch=1000, device=device - bounds, - neval=400000, - nbatch=1000, + bounds=bounds, + neval=n_eval, + nbatch=n_batch, device=device, ) +mcmc_integrator = MCMC( + 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...") 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=400000, nbatch=1000, device=device)( - unit_circle_integrand +vegas_map.train(20000, unit_circle_integrand, alpha=0.5) +vegas_integrator = MonteCarlo( + 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=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}") -print(f" Error: {res.sdev}") +vegasmcmc_integrator = MCMC( + 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=400000, nbatch=1000, nburnin=100, device=device) res = mcmc_integrator(half_sphere_integrand, mix_rate=0.5) +print("MCMC Integral results:", res) + +# train the vegas map +vegas_map.train(20000, half_sphere_integrand, epoch=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(" 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: {res.mean}") -print(f" Error: {res.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, epoch=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=bounds, + neval=n_eval, + nbatch=n_batch, + # nbatch=n_eval, + device=device, +) +mcmc_integrator = MCMC( + bounds=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], + " 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], + " 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, epoch=10, alpha=0.5) +vegas_map.train(20000, integrand_list1, epoch=10, alpha=0.5) + +print("VEGAS Integral results:") +vegas_integrator = MonteCarlo( + 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( + 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/requirements.txt b/src/requirements.txt new file mode 100644 index 0000000..6f8bf5b --- /dev/null +++ b/src/requirements.txt @@ -0,0 +1,5 @@ +torch +numpy +normflows +vegas +mpmath diff --git a/src/vegas_test.py b/src/vegas_test.py new file mode 100644 index 0000000..4ca34a5 --- /dev/null +++ b/src/vegas_test.py @@ -0,0 +1,108 @@ +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) + + +def func(x): + return torch.log(x[:, 0]) / torch.sqrt(x[:, 0]) + + +ninc = 1000 +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))") + +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, epoch=10, alpha=0.5) +# print(vegas_map.extract_grid()) + +print("VEGAS Integral results:") +vegas_integrator = MonteCarlo( + 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( + 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], +)