restyled parts of the code

Author: scarlehoff

.github/workflows/pytest.yml

@@ -9,7 +9,7 @@ jobs:
     strategy:
       max-parallel: 3
       matrix:
-        python-version: [3.9, '3.10']
+        python-version: [3.9, '3.12']
 
     steps:
     - uses: actions/checkout@v1

pyproject.toml

@@ -43,3 +43,16 @@ distribute = [
 'distributed',
 'dask-jobqueue',
 ]
+
+[tool.black]
+line-length = 100
+skip_magic_trailing_comma = true
+
+[tool.isort]
+atomic = true
+line_length = 120
+profile = "black" # https://black.readthedocs.io/en/stable/guides/using_black_with_other_tools.html#custom-configuration
+skip_gitignore = true
+known_first_party = ["bgtrees"]
+force_sort_within_sections = true
+

src/vegasflow/__init__.py

@@ -1,10 +1,10 @@
 """Monte Carlo integration with Tensorflow"""
 
-from vegasflow.configflow import int_me, float_me, run_eager, DTYPE, DTYPEINT
+from vegasflow.configflow import DTYPE, DTYPEINT, float_me, int_me, run_eager
+from vegasflow.plain import PlainFlow, plain_sampler, plain_wrapper
 
 # Expose the main interfaces
-from vegasflow.vflow import VegasFlow, vegas_wrapper, vegas_sampler
-from vegasflow.plain import PlainFlow, plain_wrapper, plain_sampler
-from vegasflow.vflowplus import VegasFlowPlus, vegasflowplus_wrapper, vegasflowplus_sampler
+from vegasflow.vflow import VegasFlow, vegas_sampler, vegas_wrapper
+from vegasflow.vflowplus import VegasFlowPlus, vegasflowplus_sampler, vegasflowplus_wrapper
 
 __version__ = "1.4.0"

src/vegasflow/configflow.py

@@ -1,9 +1,12 @@
 """
 Define some constants, header style
 """
+
+import logging
+
 # Most of this can be moved to a yaml file without loss of generality
 import os
-import logging
+
 import numpy as np
 
 # Some global parameters
@@ -94,6 +97,8 @@ def run_eager(flag=True):
 
 
 FMAX = tf.constant(np.finfo(np.float64).max, dtype=DTYPE)
+
+
 # The wrappers below transform tensors and array to the correct type
 def int_me(i):
     """Cast the input to the `DTYPEINT` type"""

src/vegasflow/monte_carlo.py

@@ -32,28 +32,29 @@
     as it is the first one found idle.
 """
 
-import inspect
-import time
+from abc import ABC, abstractmethod
 import copy
-import threading
+import inspect
 import logging
-from abc import abstractmethod, ABC
+import threading
+import time
+
 import joblib
 import numpy as np
 import tensorflow as tf
+
 from vegasflow.configflow import (
-    MAX_EVENTS_LIMIT,
     DEFAULT_ACTIVE_DEVICES,
     DTYPE,
     DTYPEINT,
+    MAX_EVENTS_LIMIT,
     TECH_CUT,
     float_me,
-    int_me,
     fone,
     fzero,
+    int_me,
 )
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -304,7 +305,8 @@ def _apply_integration_limits(self, rand):
 
     def _can_run_vectorial(self, expected_shape=None):
         """Accepting vectorial integrands depends on the algorithm,
-        if an algorithm can run on vectorial algorithms it should implement this method and return True"""
+        if an algorithm can run on vectorial algorithms it should implement this method and return True
+        """
         return self._CAN_RUN_VECTORIAL
 
     #### Integration management
@@ -561,7 +563,7 @@ def compile(self, integrand, compilable=True, signature=None, trace=False, check
                     signature = False
 
             compile_options = {
-                "experimental_autograph_options": tf.autograph.experimental.Feature.ALL,
+                "experimental_autograph_options": tf.autograph.experimental.Feature.ALL
             }
             if signature is None:
                 signature = autodiscover_signature
@@ -606,14 +608,14 @@ def batch_events(**kwargs):
                 )
             test_array = tf.random.uniform((event_size, self.n_dim), dtype=DTYPE)
             wgt = tf.random.uniform((event_size,), dtype=DTYPE)
-            res_tmp = new_integrand(test_array, weight=wgt)#.numpy()
+            res_tmp = new_integrand(test_array, weight=wgt)  # .numpy()
             res_shape = res_tmp.shape
 
             expected_shape = (event_size,)
 
             if len(res_shape) == 2:
                 self._vectorial = True
-                expected_shape = res_tmp.reshape(event_size, -1).shape
+                expected_shape = tf.reshape(res_tmp, (event_size, -1)).shape
                 if not self._can_run_vectorial(expected_shape):
                     raise NotImplementedError(
                         f"""The {self.__class__.__name__} algorithm does not support vectorial integrands

src/vegasflow/plain.py

@@ -2,10 +2,11 @@
     Plain implementation of the plainest possible MonteCarlo
 """
 
-from vegasflow.configflow import fone, fzero
-from vegasflow.monte_carlo import MonteCarloFlow, wrapper, sampler
 import tensorflow as tf
 
+from vegasflow.configflow import fone, fzero
+from vegasflow.monte_carlo import MonteCarloFlow, sampler, wrapper
+
 
 class PlainFlow(MonteCarloFlow):
     """

src/vegasflow/tests/test_algs.py

@@ -7,14 +7,16 @@
 everything works """
 import json
 import tempfile
-import pytest
+
 import numpy as np
+import pytest
+import tensorflow as tf
+
+from vegasflow import plain_sampler, vegas_sampler
 from vegasflow.configflow import DTYPE, run_eager
-from vegasflow.vflow import VegasFlow
 from vegasflow.plain import PlainFlow
+from vegasflow.vflow import VegasFlow
 from vegasflow.vflowplus import VegasFlowPlus
-from vegasflow import plain_sampler, vegas_sampler
-import tensorflow as tf
 
 # Test setup
 dim = 2

src/vegasflow/tests/test_config.py

@@ -2,11 +2,13 @@
     Test that the configuration is consistent
 """
 
+import importlib
 import os
+
 import numpy as np
-import importlib
+
 import vegasflow.configflow
-from vegasflow.configflow import DTYPE, DTYPEINT, int_me, float_me
+from vegasflow.configflow import DTYPE, DTYPEINT, float_me, int_me
 
 
 def test_int_me():

src/vegasflow/tests/test_gradients.py

@@ -4,11 +4,10 @@
 
 import numpy as np
 from pytest import mark
-
-from vegasflow import float_me, run_eager
-from vegasflow import VegasFlow, VegasFlowPlus, PlainFlow
 import tensorflow as tf
 
+from vegasflow import PlainFlow, VegasFlow, VegasFlowPlus, float_me, run_eager
+
 
 def generate_integrand(variable):
     """Generate an integrand that depends on an input variable"""

src/vegasflow/tests/test_misc.py

@@ -1,13 +1,14 @@
 """
     Miscellaneous tests that don't really fit anywhere else
 """
-import pytest
-import numpy as np
 
-from vegasflow import VegasFlow, VegasFlowPlus, PlainFlow
+import numpy as np
+import pytest
 import tensorflow as tf
 
-from .test_algs import instance_and_compile, check_is_one
+from vegasflow import PlainFlow, VegasFlow, VegasFlowPlus
+
+from .test_algs import check_is_one, instance_and_compile
 
 
 def _vector_integrand(xarr, weight=None):