We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
fix progress bar and simulation batch size in simulate_in_batches
Two problems related to `simulate_in_batches`: 

1) The progress bar fills up by sending jobs to the workers, i.e., before these are actually done. E.g., when running 100 parameters with batch size 10 and 10 workers, the progress bar is done instantly, but simulations may continue for minutes or hours. We should find a way to fill up the pbar only when the jobs are done. 

2) The simulation batch size defaults to 1. This can make things very slow when one uses a simulator that cannot simulate multiple parameters at once (vectorised) and one use multiple workers. In this scenario each worker deals with a single simulation -- creating a lot of overhead. It would be much better (for me it gave a 10x speed up) to write a simple "batch"-simulator with a for loop and set the simulation batch size to 10 or 100 so that each worker call will run at least 10 or 100 simulations before going back. Actually, it could make sense to just set 
```
sim_batch_size = int(num_simulations / num_workers)
``` 
in order to minimise overhead. 
The only downside would be that the pbar becomes less informative. So one could think about reducing the batch size a bit to obtain more "steps" in the pbar. 

Any thoughts? 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `setup.py`
Content:
```
1 #!/usr/bin/env python
2 # -*- coding: utf-8 -*-
3 #
4 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
5 # under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.
6 #
7 # Note: To use the 'upload' functionality of this file, you must:
8 #   $ pipenv install twine --dev
9 
10 import io
11 import os
12 import sys
13 from shutil import rmtree
14 
15 from setuptools import find_packages, setup, Command
16 
17 # Package meta-data.
18 NAME = "sbi"
19 DESCRIPTION = "Simulation-based inference."
20 KEYWORDS = "bayesian parameter inference system_identification simulator PyTorch"
21 URL = "https://github.com/mackelab/sbi"
22 EMAIL = "sbi@mackelab.org"
23 AUTHOR = "Álvaro Tejero-Cantero, Jakob H. Macke, Jan-Matthis Lückmann, Conor M. Durkan, Michael Deistler, Jan Bölts"
24 REQUIRES_PYTHON = ">=3.6.0"
25 
26 REQUIRED = [
27     "joblib",
28     "matplotlib",
29     "numpy",
30     "pillow",
31     "pyknos>=0.14.0",
32     "pyro-ppl>=1.3.1",
33     "scipy",
34     "tensorboard",
35     "torch>=1.6.0",
36     "tqdm",
37 ]
38 
39 EXTRAS = {
40     "dev": [
41         "autoflake",
42         "black",
43         "deepdiff",
44         "flake8",
45         "isort",
46         "jupyter",
47         "mkdocs",
48         "mkdocs-material",
49         "markdown-include",
50         "mkdocs-redirects",
51         "mkdocstrings",
52         "nbconvert",
53         "pep517",
54         "pytest",
55         "pyyaml",
56         "scikit-learn",
57         "torchtestcase",
58         "twine",
59     ],
60 }
61 
62 here = os.path.abspath(os.path.dirname(__file__))
63 
64 # Import the README and use it as the long-description.
65 try:
66     with io.open(os.path.join(here, "README.md"), encoding="utf-8") as f:
67         long_description = "\n" + f.read()
68 except FileNotFoundError:
69     long_description = DESCRIPTION
70 
71 # Load the package's __version__.py module as a dictionary.
72 about = {}
73 project_slug = NAME.lower().replace("-", "_").replace(" ", "_")
74 with open(os.path.join(here, project_slug, "__version__.py")) as f:
75     exec(f.read(), about)
76 
77 
78 class UploadCommand(Command):
79     """Support setup.py upload."""
80 
81     description = "Build and publish the package."
82     user_options = []
83 
84     @staticmethod
85     def status(s):
86         """Prints things in bold."""
87         print("\033[1m{0}\033[0m".format(s))
88 
89     def initialize_options(self):
90         pass
91 
92     def finalize_options(self):
93         pass
94 
95     def run(self):
96         try:
97             self.status("Removing previous builds…")
98             rmtree(os.path.join(here, "dist"))
99         except OSError:
100             pass
101 
102         self.status("Building Source and Wheel (universal) distribution…")
103         os.system("{0} setup.py sdist bdist_wheel --universal".format(sys.executable))
104 
105         self.status("Uploading the package to PyPI via Twine…")
106         os.system("twine upload dist/*")
107 
108         self.status("Pushing git tags…")
109         os.system("git tag v{0}".format(about["__version__"]))
110         os.system("git push --tags")
111 
112         sys.exit()
113 
114 
115 setup(
116     name=NAME,
117     version=about["__version__"],
118     description=DESCRIPTION,
119     keywords=KEYWORDS,
120     long_description=long_description,
121     long_description_content_type="text/markdown",
122     author=AUTHOR,
123     author_email=EMAIL,
124     python_requires=REQUIRES_PYTHON,
125     url=URL,
126     packages=find_packages(exclude=["tests", "*.tests", "*.tests.*", "tests.*"]),
127     install_requires=REQUIRED,
128     extras_require=EXTRAS,
129     include_package_data=True,
130     license="AGPLv3",
131     classifiers=[
132         # Full list: https://pypi.python.org/pypi?%3Aaction=list_classifiers
133         "Development Status :: 3 - Alpha",
134         "License :: OSI Approved :: GNU Affero General Public License v3 or later (AGPLv3+)",
135         "Intended Audience :: Developers",
136         "Intended Audience :: Education",
137         "Intended Audience :: Science/Research",
138         "Topic :: Adaptive Technologies",
139         "Topic :: Scientific/Engineering",
140         "Topic :: Scientific/Engineering :: Artificial Intelligence",
141         "Topic :: Scientific/Engineering :: Mathematics",
142         "Programming Language :: Python",
143         "Programming Language :: Python :: 3",
144         "Programming Language :: Python :: 3.6",
145         "Programming Language :: Python :: 3.7",
146         "Programming Language :: Python :: 3.8",
147     ],
148     # $ setup.py publish support.
149     cmdclass=dict(upload=UploadCommand),
150 )
151
```

Path: `sbi/simulators/simutils.py`
Content:
```
1 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
2 # under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.
3 
4 
5 from typing import Callable
6 
7 import torch
8 from torch import Tensor
9 from tqdm.auto import tqdm
10 from joblib import Parallel, delayed
11 
12 
13 def simulate_in_batches(
14     simulator: Callable,
15     theta: Tensor,
16     sim_batch_size: int = 1,
17     num_workers: int = 1,
18     show_progress_bars: bool = True,
19 ) -> Tensor:
20     r"""
21     Return simulations $x$ for parameters $\theta$ conducted batchwise.
22 
23     Parameters are batched with size `sim_batch_size` (default whole theta at once).
24     Multiprocessing is used when `num_workers > 1`.
25 
26     Args:
27         simulator: Simulator callable (a function or a class with `__call__`).
28         theta: All parameters $\theta$ sampled from prior or posterior.
29         sim_batch_size: Number of simulations per batch. Default is to simulate
30             the entire theta in a single batch.
31         num_workers: Number of workers for multiprocessing.
32         show_progress_bars: Whether to show a progress bar during simulation.
33 
34     Returns:
35         Parameters theta and simulations $x$.
36     """
37 
38     num_sims, *_ = theta.shape
39 
40     if num_sims == 0:
41         x = torch.tensor([])
42     elif sim_batch_size is not None and sim_batch_size < num_sims:
43         # Dev note: pyright complains of torch.split lacking a type stub
44         # as of PyTorch 1.4.0, see https://github.com/microsoft/pyright/issues/291
45         batches = torch.split(theta, sim_batch_size, dim=0)
46 
47         if num_workers > 1:
48             # Parallelize the sequence of batches across workers.
49             # TODO: This usage of tqdm tracks the dispatching of jobs instead of the
50             # moment when they are done, resulting in waiting time at 100% in case the
51             # last jobs takes long. A potential solution can be found here: https://
52             # stackoverflow.com/a/61689175
53             simulation_outputs = Parallel(n_jobs=num_workers)(
54                 delayed(simulator)(batch)
55                 for batch in tqdm(
56                     batches,
57                     disable=not show_progress_bars,
58                     desc=f"Running {num_sims} simulations in {len(batches)} batches.",
59                     total=len(batches),
60                 )
61             )
62         else:
63             pbar = tqdm(
64                 total=num_sims,
65                 disable=not show_progress_bars,
66                 desc=f"Running {num_sims} simulations.",
67             )
68 
69             with pbar:
70                 simulation_outputs = []
71                 for batch in batches:
72                     simulation_outputs.append(simulator(batch))
73                     pbar.update(sim_batch_size)
74 
75         x = torch.cat(simulation_outputs, dim=0)
76     else:
77         x = simulator(theta)
78 
79     return x
80
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/sbi/simulators/simutils.py b/sbi/simulators/simutils.py
--- a/sbi/simulators/simutils.py
+++ b/sbi/simulators/simutils.py
@@ -8,6 +8,8 @@
 from torch import Tensor
 from tqdm.auto import tqdm
 from joblib import Parallel, delayed
+import contextlib
+import joblib
 
 
 def simulate_in_batches(
@@ -46,19 +48,19 @@
 
         if num_workers > 1:
             # Parallelize the sequence of batches across workers.
-            # TODO: This usage of tqdm tracks the dispatching of jobs instead of the
-            # moment when they are done, resulting in waiting time at 100% in case the
-            # last jobs takes long. A potential solution can be found here: https://
-            # stackoverflow.com/a/61689175
-            simulation_outputs = Parallel(n_jobs=num_workers)(
-                delayed(simulator)(batch)
-                for batch in tqdm(
+            # We use the solution proposed here: https://stackoverflow.com/a/61689175
+            # to update the pbar only after the workers finished a task.
+            with tqdm_joblib(
+                tqdm(
                     batches,
                     disable=not show_progress_bars,
                     desc=f"Running {num_sims} simulations in {len(batches)} batches.",
                     total=len(batches),
                 )
-            )
+            ) as progress_bar:
+                simulation_outputs = Parallel(n_jobs=num_workers)(
+                    delayed(simulator)(batch) for batch in batches
+                )
         else:
             pbar = tqdm(
                 total=num_sims,
@@ -77,3 +79,28 @@
         x = simulator(theta)
 
     return x
+
+
+@contextlib.contextmanager
+def tqdm_joblib(tqdm_object):
+    """Context manager to patch joblib to report into tqdm progress bar given as
+    argument
+
+    This wrapped context manager obtains the number of finished tasks from the tqdm
+    print function and uses it to update the pbar, as suggested in
+    https://stackoverflow.com/a/61689175. See #419, #421
+    """
+
+    def tqdm_print_progress(self):
+        if self.n_completed_tasks > tqdm_object.n:
+            n_completed = self.n_completed_tasks - tqdm_object.n
+            tqdm_object.update(n=n_completed)
+
+    original_print_progress = joblib.parallel.Parallel.print_progress
+    joblib.parallel.Parallel.print_progress = tqdm_print_progress
+
+    try:
+        yield tqdm_object
+    finally:
+        joblib.parallel.Parallel.print_progress = original_print_progress
+        tqdm_object.close()
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -24,7 +24,7 @@
 REQUIRES_PYTHON = ">=3.6.0"
 
 REQUIRED = [
-    "joblib",
+    "joblib>=1.0.0",
     "matplotlib",
     "numpy",
     "pillow",

{"golden_diff": "diff --git a/sbi/simulators/simutils.py b/sbi/simulators/simutils.py\n--- a/sbi/simulators/simutils.py\n+++ b/sbi/simulators/simutils.py\n@@ -8,6 +8,8 @@\n from torch import Tensor\n from tqdm.auto import tqdm\n from joblib import Parallel, delayed\n+import contextlib\n+import joblib\n \n \n def simulate_in_batches(\n@@ -46,19 +48,19 @@\n \n         if num_workers > 1:\n             # Parallelize the sequence of batches across workers.\n-            # TODO: This usage of tqdm tracks the dispatching of jobs instead of the\n-            # moment when they are done, resulting in waiting time at 100% in case the\n-            # last jobs takes long. A potential solution can be found here: https://\n-            # stackoverflow.com/a/61689175\n-            simulation_outputs = Parallel(n_jobs=num_workers)(\n-                delayed(simulator)(batch)\n-                for batch in tqdm(\n+            # We use the solution proposed here: https://stackoverflow.com/a/61689175\n+            # to update the pbar only after the workers finished a task.\n+            with tqdm_joblib(\n+                tqdm(\n                     batches,\n                     disable=not show_progress_bars,\n                     desc=f\"Running {num_sims} simulations in {len(batches)} batches.\",\n                     total=len(batches),\n                 )\n-            )\n+            ) as progress_bar:\n+                simulation_outputs = Parallel(n_jobs=num_workers)(\n+                    delayed(simulator)(batch) for batch in batches\n+                )\n         else:\n             pbar = tqdm(\n                 total=num_sims,\n@@ -77,3 +79,28 @@\n         x = simulator(theta)\n \n     return x\n+\n+\n+@contextlib.contextmanager\n+def tqdm_joblib(tqdm_object):\n+    \"\"\"Context manager to patch joblib to report into tqdm progress bar given as\n+    argument\n+\n+    This wrapped context manager obtains the number of finished tasks from the tqdm\n+    print function and uses it to update the pbar, as suggested in\n+    https://stackoverflow.com/a/61689175. See #419, #421\n+    \"\"\"\n+\n+    def tqdm_print_progress(self):\n+        if self.n_completed_tasks > tqdm_object.n:\n+            n_completed = self.n_completed_tasks - tqdm_object.n\n+            tqdm_object.update(n=n_completed)\n+\n+    original_print_progress = joblib.parallel.Parallel.print_progress\n+    joblib.parallel.Parallel.print_progress = tqdm_print_progress\n+\n+    try:\n+        yield tqdm_object\n+    finally:\n+        joblib.parallel.Parallel.print_progress = original_print_progress\n+        tqdm_object.close()\ndiff --git a/setup.py b/setup.py\n--- a/setup.py\n+++ b/setup.py\n@@ -24,7 +24,7 @@\n REQUIRES_PYTHON = \">=3.6.0\"\n \n REQUIRED = [\n-    \"joblib\",\n+    \"joblib>=1.0.0\",\n     \"matplotlib\",\n     \"numpy\",\n     \"pillow\",\n", "issue": "fix progress bar and simulation batch size in simulate_in_batches\nTwo problems related to `simulate_in_batches`: \r\n\r\n1) The progress bar fills up by sending jobs to the workers, i.e., before these are actually done. E.g., when running 100 parameters with batch size 10 and 10 workers, the progress bar is done instantly, but simulations may continue for minutes or hours. We should find a way to fill up the pbar only when the jobs are done. \r\n\r\n2) The simulation batch size defaults to 1. This can make things very slow when one uses a simulator that cannot simulate multiple parameters at once (vectorised) and one use multiple workers. In this scenario each worker deals with a single simulation -- creating a lot of overhead. It would be much better (for me it gave a 10x speed up) to write a simple \"batch\"-simulator with a for loop and set the simulation batch size to 10 or 100 so that each worker call will run at least 10 or 100 simulations before going back. Actually, it could make sense to just set \r\n```\r\nsim_batch_size = int(num_simulations / num_workers)\r\n``` \r\nin order to minimise overhead. \r\nThe only downside would be that the pbar becomes less informative. So one could think about reducing the batch size a bit to obtain more \"steps\" in the pbar. \r\n\r\nAny thoughts? \n", "before_files": [{"content": "#!/usr/bin/env python\n# -*- coding: utf-8 -*-\n#\n# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed\n# under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.\n#\n# Note: To use the 'upload' functionality of this file, you must:\n#   $ pipenv install twine --dev\n\nimport io\nimport os\nimport sys\nfrom shutil import rmtree\n\nfrom setuptools import find_packages, setup, Command\n\n# Package meta-data.\nNAME = \"sbi\"\nDESCRIPTION = \"Simulation-based inference.\"\nKEYWORDS = \"bayesian parameter inference system_identification simulator PyTorch\"\nURL = \"https://github.com/mackelab/sbi\"\nEMAIL = \"sbi@mackelab.org\"\nAUTHOR = \"\u00c1lvaro Tejero-Cantero, Jakob H. Macke, Jan-Matthis L\u00fcckmann, Conor M. Durkan, Michael Deistler, Jan B\u00f6lts\"\nREQUIRES_PYTHON = \">=3.6.0\"\n\nREQUIRED = [\n    \"joblib\",\n    \"matplotlib\",\n    \"numpy\",\n    \"pillow\",\n    \"pyknos>=0.14.0\",\n    \"pyro-ppl>=1.3.1\",\n    \"scipy\",\n    \"tensorboard\",\n    \"torch>=1.6.0\",\n    \"tqdm\",\n]\n\nEXTRAS = {\n    \"dev\": [\n        \"autoflake\",\n        \"black\",\n        \"deepdiff\",\n        \"flake8\",\n        \"isort\",\n        \"jupyter\",\n        \"mkdocs\",\n        \"mkdocs-material\",\n        \"markdown-include\",\n        \"mkdocs-redirects\",\n        \"mkdocstrings\",\n        \"nbconvert\",\n        \"pep517\",\n        \"pytest\",\n        \"pyyaml\",\n        \"scikit-learn\",\n        \"torchtestcase\",\n        \"twine\",\n    ],\n}\n\nhere = os.path.abspath(os.path.dirname(__file__))\n\n# Import the README and use it as the long-description.\ntry:\n    with io.open(os.path.join(here, \"README.md\"), encoding=\"utf-8\") as f:\n        long_description = \"\\n\" + f.read()\nexcept FileNotFoundError:\n    long_description = DESCRIPTION\n\n# Load the package's __version__.py module as a dictionary.\nabout = {}\nproject_slug = NAME.lower().replace(\"-\", \"_\").replace(\" \", \"_\")\nwith open(os.path.join(here, project_slug, \"__version__.py\")) as f:\n    exec(f.read(), about)\n\n\nclass UploadCommand(Command):\n    \"\"\"Support setup.py upload.\"\"\"\n\n    description = \"Build and publish the package.\"\n    user_options = []\n\n    @staticmethod\n    def status(s):\n        \"\"\"Prints things in bold.\"\"\"\n        print(\"\\033[1m{0}\\033[0m\".format(s))\n\n    def initialize_options(self):\n        pass\n\n    def finalize_options(self):\n        pass\n\n    def run(self):\n        try:\n            self.status(\"Removing previous builds\u2026\")\n            rmtree(os.path.join(here, \"dist\"))\n        except OSError:\n            pass\n\n        self.status(\"Building Source and Wheel (universal) distribution\u2026\")\n        os.system(\"{0} setup.py sdist bdist_wheel --universal\".format(sys.executable))\n\n        self.status(\"Uploading the package to PyPI via Twine\u2026\")\n        os.system(\"twine upload dist/*\")\n\n        self.status(\"Pushing git tags\u2026\")\n        os.system(\"git tag v{0}\".format(about[\"__version__\"]))\n        os.system(\"git push --tags\")\n\n        sys.exit()\n\n\nsetup(\n    name=NAME,\n    version=about[\"__version__\"],\n    description=DESCRIPTION,\n    keywords=KEYWORDS,\n    long_description=long_description,\n    long_description_content_type=\"text/markdown\",\n    author=AUTHOR,\n    author_email=EMAIL,\n    python_requires=REQUIRES_PYTHON,\n    url=URL,\n    packages=find_packages(exclude=[\"tests\", \"*.tests\", \"*.tests.*\", \"tests.*\"]),\n    install_requires=REQUIRED,\n    extras_require=EXTRAS,\n    include_package_data=True,\n    license=\"AGPLv3\",\n    classifiers=[\n        # Full list: https://pypi.python.org/pypi?%3Aaction=list_classifiers\n        \"Development Status :: 3 - Alpha\",\n        \"License :: OSI Approved :: GNU Affero General Public License v3 or later (AGPLv3+)\",\n        \"Intended Audience :: Developers\",\n        \"Intended Audience :: Education\",\n        \"Intended Audience :: Science/Research\",\n        \"Topic :: Adaptive Technologies\",\n        \"Topic :: Scientific/Engineering\",\n        \"Topic :: Scientific/Engineering :: Artificial Intelligence\",\n        \"Topic :: Scientific/Engineering :: Mathematics\",\n        \"Programming Language :: Python\",\n        \"Programming Language :: Python :: 3\",\n        \"Programming Language :: Python :: 3.6\",\n        \"Programming Language :: Python :: 3.7\",\n        \"Programming Language :: Python :: 3.8\",\n    ],\n    # $ setup.py publish support.\n    cmdclass=dict(upload=UploadCommand),\n)\n", "path": "setup.py"}, {"content": "# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed\n# under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.\n\n\nfrom typing import Callable\n\nimport torch\nfrom torch import Tensor\nfrom tqdm.auto import tqdm\nfrom joblib import Parallel, delayed\n\n\ndef simulate_in_batches(\n    simulator: Callable,\n    theta: Tensor,\n    sim_batch_size: int = 1,\n    num_workers: int = 1,\n    show_progress_bars: bool = True,\n) -> Tensor:\n    r\"\"\"\n    Return simulations $x$ for parameters $\\theta$ conducted batchwise.\n\n    Parameters are batched with size `sim_batch_size` (default whole theta at once).\n    Multiprocessing is used when `num_workers > 1`.\n\n    Args:\n        simulator: Simulator callable (a function or a class with `__call__`).\n        theta: All parameters $\\theta$ sampled from prior or posterior.\n        sim_batch_size: Number of simulations per batch. Default is to simulate\n            the entire theta in a single batch.\n        num_workers: Number of workers for multiprocessing.\n        show_progress_bars: Whether to show a progress bar during simulation.\n\n    Returns:\n        Parameters theta and simulations $x$.\n    \"\"\"\n\n    num_sims, *_ = theta.shape\n\n    if num_sims == 0:\n        x = torch.tensor([])\n    elif sim_batch_size is not None and sim_batch_size < num_sims:\n        # Dev note: pyright complains of torch.split lacking a type stub\n        # as of PyTorch 1.4.0, see https://github.com/microsoft/pyright/issues/291\n        batches = torch.split(theta, sim_batch_size, dim=0)\n\n        if num_workers > 1:\n            # Parallelize the sequence of batches across workers.\n            # TODO: This usage of tqdm tracks the dispatching of jobs instead of the\n            # moment when they are done, resulting in waiting time at 100% in case the\n            # last jobs takes long. A potential solution can be found here: https://\n            # stackoverflow.com/a/61689175\n            simulation_outputs = Parallel(n_jobs=num_workers)(\n                delayed(simulator)(batch)\n                for batch in tqdm(\n                    batches,\n                    disable=not show_progress_bars,\n                    desc=f\"Running {num_sims} simulations in {len(batches)} batches.\",\n                    total=len(batches),\n                )\n            )\n        else:\n            pbar = tqdm(\n                total=num_sims,\n                disable=not show_progress_bars,\n                desc=f\"Running {num_sims} simulations.\",\n            )\n\n            with pbar:\n                simulation_outputs = []\n                for batch in batches:\n                    simulation_outputs.append(simulator(batch))\n                    pbar.update(sim_batch_size)\n\n        x = torch.cat(simulation_outputs, dim=0)\n    else:\n        x = simulator(theta)\n\n    return x\n", "path": "sbi/simulators/simutils.py"}], "after_files": [{"content": "#!/usr/bin/env python\n# -*- coding: utf-8 -*-\n#\n# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed\n# under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.\n#\n# Note: To use the 'upload' functionality of this file, you must:\n#   $ pipenv install twine --dev\n\nimport io\nimport os\nimport sys\nfrom shutil import rmtree\n\nfrom setuptools import find_packages, setup, Command\n\n# Package meta-data.\nNAME = \"sbi\"\nDESCRIPTION = \"Simulation-based inference.\"\nKEYWORDS = \"bayesian parameter inference system_identification simulator PyTorch\"\nURL = \"https://github.com/mackelab/sbi\"\nEMAIL = \"sbi@mackelab.org\"\nAUTHOR = \"\u00c1lvaro Tejero-Cantero, Jakob H. Macke, Jan-Matthis L\u00fcckmann, Conor M. Durkan, Michael Deistler, Jan B\u00f6lts\"\nREQUIRES_PYTHON = \">=3.6.0\"\n\nREQUIRED = [\n    \"joblib>=1.0.0\",\n    \"matplotlib\",\n    \"numpy\",\n    \"pillow\",\n    \"pyknos>=0.14.0\",\n    \"pyro-ppl>=1.3.1\",\n    \"scipy\",\n    \"tensorboard\",\n    \"torch>=1.6.0\",\n    \"tqdm\",\n]\n\nEXTRAS = {\n    \"dev\": [\n        \"autoflake\",\n        \"black\",\n        \"deepdiff\",\n        \"flake8\",\n        \"isort\",\n        \"jupyter\",\n        \"mkdocs\",\n        \"mkdocs-material\",\n        \"markdown-include\",\n        \"mkdocs-redirects\",\n        \"mkdocstrings\",\n        \"nbconvert\",\n        \"pep517\",\n        \"pytest\",\n        \"pyyaml\",\n        \"scikit-learn\",\n        \"torchtestcase\",\n        \"twine\",\n    ],\n}\n\nhere = os.path.abspath(os.path.dirname(__file__))\n\n# Import the README and use it as the long-description.\ntry:\n    with io.open(os.path.join(here, \"README.md\"), encoding=\"utf-8\") as f:\n        long_description = \"\\n\" + f.read()\nexcept FileNotFoundError:\n    long_description = DESCRIPTION\n\n# Load the package's __version__.py module as a dictionary.\nabout = {}\nproject_slug = NAME.lower().replace(\"-\", \"_\").replace(\" \", \"_\")\nwith open(os.path.join(here, project_slug, \"__version__.py\")) as f:\n    exec(f.read(), about)\n\n\nclass UploadCommand(Command):\n    \"\"\"Support setup.py upload.\"\"\"\n\n    description = \"Build and publish the package.\"\n    user_options = []\n\n    @staticmethod\n    def status(s):\n        \"\"\"Prints things in bold.\"\"\"\n        print(\"\\033[1m{0}\\033[0m\".format(s))\n\n    def initialize_options(self):\n        pass\n\n    def finalize_options(self):\n        pass\n\n    def run(self):\n        try:\n            self.status(\"Removing previous builds\u2026\")\n            rmtree(os.path.join(here, \"dist\"))\n        except OSError:\n            pass\n\n        self.status(\"Building Source and Wheel (universal) distribution\u2026\")\n        os.system(\"{0} setup.py sdist bdist_wheel --universal\".format(sys.executable))\n\n        self.status(\"Uploading the package to PyPI via Twine\u2026\")\n        os.system(\"twine upload dist/*\")\n\n        self.status(\"Pushing git tags\u2026\")\n        os.system(\"git tag v{0}\".format(about[\"__version__\"]))\n        os.system(\"git push --tags\")\n\n        sys.exit()\n\n\nsetup(\n    name=NAME,\n    version=about[\"__version__\"],\n    description=DESCRIPTION,\n    keywords=KEYWORDS,\n    long_description=long_description,\n    long_description_content_type=\"text/markdown\",\n    author=AUTHOR,\n    author_email=EMAIL,\n    python_requires=REQUIRES_PYTHON,\n    url=URL,\n    packages=find_packages(exclude=[\"tests\", \"*.tests\", \"*.tests.*\", \"tests.*\"]),\n    install_requires=REQUIRED,\n    extras_require=EXTRAS,\n    include_package_data=True,\n    license=\"AGPLv3\",\n    classifiers=[\n        # Full list: https://pypi.python.org/pypi?%3Aaction=list_classifiers\n        \"Development Status :: 3 - Alpha\",\n        \"License :: OSI Approved :: GNU Affero General Public License v3 or later (AGPLv3+)\",\n        \"Intended Audience :: Developers\",\n        \"Intended Audience :: Education\",\n        \"Intended Audience :: Science/Research\",\n        \"Topic :: Adaptive Technologies\",\n        \"Topic :: Scientific/Engineering\",\n        \"Topic :: Scientific/Engineering :: Artificial Intelligence\",\n        \"Topic :: Scientific/Engineering :: Mathematics\",\n        \"Programming Language :: Python\",\n        \"Programming Language :: Python :: 3\",\n        \"Programming Language :: Python :: 3.6\",\n        \"Programming Language :: Python :: 3.7\",\n        \"Programming Language :: Python :: 3.8\",\n    ],\n    # $ setup.py publish support.\n    cmdclass=dict(upload=UploadCommand),\n)\n", "path": "setup.py"}, {"content": "# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed\n# under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.\n\n\nfrom typing import Callable\n\nimport torch\nfrom torch import Tensor\nfrom tqdm.auto import tqdm\nfrom joblib import Parallel, delayed\nimport contextlib\nimport joblib\n\n\ndef simulate_in_batches(\n    simulator: Callable,\n    theta: Tensor,\n    sim_batch_size: int = 1,\n    num_workers: int = 1,\n    show_progress_bars: bool = True,\n) -> Tensor:\n    r\"\"\"\n    Return simulations $x$ for parameters $\\theta$ conducted batchwise.\n\n    Parameters are batched with size `sim_batch_size` (default whole theta at once).\n    Multiprocessing is used when `num_workers > 1`.\n\n    Args:\n        simulator: Simulator callable (a function or a class with `__call__`).\n        theta: All parameters $\\theta$ sampled from prior or posterior.\n        sim_batch_size: Number of simulations per batch. Default is to simulate\n            the entire theta in a single batch.\n        num_workers: Number of workers for multiprocessing.\n        show_progress_bars: Whether to show a progress bar during simulation.\n\n    Returns:\n        Parameters theta and simulations $x$.\n    \"\"\"\n\n    num_sims, *_ = theta.shape\n\n    if num_sims == 0:\n        x = torch.tensor([])\n    elif sim_batch_size is not None and sim_batch_size < num_sims:\n        # Dev note: pyright complains of torch.split lacking a type stub\n        # as of PyTorch 1.4.0, see https://github.com/microsoft/pyright/issues/291\n        batches = torch.split(theta, sim_batch_size, dim=0)\n\n        if num_workers > 1:\n            # Parallelize the sequence of batches across workers.\n            # We use the solution proposed here: https://stackoverflow.com/a/61689175\n            # to update the pbar only after the workers finished a task.\n            with tqdm_joblib(\n                tqdm(\n                    batches,\n                    disable=not show_progress_bars,\n                    desc=f\"Running {num_sims} simulations in {len(batches)} batches.\",\n                    total=len(batches),\n                )\n            ) as progress_bar:\n                simulation_outputs = Parallel(n_jobs=num_workers)(\n                    delayed(simulator)(batch) for batch in batches\n                )\n        else:\n            pbar = tqdm(\n                total=num_sims,\n                disable=not show_progress_bars,\n                desc=f\"Running {num_sims} simulations.\",\n            )\n\n            with pbar:\n                simulation_outputs = []\n                for batch in batches:\n                    simulation_outputs.append(simulator(batch))\n                    pbar.update(sim_batch_size)\n\n        x = torch.cat(simulation_outputs, dim=0)\n    else:\n        x = simulator(theta)\n\n    return x\n\n\n@contextlib.contextmanager\ndef tqdm_joblib(tqdm_object):\n    \"\"\"Context manager to patch joblib to report into tqdm progress bar given as\n    argument\n\n    This wrapped context manager obtains the number of finished tasks from the tqdm\n    print function and uses it to update the pbar, as suggested in\n    https://stackoverflow.com/a/61689175. See #419, #421\n    \"\"\"\n\n    def tqdm_print_progress(self):\n        if self.n_completed_tasks > tqdm_object.n:\n            n_completed = self.n_completed_tasks - tqdm_object.n\n            tqdm_object.update(n=n_completed)\n\n    original_print_progress = joblib.parallel.Parallel.print_progress\n    joblib.parallel.Parallel.print_progress = tqdm_print_progress\n\n    try:\n        yield tqdm_object\n    finally:\n        joblib.parallel.Parallel.print_progress = original_print_progress\n        tqdm_object.close()\n", "path": "sbi/simulators/simutils.py"}]}