xobjects 0.6.7__tar.gz → 0.6.9__tar.gz

{xobjects-0.6.7/xobjects.egg-info → xobjects-0.6.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xobjects
-Version: 0.6.7
+Version: 0.6.9
 Summary: In-memory serialization and code generator for CPU and GPU
 Home-page: https://xsuite.readthedocs.io/
 Download-URL: https://pypi.python.org/pypi/xobjects

xobjects-0.6.9/tests/test_general.py

@@ -0,0 +1,266 @@
+import os
+import subprocess
+import sys
+
+import pytest
+
+import xobjects as xo
+from xobjects.general import Print
+from xobjects.test_helpers import allow_kernel_compilation
+
+
+def test_print_mode_default(capsys):
+    printer = Print()
+
+    printer('visible')
+
+    assert capsys.readouterr().out == 'visible\n'
+
+
+def test_print_mode_suppressed(capsys):
+    printer = Print()
+    with xo.settings.override(print_mode='suppress'):
+        printer('hidden')
+
+    assert capsys.readouterr().out == ''
+
+
+def test_settings_control_xsuite_printer(capsys):
+    with xo.settings.override(print_mode='suppress'):
+        xo._print('hidden')
+
+    assert capsys.readouterr().out == ''
+
+
+def test_python_setting_overrides_environment_default(capsys):
+    with xo.settings.override(print_mode='print'):
+        xo._print('visible')
+
+    assert capsys.readouterr().out == 'visible\n'
+
+
+def test_invalid_print_mode_setting():
+    with pytest.raises(ValueError, match='XSUITE_PRINT_MODE.*print.*suppress'):
+        xo.settings.print_mode = 'invalid'
+
+
+def test_settings_override_restores_value_after_error():
+    original = xo.settings.print_mode
+
+    with pytest.raises(RuntimeError):
+        with xo.settings.override(print_mode='suppress'):
+            assert xo.settings.print_mode == 'suppress'
+            raise RuntimeError
+
+    assert xo.settings.print_mode == original
+
+
+def test_settings_are_discoverable():
+    expected_settings = {
+        'print_mode',
+        'progress_indicator',
+        'allow_kernel_compilation',
+        'force_kernel_compilation',
+        'show_kernel_diagnostics',
+        'cffi_forbid_compile',
+        'cffi_keep_build_files',
+        'cuda_backend',
+        'cuda_fast_compile',
+        'cuda_compiler',
+    }
+
+    assert expected_settings <= set(dir(xo.settings))
+    for name in expected_settings:
+        assert f'{name}=' in repr(xo.settings)
+        assert name in type(xo.settings).__doc__
+
+
+def test_print_mode_environment_variable_is_startup_default():
+    environment = os.environ.copy()
+    environment['XSUITE_PRINT_MODE'] = 'suppress'
+    code = (
+        'import xobjects as xo; '
+        'assert xo.settings.print_mode == "suppress"; '
+        'xo._print("hidden")')
+
+    completed = subprocess.run(
+        [sys.executable, '-c', code],
+        env=environment,
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+
+    assert completed.stdout == ''
+
+
+def test_print_mode_environment_variable_can_be_overridden_in_python():
+    environment = os.environ.copy()
+    environment['XSUITE_PRINT_MODE'] = 'suppress'
+    code = (
+        'import xobjects as xo; '
+        'xo.settings.print_mode = "print"; '
+        'assert xo.settings.print_mode == "print"; '
+        'xo._print("visible")')
+
+    completed = subprocess.run(
+        [sys.executable, '-c', code],
+        env=environment,
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+
+    assert completed.stdout == 'visible\n'
+
+
+def test_environment_change_after_import_does_not_override_python_setting(
+        capsys, monkeypatch):
+    with xo.settings.override(print_mode='print'):
+        monkeypatch.setenv('XSUITE_PRINT_MODE', 'suppress')
+        xo._print('visible')
+
+    assert capsys.readouterr().out == 'visible\n'
+
+
+@pytest.mark.parametrize('value', ['1', 'true', 'YES', 'on'])
+def test_boolean_environment_true_values(value):
+    environment = os.environ.copy()
+    environment['XSUITE_ALLOW_KERNEL_COMPILATION'] = value
+    code = (
+        'import xobjects as xo; '
+        'assert xo.settings.allow_kernel_compilation is True; '
+        'xo.settings.allow_kernel_compilation = False; '
+        'assert xo.settings.allow_kernel_compilation is False')
+
+    subprocess.run(
+        [sys.executable, '-c', code],
+        env=environment,
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+
+
+@pytest.mark.parametrize('value', ['0', 'false', 'NO', 'off'])
+def test_boolean_environment_false_values(value):
+    environment = os.environ.copy()
+    environment['XSUITE_FORCE_KERNEL_COMPILATION'] = value
+    code = (
+        'import xobjects as xo; '
+        'assert xo.settings.force_kernel_compilation is False')
+
+    subprocess.run(
+        [sys.executable, '-c', code],
+        env=environment,
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+
+
+def test_invalid_boolean_environment_value():
+    environment = os.environ.copy()
+    environment['XSUITE_CFFI_FORBID_COMPILE'] = 'sometimes'
+
+    completed = subprocess.run(
+        [sys.executable, '-c', 'import xobjects'],
+        env=environment,
+        capture_output=True,
+        text=True,
+    )
+
+    assert completed.returncode != 0
+    assert 'Invalid boolean value' in completed.stderr
+    assert 'xobjects.settings.cffi_forbid_compile' in completed.stderr
+    assert 'XSUITE_CFFI_FORBID_COMPILE' in completed.stderr
+
+
+def test_runtime_settings_environment_defaults():
+    environment = os.environ.copy()
+    environment.update({
+        'XSUITE_PROGRESS_INDICATOR': 'text',
+        'XSUITE_FORCE_KERNEL_COMPILATION': 'yes',
+        'XSUITE_SHOW_KERNEL_DIAGNOSTICS': 'on',
+        'XSUITE_CFFI_FORBID_COMPILE': 'true',
+        'XSUITE_CFFI_KEEP_BUILD_FILES': '1',
+        'XSUITE_CUDA_BACKEND': 'clang',
+        'XSUITE_CUDA_FAST_COMPILE': 'false',
+        'XSUITE_CUDA_COMPILER': '/path/to/clang++',
+    })
+    code = (
+        'import xobjects as xo; '
+        'assert xo.settings.progress_indicator == "text"; '
+        'assert xo.settings.force_kernel_compilation is True; '
+        'assert xo.settings.show_kernel_diagnostics is True; '
+        'assert xo.settings.cffi_forbid_compile is True; '
+        'assert xo.settings.cffi_keep_build_files is True; '
+        'assert xo.settings.cuda_backend == "clang"; '
+        'assert xo.settings.cuda_fast_compile is False; '
+        'assert xo.settings.cuda_compiler == "/path/to/clang++"')
+
+    subprocess.run(
+        [sys.executable, '-c', code],
+        env=environment,
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+
+
+@pytest.mark.parametrize(
+    'allow, force, compilation_allowed',
+    [
+        (False, False, False),
+        (True, False, True),
+        (False, True, True),
+        (True, True, True),
+    ],
+)
+def test_kernel_compilation_settings(allow, force, compilation_allowed):
+    with xo.settings.override(
+        allow_kernel_compilation=allow,
+        force_kernel_compilation=force,
+    ):
+        assert xo.context_cpu.kernel_compilation_allowed(
+            xo.ContextCpu()) is compilation_allowed
+
+
+def test_user_context_environment_variable(monkeypatch):
+    monkeypatch.setenv('XOBJECTS_USER_CONTEXT', 'ContextCpu:auto')
+    context = xo.get_user_context()
+
+    assert context.openmp_enabled
+
+
+def test_test_contexts_environment_variable(monkeypatch):
+    monkeypatch.setenv(
+        'XOBJECTS_TEST_CONTEXTS',
+        'ContextCpu;ContextCpu:auto',
+    )
+
+    contexts = list(xo.context.get_test_contexts())
+
+    assert len(contexts) == 2
+    assert contexts[0].openmp_enabled is False
+    assert contexts[1].openmp_enabled is True
+
+
+def test_cffi_forbid_compile_setting():
+    with xo.settings.override(cffi_forbid_compile=True):
+        with pytest.raises(RuntimeError) as err:
+            xo.ContextCpu().build_kernels({})
+
+    message = str(err.value)
+    assert 'xobjects.settings.cffi_forbid_compile' in message
+    assert 'XSUITE_CFFI_FORBID_COMPILE' in message
+
+
+def test_allow_kernel_compilation_decorator_restores_state():
+    @allow_kernel_compilation(skip_when_forbid_compile=False)
+    def decorated():
+        assert xo.settings.allow_kernel_compilation is True
+
+    with xo.settings.override(allow_kernel_compilation=False):
+        decorated()
+        assert xo.settings.allow_kernel_compilation is False

{xobjects-0.6.7 → xobjects-0.6.9}/tests/test_struct.py

@@ -6,7 +6,11 @@ import cffi
 import numpy as np
 
 import xobjects as xo
-from xobjects.test_helpers import for_all_test_contexts, requires_context
+from xobjects.test_helpers import (
+    allow_kernel_compilation,
+    for_all_test_contexts,
+    requires_context,
+)
 
 
 def test_static_struct_def():
@@ -299,6 +303,7 @@ double mult_four(MyStruct stru) {
 
 
 @requires_context("ContextCpu")
+@allow_kernel_compilation
 def test_compile_kernels_only_if_needed(tmp_path, mocker):
     """Test the use case of xtrack.
 
@@ -348,6 +353,7 @@ def test_compile_kernels_only_if_needed(tmp_path, mocker):
 
 
 @requires_context("ContextCpu")
+@allow_kernel_compilation
 def test_thisclass_placeholder_on_struct():
     test_context = xo.ContextCpu()
 

{xobjects-0.6.7 → xobjects-0.6.9}/xobjects/__init__.py

@@ -35,6 +35,7 @@ from .hybrid_class import JEncoder, HybridClass, MetaHybridClass
 from .linkedarray import BypassLinked
 
 from .general import _print
+from .settings import settings
 
 from .general import assert_allclose
 

xobjects-0.6.9/xobjects/_version.py

@@ -0,0 +1 @@
+__version__ = "0.6.9"

{xobjects-0.6.7 → xobjects-0.6.9}/xobjects/context.py

@@ -207,6 +207,7 @@ class ModuleNotAvailable(object):
 class XContext(ABC):
     minimum_alignment = 1
     allow_prebuilt_kernels = False
+    allow_kernel_compilation = False
 
     def __init__(self):
         self._kernels = KernelDict()
@@ -776,8 +777,8 @@ def get_test_contexts():
 
 def get_user_context():
     """
-    Get the context specfied by the enviroment variable XOBJECTS_USER_CONTEXT.
-    If not present use ContextCpu().
+    Get the context specified by the environment variable
+    ``XOBJECTS_USER_CONTEXT``. If it is not set, use ``ContextCpu()``.
 
     Examples:
        ContextPyopencl:0.0  -> ContextPyopencl(device="0.0")

{xobjects-0.6.7 → xobjects-0.6.9}/xobjects/context_cpu.py

@@ -13,42 +13,42 @@ from typing import Callable, Dict, List, Sequence, Tuple
 import weakref
 
 from .general import _print
+from .settings import settings
 
 import numpy as np
 import scipy as sp
 
-_forbid_compile = False
 _suppress_warnings = False
-allow_no_prebuilt_kernel = False
 
 
-def _class_allows_no_prebuilt_kernel(cls):
-    return (
-        getattr(cls, "allow_no_prebuilt_kernel", False)
-        or getattr(
-            getattr(cls, "_DressingClass", None),
-            "allow_no_prebuilt_kernel",
-            False,
-        )
-        or getattr(
-            getattr(cls, "_XoStruct", None), "allow_no_prebuilt_kernel", False
-        )
-    )
+def _class_allows_kernel_compilation(cls):
+    for candidate in (
+        cls,
+        getattr(cls, "_DressingClass", None),
+        getattr(cls, "_XoStruct", None),
+    ):
+        if candidate is None:
+            continue
+        if getattr(candidate, "allow_kernel_compilation", False):
+            return True
+        # Compatibility with classes from packages that have not migrated yet.
+        if getattr(candidate, "allow_no_prebuilt_kernel", False):
+            return True
+    return False
 
 
-def allow_no_prebuilt_kernel_enabled(context=None, classes=()):
+def kernel_compilation_allowed(context=None, classes=()):
     if classes is None:
         classes = ()
     elif isinstance(classes, type):
         classes = (classes,)
 
-    if os.environ.get("XSUITE_ALLOW_NO_PREBUILT_KERNELS") is not None:
-        return True
-    if allow_no_prebuilt_kernel:
+    if (settings.allow_kernel_compilation
+            or settings.force_kernel_compilation):
         return True
-    if any(_class_allows_no_prebuilt_kernel(cls) for cls in classes):
+    if any(_class_allows_kernel_compilation(cls) for cls in classes):
         return True
-    return getattr(context, "allow_no_prebuilt_kernel", False)
+    return getattr(context, "allow_kernel_compilation", False)
 
 
 def _is_serial_cpu_context(context):
@@ -58,20 +58,20 @@ def _is_serial_cpu_context(context):
 
 
 def require_prebuilt_kernel(context=None, classes=()):
-    return not allow_no_prebuilt_kernel_enabled(
+    return not kernel_compilation_allowed(
         context, classes=classes
     ) and _is_serial_cpu_context(context)
 
 
-def no_prebuilt_kernel_jit_message():
+def kernel_compilation_help_message():
     return (
         "To allow just-in-time compilation instead, as in older Xsuite "
-        "versions, set the environment variable "
-        "`XSUITE_ALLOW_NO_PREBUILT_KERNELS`, set "
-        "`xobjects.context_cpu.allow_no_prebuilt_kernel = True`, or set "
-        "`context.allow_no_prebuilt_kernel = True`. Classes that require "
+        "versions, set `xobjects.settings.allow_kernel_compilation = True`, "
+        "or equivalently set the environment variable "
+        "`XSUITE_ALLOW_KERNEL_COMPILATION=1`, or set "
+        "`context.allow_kernel_compilation = True`. Classes that require "
         "just-in-time compilation can also define "
-        "`allow_no_prebuilt_kernel = True` as a class attribute. Using "
+        "`allow_kernel_compilation = True` as a class attribute. Using "
         "just-in-time compilation instead of prebuilt kernels may require "
         "lengthy compilation whenever a different kernel is needed."
     )
@@ -369,13 +369,11 @@ class ContextCpu(XContext):
             cdefs = "\n".join(cls._gen_c_decl({}) for cls in classes)
             cdefs += "\n" + extra_cdef
 
-            if _forbid_compile:
-                raise RuntimeError("Compilation is forbidden")
-
-            if os.environ.get("XOBJECTS_FORBID_COMPILE"):
+            if settings.cffi_forbid_compile:
                 raise RuntimeError(
-                    "Compilation is forbidden by the environment variable "
-                    "XOBJECTS_FORBID_COMPILE"
+                    "CFFI compilation is forbidden by "
+                    "xobjects.settings.cffi_forbid_compile or equivalently "
+                    "the environment variable XSUITE_CFFI_FORBID_COMPILE."
                 )
 
             so_file = self.compile_kernel(
@@ -536,7 +534,7 @@ class ContextCpu(XContext):
             return Path(output_file)
         finally:
             # Clean temp files
-            if "XOBJECTS_KEEP_BUILD_FILES" not in os.environ:
+            if not settings.cffi_keep_build_files:
                 files_to_remove = [
                     module_name + ".c",
                     module_name + ".o",
@@ -961,8 +959,8 @@ class FFTCpu(object):
                 direction="FFTW_BACKWARD",
                 flags=("FFTW_MEASURE",),
             )
-            print(f"fftw simd_aligned={self.fftw.simd_aligned}")
-            print(f"ifftw simd_aligned={self.fftw.simd_aligned}")
+            _print(f"fftw simd_aligned={self.fftw.simd_aligned}")
+            _print(f"ifftw simd_aligned={self.fftw.simd_aligned}")
         else:
             # I perform one fft to have numpy cache the plan
             _ = np.fft.ifftn(np.fft.fftn(data, axes=axes), axes=axes)

{xobjects-0.6.7 → xobjects-0.6.9}/xobjects/context_cupy.py

@@ -25,23 +25,21 @@ from .context import (
     sources_from_classes,
 )
 from .linkedarray import BaseLinkedArray
+from .settings import settings
 from .specialize_source import specialize_source
 
 log = logging.getLogger(__name__)
 
-no_fast_compile = False
-"""Disable NVRTC fast compile tuning when building CUDA kernels.
-
-When set to ``True``, ``ContextCupy`` does not pass ``--Ofast-compile=min``
-to NVRTC. The ``XO_CUDA_NO_FAST_COMPILE`` environment variable provides the
-same behavior.
-"""
-
 try:
     import cupy
     import cupyx.scipy
     import cupyx.scipy.interpolate
-    import cupyx.scipy.signal
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            'ignore',
+            message=r'cupyx\.jit\.rawkernel is experimental\.',
+            category=FutureWarning)
+        import cupyx.scipy.signal
     import cupyx.scipy.special
     import cupyx.scipy.stats
     from cupyx.scipy import fftpack as cufftp
@@ -394,12 +392,12 @@ class ContextCupy(XContext):
     Creates a Cupy Context object, that allows performing the computations
     on nVidia GPUs.
 
-    The module-level flag ``xobjects.context_cupy.no_fast_compile`` controls
-    whether NVRTC fast compile tuning is disabled. By default it is ``False``,
-    so CUDA kernels built with NVRTC >= 12.9 use ``--Ofast-compile=min`` to
-    reduce compilation time and memory usage, at the cost of some runtime
-    performance. Set it to ``True`` to disable this option. The environment
-    variable ``XO_CUDA_NO_FAST_COMPILE`` also disables it.
+    ``xobjects.settings.cuda_fast_compile`` controls whether NVRTC fast compile
+    tuning is enabled. By default it is ``True``, so CUDA kernels built with
+    NVRTC >= 12.9 use ``--Ofast-compile=min`` to reduce compilation time and
+    memory usage, at the cost of some runtime performance. Set
+    ``xobjects.settings.cuda_fast_compile = False``, or equivalently the
+    environment variable ``XSUITE_CUDA_FAST_COMPILE=0``, to disable it.
 
     Args:
         default_block_size (int):  CUDA thread size that is used by default
@@ -434,8 +432,8 @@ class ContextCupy(XContext):
         self.default_block_size = default_block_size
         self.default_shared_mem_size_bytes = default_shared_mem_size_bytes
 
-        if not backend:
-            backend = os.environ.get("XO_CUDA_BACKEND", "nvrtc")
+        if backend is None:
+            backend = settings.cuda_backend
 
         if backend not in ["nvrtc", "clang"]:
             raise ValueError(
@@ -511,9 +509,7 @@ class ContextCupy(XContext):
         if self.backend == "nvrtc":
             # NVRTC (default): add NVRTC-specific flags
             nvrtc_args = (*extra_compile_args,)
-            fast_compile = not (
-                no_fast_compile or os.environ.get("XO_CUDA_NO_FAST_COMPILE")
-            )
+            fast_compile = settings.cuda_fast_compile
             if nvrtc and nvrtc.getVersion() >= (12, 9):
                 # If supported, skip prohibitively heavy optimisations (e.g.
                 # involving cloning). This it at the expense of <20%
@@ -553,7 +549,7 @@ class ContextCupy(XContext):
         return out_kernels
 
     def _find_clang(self):
-        override = os.environ.get("XO_CUDA_CLANG")
+        override = settings.cuda_compiler
         if override:
             return override
 
@@ -562,8 +558,11 @@ class ContextCupy(XContext):
             return found
 
         raise RuntimeError(
-            "clang++ for the CUDA context not found. Either install clang so that 'clang++' is on PATH,"
-            "or set the XO_CUDA_CLANG variable to the desired clang++ executable."
+            "clang++ for the CUDA context not found. Either install clang so "
+            "that 'clang++' is on PATH, or set "
+            "xobjects.settings.cuda_compiler, or equivalently the environment "
+            "variable XSUITE_CUDA_COMPILER, to the desired clang++ "
+            "executable."
         )
 
     def _build_module_with_clang(self, source, extra_compile_args=()):

{xobjects-0.6.7 → xobjects-0.6.9}/xobjects/context_pyopencl.py

@@ -21,6 +21,7 @@ from .context import (
 )
 from .linkedarray import BaseLinkedArray
 from .specialize_source import specialize_source
+from .general import _print
 
 log = logging.getLogger(__name__)
 
@@ -100,9 +101,9 @@ class ContextPyopencl(XContext):
     @classmethod
     def print_devices(cls):
         for ip, platform in enumerate(cl.get_platforms()):
-            print(f"Platform {ip}  : {platform.name}")
+            _print(f"Platform {ip}  : {platform.name}")
             for id, device in enumerate(platform.get_devices()):
-                print(f"Device   {ip}.{id}: {device.name}")
+                _print(f"Device   {ip}.{id}: {device.name}")
 
     def __init__(
         self, device=None, patch_pyopencl_array=True, minimum_alignment=None

{xobjects-0.6.7 → xobjects-0.6.9}/xobjects/general.py

@@ -5,13 +5,21 @@
 from numpy.testing import assert_allclose as np_assert_allclose
 import numpy as np
 
+from .settings import settings
+
 
 class Print:
-    suppress = False
+    """Configurable wrapper around :func:`print` used by Xsuite.
+
+    The behavior is controlled by ``xobjects.settings.print_mode``, or
+    equivalently the environment variable ``XSUITE_PRINT_MODE``.
+    """
 
     def __call__(self, *args, **kwargs):
-        if not self.suppress:
-            print(*args, **kwargs)
+        if settings.print_mode == 'suppress':
+            return
+
+        print(*args, **kwargs)
 
 
 _print = Print()