xobjects 0.4.7__tar.gz → 0.5.1__tar.gz

{xobjects-0.4.7/xobjects.egg-info → xobjects-0.5.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xobjects
-Version: 0.4.7
+Version: 0.5.1
 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.5.1/xobjects/_version.py

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

{xobjects-0.4.7 → xobjects-0.5.1}/xobjects/context.py

@@ -254,6 +254,7 @@ class XContext(ABC):
         extra_classes: Sequence[Type] = (),
         extra_headers: Sequence[SourceType] = (),
         compile: bool = True,  # noqa
+        extra_compile_args: Sequence[str] = (),
     ):
         """
         Adds user-defined kernels to the context. The kernel source
@@ -333,6 +334,7 @@ class XContext(ABC):
             extra_classes=extra_classes,
             extra_headers=extra_headers,
             compile=compile,
+            extra_compile_args=extra_compile_args,
         )
         self.kernels.update(generated_kernels)
 
@@ -348,6 +350,7 @@ class XContext(ABC):
         extra_classes: Sequence[Type],
         extra_headers: Sequence[SourceType],
         compile: bool,
+        extra_compile_args: Sequence[str],
     ) -> Dict[Tuple[str, tuple], KernelType]:
         pass
 

{xobjects-0.4.7 → xobjects-0.5.1}/xobjects/context_cpu.py

@@ -169,8 +169,8 @@ class ContextCpu(XContext):
         specialize=True,
         apply_to_source=(),
         save_source_as=None,
-        extra_compile_args: Sequence[str] = ("-O3", "-Wno-unused-function"),
-        extra_link_args: Sequence[str] = ("-O3",),
+        extra_compile_args: Sequence[str] = (),
+        extra_link_args: Sequence[str] = (),
         extra_cdef="",
         extra_classes=(),
         extra_headers=(),
@@ -271,13 +271,16 @@ class ContextCpu(XContext):
         specialize=True,
         apply_to_source=(),
         save_source_as=None,
-        extra_compile_args=("-O3", "-Wno-unused-function"),
-        extra_link_args=("-O3",),
+        extra_compile_args=(),
+        extra_link_args=(),
         extra_cdef="",
         extra_classes=(),
         extra_headers=(),
         compile=True,  # noqa
     ) -> Dict[Tuple[str, tuple], "KernelCpu"]:
+        extra_compile_args += ("-O3", "-Wno-unused-function")
+        extra_link_args += ("-O3",)
+
         # Determine names and paths
         clean_up_so = not module_name
         module_name = module_name or str(uuid.uuid4().hex)
@@ -409,20 +412,19 @@ class ContextCpu(XContext):
             ffi_interface.cdef("int omp_get_max_threads();")
 
         # Compile
-        xtr_compile_args = ["-std=c99"]
-        xtr_link_args = ["-std=c99"]
+        xtr_compile_args = ["-std=c99", "-DXO_CONTEXT_CPU"]
+        xtr_link_args = ["-std=c99", "-DXO_CONTEXT_CPU"]
         xtr_compile_args += extra_compile_args
         xtr_link_args += extra_link_args
 
         if self.openmp_enabled:
             xtr_compile_args.append("-fopenmp")
             xtr_link_args.append("-fopenmp")
-
-            # https://mac.r-project.org/openmp/
-            # on macos comment the above and uncomment the below flags to compile OpenMP with Xcode clang:
-            # xtr_compile_args.append("-Xclang")
-            # xtr_compile_args.append("-fopenmp")
-            # xtr_link_args.append("-lomp")
+            xtr_compile_args.append("-DXO_CONTEXT_CPU_OPENMP")
+            xtr_link_args.append("-DXO_CONTEXT_CPU_OPENMP")
+        else:
+            xtr_compile_args.append("-DXO_CONTEXT_CPU_SERIAL")
+            xtr_link_args.append("-DXO_CONTEXT_CPU_SERIAL")
 
         if os.name == "nt":  # windows
             # TODO: to be handled properly

{xobjects-0.4.7 → xobjects-0.5.1}/xobjects/context_cupy.py

@@ -416,6 +416,7 @@ class ContextCupy(XContext):
         specialize=True,
         apply_to_source=(),
         save_source_as=None,
+        extra_compile_args=(),
         extra_cdef=None,
         extra_classes=(),
         extra_headers=(),
@@ -454,7 +455,10 @@ class ContextCupy(XContext):
             with open(save_source_as, "w") as fid:
                 fid.write(specialized_source)
 
-        module = cupy.RawModule(code=specialized_source)
+        extra_compile_args = (*extra_compile_args, "-DXO_CONTEXT_CUDA")
+        module = cupy.RawModule(
+            code=specialized_source, options=extra_compile_args
+        )
 
         out_kernels = {}
         for pyname, kernel in kernel_descriptions.items():

{xobjects-0.4.7 → xobjects-0.5.1}/xobjects/context_pyopencl.py

@@ -180,6 +180,7 @@ class ContextPyopencl(XContext):
         specialize=True,
         apply_to_source=(),
         save_source_as=None,
+        extra_compile_args=(),
         extra_cdef=None,
         extra_classes=(),
         extra_headers=(),
@@ -218,8 +219,13 @@ class ContextPyopencl(XContext):
             with open(save_source_as, "w") as fid:
                 fid.write(specialized_source)
 
+        extra_compile_args = (
+            *extra_compile_args,
+            "-cl-std=CL2.0",
+            "-DXO_CONTEXT_CL",
+        )
         prg = cl.Program(self.context, specialized_source).build(
-            options="-cl-std=CL2.0",
+            options=extra_compile_args,
         )
 
         out_kernels = {}

xobjects-0.5.1/xobjects/general.py

@@ -0,0 +1,52 @@
+# copyright ################################# #
+# This file is part of the Xobjects Package.  #
+# Copyright (c) CERN, 2024.                   #
+# ########################################### #
+from numpy.testing import assert_allclose as np_assert_allclose
+import numpy as np
+
+
+class Print:
+    suppress = False
+
+    def __call__(self, *args, **kwargs):
+        if not self.suppress:
+            print(*args, **kwargs)
+
+
+_print = Print()
+
+
+def assert_allclose(a, b, rtol=0, atol=0, max_outliers=0):
+    if hasattr(a, "get"):
+        a = a.get()
+    if hasattr(b, "get"):
+        b = b.get()
+    try:
+        a = np.squeeze(a)
+    except:
+        pass
+    try:
+        b = np.squeeze(b)
+    except:
+        pass
+
+    try:
+        np_assert_allclose(a, b, rtol=rtol, atol=atol)
+    except AssertionError as e:
+        if max_outliers == 0:
+            raise e
+        if not allclose_with_outliers(a, b, rtol, atol, max_outliers):
+            raise AssertionError(
+                "Arrays are not close enough, even with outliers allowed."
+            ) from e
+
+
+def allclose_with_outliers(a, b, rtol=1e-7, atol=0, max_outliers=0):
+    a = np.asanyarray(a)
+    b = np.asanyarray(b)
+    diff = np.abs(a - b)
+    allowed = atol + rtol * np.abs(b)
+    mask = diff > allowed
+    num_outliers = np.count_nonzero(mask)
+    return num_outliers <= max_outliers

{xobjects-0.4.7 → xobjects-0.5.1}/xobjects/hybrid_class.py

@@ -96,7 +96,7 @@ class JEncoder(json.JSONEncoder):
 def _build_xofields_dict(bases, data):
     if "_xofields" in data.keys():
         xofields = data["_xofields"].copy()
-    elif any(map(lambda b: hasattr(b, "_xofields"), bases)):
+    elif any(hasattr(b, "_xofields") for b in bases):
         n_filled = 0
         for bb in bases:
             if hasattr(bb, "_xofields") and len(bb._xofields.keys()) > 0:

{xobjects-0.4.7 → xobjects-0.5.1}/xobjects/struct.py

@@ -473,6 +473,7 @@ class Struct(metaclass=MetaStruct):
         apply_to_source=(),
         save_source_as=None,
         extra_classes=(),
+        extra_compile_args=(),
     ):
         if only_if_needed:
             all_found = True
@@ -489,6 +490,7 @@ class Struct(metaclass=MetaStruct):
             extra_classes=[cls] + list(extra_classes),
             apply_to_source=apply_to_source,
             save_source_as=save_source_as,
+            extra_compile_args=extra_compile_args,
         )
 
     def compile_kernels(
@@ -497,6 +499,7 @@ class Struct(metaclass=MetaStruct):
         apply_to_source=(),
         save_source_as=None,
         extra_classes=(),
+        extra_compile_args=(),
     ):
         self.compile_class_kernels(
             context=self._context,
@@ -504,6 +507,7 @@ class Struct(metaclass=MetaStruct):
             apply_to_source=apply_to_source,
             save_source_as=save_source_as,
             extra_classes=extra_classes,
+            extra_compile_args=extra_compile_args,
         )
 
 

{xobjects-0.4.7 → xobjects-0.5.1/xobjects.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xobjects
-Version: 0.4.7
+Version: 0.5.1
 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.4.7/xobjects/_version.py

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

xobjects-0.4.7/xobjects/general.py

@@ -1,33 +0,0 @@
-# copyright ################################# #
-# This file is part of the Xobjects Package.  #
-# Copyright (c) CERN, 2024.                   #
-# ########################################### #
-from numpy.testing import assert_allclose as np_assert_allclose
-import numpy as np
-
-
-class Print:
-    suppress = False
-
-    def __call__(self, *args, **kwargs):
-        if not self.suppress:
-            print(*args, **kwargs)
-
-
-_print = Print()
-
-
-def assert_allclose(a, b, rtol=1e-7, atol=1e-7):
-    if hasattr(a, "get"):
-        a = a.get()
-    if hasattr(b, "get"):
-        b = b.get()
-    try:
-        a = np.squeeze(a)
-    except:
-        pass
-    try:
-        b = np.squeeze(b)
-    except:
-        pass
-    np_assert_allclose(a, b, rtol=rtol, atol=atol)