xobjects 0.5.7__tar.gz → 0.5.9__tar.gz

{xobjects-0.5.7 → xobjects-0.5.9}/LICENSE

@@ -198,4 +198,4 @@
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
-   limitations under the License.
+   limitations under the License.

xobjects-0.5.9/PKG-INFO

@@ -0,0 +1,39 @@
+Metadata-Version: 2.4
+Name: xobjects
+Version: 0.5.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
+Author: Riccardo De Maria
+Author-email: Riccardo De Maria <riccardo.de.maria@cern.ch>
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://xsuite.readthedocs.io/
+Project-URL: Bug Tracker, https://github.com/xsuite/xsuite/issues
+Project-URL: Documentation, https://xsuite.readthedocs.io/
+Project-URL: Source Code, https://github.com/xsuite/xobjects
+Project-URL: Download, https://pypi.python.org/pypi/xobjects
+Requires-Python: >=3.7
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: cffi
+Requires-Dist: scipy
+Provides-Extra: tests
+Requires-Dist: pytest; extra == "tests"
+Requires-Dist: pytest-mock; extra == "tests"
+Dynamic: author
+Dynamic: download-url
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: requires-python
+
+# XObjects
+
+Library to create and manipulate serialized object in CPU and GPU memory efficiently.
+
+The library defines:
+-  Contexts: create buffers and manage computations e.g. BytearrayContext, CLContext, ...
+-  Buffers: reference to internal memory buffer with an allocator e.g.
+   BytearrayContext, CLBuffer, ...
+-  Data types: define users define types to create and manipulate objects on
+   buffers and generate C api, e.g. Float64, Int64, String, Struct, Array, Union

xobjects-0.5.9/README.md

@@ -0,0 +1,10 @@
+# XObjects
+
+Library to create and manipulate serialized object in CPU and GPU memory efficiently.
+
+The library defines:
+-  Contexts: create buffers and manage computations e.g. BytearrayContext, CLContext, ...
+-  Buffers: reference to internal memory buffer with an allocator e.g.
+   BytearrayContext, CLBuffer, ...
+-  Data types: define users define types to create and manipulate objects on
+   buffers and generate C api, e.g. Float64, Int64, String, Struct, Array, Union

xobjects-0.5.9/pyproject.toml

@@ -0,0 +1,50 @@
+[build-system]
+requires = [
+    "setuptools>=77.0"
+]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "xobjects"
+dynamic = ["version"]
+description = "In-memory serialization and code generator for CPU and GPU"
+readme = "README.md"
+authors = [
+    { name = "Riccardo De Maria", email = "riccardo.de.maria@cern.ch" }
+]
+license = "Apache-2.0"
+requires-python = ">=3.9"
+dependencies = [
+    "numpy",
+    "cffi",
+    "scipy"
+]
+[project.urls]
+Homepage = "https://xsuite.readthedocs.io/"
+"Bug Tracker" = "https://github.com/xsuite/xsuite/issues"
+Documentation = "https://xsuite.readthedocs.io/"
+"Source Code" = "https://github.com/xsuite/xobjects"
+"Download" = "https://pypi.python.org/pypi/xobjects"
+
+[project.optional-dependencies]
+tests = ["pytest", "pytest-mock"]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["xobjects*"]
+
+[tool.setuptools.dynamic]
+version = {attr = "xobjects._version.__version__"}
+
+[tool.black]
+line-length = 79
+target-version = ['py310', 'py311', 'py312']
+include = '\.pyi?$'
+
+[project.entry-points.xobjects]
+include = "xobjects"
+
+[pytest]
+markers = [
+    "context_dependent: marks test as one that depends on the execution context",
+]

{xobjects-0.5.7 → xobjects-0.5.9}/tests/test_capi.py

@@ -11,7 +11,7 @@ import numpy as np
 import cffi
 
 import xobjects as xo
-
+from xobjects.test_helpers import for_all_test_contexts
 
 ffi = cffi.FFI()
 
@@ -158,6 +158,70 @@ def test_array_dynamic_type_init_get_set(array_cls, example_shape):
             assert arr[ii].field1[idx_in_field] == 13 * vv
 
 
+@for_all_test_contexts
+@pytest.mark.parametrize(
+    "array_type",
+    [
+        xo.UInt64[3, 5, 7],
+        xo.UInt64[:, :, :],
+        xo.UInt64[:, 5, :],
+    ],
+)
+def test_array_get_shape(test_context, array_type):
+    source = """
+        #include "xobjects/headers/common.h"
+
+        GPUKERN void get_nd_and_shape(
+            ARRAY_TYPE arr,
+            GPUGLMEM int64_t* out_nd,
+            GPUGLMEM int64_t* out_shape
+        ) {
+            *out_nd = ARRAY_TYPE_nd(arr);
+            ARRAY_TYPE_shape(arr, out_shape);
+        }
+    """.replace(
+        "ARRAY_TYPE", array_type.__name__
+    )
+
+    kernels = {
+        "get_nd_and_shape": xo.Kernel(
+            c_name="get_nd_and_shape",
+            args=[
+                xo.Arg(array_type, name="arr"),
+                xo.Arg(xo.Int64, pointer=True, name="out_nd"),
+                xo.Arg(xo.Int64, pointer=True, name="out_shape"),
+            ],
+        ),
+    }
+
+    test_context.add_kernels(
+        sources=[source],
+        kernels=kernels,
+    )
+
+    instance = array_type(
+        np.array(range(3 * 5 * 7)).reshape((3, 5, 7)),
+        _context=test_context,
+    )
+
+    expected_nd = 3
+    result_nd = test_context.zeros((1,), dtype=np.int64)
+
+    expected_shape = [3, 5, 7]
+    result_shape = test_context.zeros((expected_nd,), dtype=np.int64)
+
+    test_context.kernels.get_nd_and_shape(
+        arr=instance,
+        out_nd=result_nd,
+        out_shape=result_shape,
+    )
+
+    assert result_nd[0] == expected_nd
+    assert result_shape[0] == expected_shape[0]
+    assert result_shape[1] == expected_shape[1]
+    assert result_shape[2] == expected_shape[2]
+
+
 def test_struct1():
     kernels = Struct1._gen_kernels()
     ctx = xo.ContextCpu()
@@ -539,45 +603,190 @@ def test_getp1_dyn_length_dyn_type_string_array():
         assert ord(ffi.cast("char *", s2)[8 + ii]) == ch
 
 
-def test_gpu_api():
-    for ctx in xo.context.get_test_contexts():
-        src_code = """
-        /*gpufun*/
-        void myfun(double x, double y,
-            double* z){
-            z[0] = x * y;
-            }
+@for_all_test_contexts
+def test_gpu_api(test_context):
+    src_code = """
+    /*gpufun*/
+    void myfun(double x, double y,
+        double* z){
+        z[0] = x * y;
+        }
 
-        /*gpukern*/
-        void my_mul(const int n,
-            /*gpuglmem*/ const double* x1,
-            /*gpuglmem*/ const double* x2,
-            /*gpuglmem*/       double* y) {
-            int tid = 0 //vectorize_over tid n
-            double z;
-            myfun(x1[tid], x2[tid], &z);
-            y[tid] = z;
-            //end_vectorize
+    /*gpukern*/
+    void my_mul(const int n,
+        /*gpuglmem*/ const double* x1,
+        /*gpuglmem*/ const double* x2,
+        /*gpuglmem*/       double* y) {
+        int tid = 0 //vectorize_over tid n
+        double z;
+        myfun(x1[tid], x2[tid], &z);
+        y[tid] = z;
+        //end_vectorize
+        }
+    """
+
+    kernel_descriptions = {
+        "my_mul": xo.Kernel(
+            args=[
+                xo.Arg(xo.Int32, name="n"),
+                xo.Arg(xo.Float64, pointer=True, const=True, name="x1"),
+                xo.Arg(xo.Float64, pointer=True, const=True, name="x2"),
+                xo.Arg(xo.Float64, pointer=True, const=False, name="y"),
+            ],
+            n_threads="n",
+        ),
+    }
+
+    test_context.add_kernels(
+        sources=[src_code],
+        kernels=kernel_descriptions,
+        save_source_as=None,
+        compile=True,
+        extra_classes=[xo.String[:]],
+    )
+
+
+@for_all_test_contexts
+def test_array_of_arrays(test_context):
+    cell_ids = [3, 5, 7]
+    particle_per_cell = [
+        [1, 8],
+        [9, 3, 2],
+        [4, 5, 6, 7],
+    ]
+
+    class Cells(xo.Struct):
+        ids = xo.Int64[:]
+        particles = xo.Int64[:][:]
+
+    cells = Cells(
+        ids=cell_ids, particles=particle_per_cell, _context=test_context
+    )
+
+    # Data layout (displayed as uint64):
+    #
+    #   [0] 216 (cells size)
+    #   [8]  56 (offset field 2 -- particles field)
+    #  [16] cell_ids data:
+    #        [0] 40 (cell_ids size)
+    #        [8]  3 (cell_ids length)
+    #       [16] {3, 5, 7} (cell_ids elements)
+    #  [56] particles data:
+    #        [0] 160 (particles size)
+    #        [8]   3 (particles length)
+    #       [16]  40 (offset particles[0])
+    #       [24]  72 (offset particles[1])
+    #       [32] 112 (offset particles[2])
+    #       [40] particles[0] data:
+    #             [0] 32 (particles[0] size)
+    #             [8]  2 (particles[0] length)
+    #            [16] {1, 8} (particles[0] elements)
+    #       [72] particles[1] data:
+    #             [0] 40 (particles[1] size)
+    #             [8]  3 (particles[1] length)
+    #            [16] {9, 3, 2} (particles[1
+    #      [112] particles[2] data:
+    #             [0] 48 (particles[2] size)
+    #             [8]  4 (particles[2] length)
+    #            [16] {4, 5, 6, 7} (particles[2] elements)
+
+    src = r"""
+    #include "xobjects/headers/common.h"
+
+    static const int MAX_PARTICLES = 4;
+    static const int MAX_CELLS = 3;
+
+    GPUKERN void loop_over(
+        Cells cells,
+        GPUGLMEM uint64_t* out_counts,
+        GPUGLMEM uint64_t* out_vals,
+        GPUGLMEM uint8_t* success
+    )
+    {
+        int64_t num_cells = Cells_len_ids(cells);
+
+        for (int64_t i = 0; i < num_cells; i++) {
+            int64_t id = Cells_get_ids(cells, i);
+            int64_t count = Cells_len1_particles(cells, i);
+
+            if (i >= MAX_CELLS) {
+                *success = 0;
+                continue;
             }
-        """
-
-        kernel_descriptions = {
-            "my_mul": xo.Kernel(
-                args=[
-                    xo.Arg(xo.Int32, name="n"),
-                    xo.Arg(xo.Float64, pointer=True, const=True, name="x1"),
-                    xo.Arg(xo.Float64, pointer=True, const=True, name="x2"),
-                    xo.Arg(xo.Float64, pointer=True, const=False, name="y"),
-                ],
-                n_threads="n",
-            ),
+
+            out_counts[i] = count;
+
+            ArrNInt64 particles = Cells_getp1_particles(cells, i);
+            uint32_t num_particles = ArrNInt64_len(particles);
+
+            VECTORIZE_OVER(j, num_particles);
+                int64_t val = ArrNInt64_get(particles, j);
+
+                if (j >= MAX_PARTICLES) {
+                    *success = 0;
+                } else {
+                    out_vals[i * MAX_PARTICLES + j] = val;
+                }
+            END_VECTORIZE;
         }
+    }
+
+    GPUKERN void kernel_Cells_get_particles(
+        Cells obj,
+        int64_t i0,
+        int64_t i1,
+        GPUGLMEM int64_t* out
+    ) {
+        *out = Cells_get_particles(obj, i0, i1);
+    }
+    """
+
+    kernels = {
+        "loop_over": xo.Kernel(
+            args=[
+                xo.Arg(Cells, name="cells"),
+                xo.Arg(xo.UInt64, pointer=True, name="out_counts"),
+                xo.Arg(xo.UInt64, pointer=True, name="out_vals"),
+                xo.Arg(xo.UInt8, pointer=True, name="success"),
+            ],
+            n_threads=4,
+        ),
+        "kernel_Cells_get_particles": xo.Kernel(
+            args=[
+                xo.Arg(Cells, name="obj"),
+                xo.Arg(xo.Int64, name="i0"),
+                xo.Arg(xo.Int64, name="i1"),
+                xo.Arg(xo.Int64, pointer=True, name="out"),
+            ],
+        ),
+    }
+
+    test_context.add_kernels(
+        sources=[src],
+        kernels=kernels,
+    )
 
-        ctx.add_kernels(
-            sources=[src_code],
-            kernels=kernel_descriptions,
-            # save_src_as=f'_test_{name}.c')
-            save_source_as=None,
-            compile=True,
-            extra_classes=[xo.String[:]],
-        )
+    counts = test_context.zeros(len(cell_ids), dtype=np.uint64)
+    vals = test_context.zeros(12, dtype=np.uint64)
+    success = test_context.zeros((1,), dtype=np.uint8) + 1
+
+    for i, _ in enumerate(particle_per_cell):
+        for j, expected in enumerate(particle_per_cell[i]):
+            result = test_context.zeros(shape=(1,), dtype=np.int64)
+            test_context.kernels.kernel_Cells_get_particles(
+                obj=cells, i0=i, i1=j, out=result
+            )
+            assert result[0] == expected
+
+    test_context.kernels.loop_over(
+        cells=cells,
+        out_counts=counts,
+        out_vals=vals,
+        success=success,
+    )
+    counts = test_context.nparray_from_context_array(counts)
+    vals = test_context.nparray_from_context_array(vals)
+
+    assert success[0] == 1
+    assert np.all(counts == [2, 3, 4])
+    assert np.all(vals == [1, 8, 0, 0, 9, 3, 2, 0, 4, 5, 6, 7])

xobjects-0.5.9/tests/test_common.py

@@ -0,0 +1,179 @@
+# copyright ################################# #
+# This file is part of the Xobjects Package.  #
+# Copyright (c) CERN, 2025.                   #
+# ########################################### #
+
+import xobjects as xo
+import numpy as np
+import pytest
+
+from xobjects.test_helpers import for_all_test_contexts
+
+
+@for_all_test_contexts
+def test_common_atomicadd(test_context):
+    src = r"""
+    #include "xobjects/headers/common.h"
+    #include "xobjects/headers/atomicadd.h"
+
+    GPUKERN void test_atomic_add(GPUGLMEM double* out, int32_t iterations)
+    {
+        VECTORIZE_OVER(i, iterations);
+            // If on CPU do some work to avoid the loop being optimized out
+            #if defined(XO_CONTEXT_CPU_OPENMP)
+                usleep(10);
+            #endif
+            atomicAdd(out, 1.0);
+        END_VECTORIZE;
+    }
+    """
+
+    n_threads = 1
+    if type(test_context).__name__ in {"ContextCupy", "ContextPyopencl"}:
+        n_threads = 1000
+    elif (
+        test_context.omp_num_threads == "auto"
+        or test_context.omp_num_threads > 1
+    ):
+        n_threads = 8
+
+    test_context.add_kernels(
+        sources=[src],
+        kernels={
+            "test_atomic_add": xo.Kernel(
+                c_name="test_atomic_add",
+                args=[
+                    xo.Arg(xo.Float64, pointer=True, name="out"),
+                    xo.Arg(xo.Int32, name="iterations"),
+                ],
+                n_threads=n_threads,
+            )
+        },
+    )
+
+    expected = 1000
+    result = np.array([0], dtype=np.float64)
+    result_ctx = test_context.nparray_to_context_array(result)
+    test_context.kernels.test_atomic_add(out=result_ctx, iterations=expected)
+    result = test_context.nparray_from_context_array(result_ctx)
+
+    assert result == expected
+
+
+@for_all_test_contexts
+@pytest.mark.parametrize(
+    "overload", [True, False], ids=["overload", "no_overload"]
+)
+@pytest.mark.parametrize(
+    "ctype",
+    [
+        xo.Int8,
+        xo.Int16,
+        xo.Int32,
+        xo.Int64,
+        xo.UInt8,
+        xo.UInt16,
+        xo.UInt32,
+        xo.UInt64,
+        xo.Float32,
+        xo.Float64,
+    ],
+)
+def test_atomic(overload, ctype, test_context):
+    if overload:
+        func_name = "atomicAdd"
+    else:
+        func_name = f'atomicAdd_{ctype.__name__.lower()[0]}{ctype.__name__.split("t")[1]}'
+
+    class TestAtomic(xo.Struct):
+        val = ctype
+        _extra_c_sources = [
+            f"""
+            #include "xobjects/headers/common.h"
+            #include "xobjects/headers/atomicadd.h"
+
+            GPUKERN
+            void run_atomic_test(TestAtomic el, GPUGLMEM {ctype._c_type}* increments,
+                                 GPUGLMEM {ctype._c_type}* retvals, int length) {{
+                VECTORIZE_OVER(ii, length);
+                    GPUGLMEM {ctype._c_type}* val = TestAtomic_getp_val(el);
+                    {ctype._c_type} ret = {func_name}(val, increments[ii]);
+                    retvals[ii] = ret;
+                END_VECTORIZE;
+            }}
+            """
+        ]
+
+    kernels = {
+        "run_atomic_test": xo.Kernel(
+            c_name="run_atomic_test",
+            args=[
+                xo.Arg(TestAtomic, name="el"),
+                xo.Arg(ctype, pointer=True, name="increments"),
+                xo.Arg(ctype, pointer=True, name="retvals"),
+                xo.Arg(xo.Int32, name="length"),
+            ],
+            n_threads="length",
+        )
+    }
+
+    atomic = TestAtomic(_context=test_context, val=0)
+    assert atomic.val == 0
+
+    test_context.add_kernels(kernels=kernels)
+
+    # Test with all increments = 1, so we can check the return values easily.
+    num_steps = 10000
+    if ctype.__name__.startswith("Int") or ctype.__name__.startswith("Uint"):
+        # Less steps to avoid overflow
+        num_steps = min(num_steps, 2 ** (8 * ctype._size - 1) - 1)
+    increments = test_context.zeros(shape=(num_steps,), dtype=ctype._dtype) + 1
+    retvals = test_context.zeros(shape=(num_steps,), dtype=ctype._dtype)
+    test_context.kernels.run_atomic_test(
+        el=atomic, increments=increments, retvals=retvals, length=num_steps
+    )
+    assert atomic.val == num_steps
+    retvals = np.sort(test_context.nparray_from_context_array(retvals))
+    assert np.allclose(
+        retvals,
+        np.arange(num_steps, dtype=ctype._dtype),
+        atol=1.0e-15,
+        rtol=1.0e-15,
+    )
+
+    # Test with random increments, where we now only can check the total sum
+    # (retvals can be anything). Watch out: overflow is undefined behaviour,
+    # except for unsigned integers, so we skip this test for signed integers.
+    atomic.val = 0
+    retvals = test_context.zeros(shape=(num_steps,), dtype=ctype._dtype)
+    if ctype.__name__.startswith("Uint"):
+        low = 0
+        high = 2 ** (8 * ctype._size) - 1
+        increments = np.random.randint(
+            low, high + 1, size=num_steps, dtype=ctype._dtype
+        )
+        increments = test_context.nparray_to_context_array(increments)
+        test_context.kernels.run_atomic_test(
+            el=atomic, increments=increments, retvals=retvals, length=num_steps
+        )
+        increments = test_context.nparray_from_context_array(increments)
+        assert atomic.val == (
+            np.sum(increments).item() % (2 ** (8 * ctype._size))
+        )
+
+    elif ctype.__name__.startswith("Float"):
+        increments = np.zeros(shape=(num_steps,), dtype=ctype._dtype)
+        increments += np.random.uniform(0, 10, size=num_steps)
+        increments = test_context.nparray_to_context_array(increments)
+        test_context.kernels.run_atomic_test(
+            el=atomic, increments=increments, retvals=retvals, length=num_steps
+        )
+        increments = test_context.nparray_from_context_array(increments)
+        if ctype == xo.Float32:
+            assert np.isclose(
+                atomic.val, np.sum(increments), atol=10.0, rtol=1.0e-4
+            )
+        else:
+            assert np.isclose(
+                atomic.val, np.sum(increments), atol=1.0e-6, rtol=1.0e-12
+            )

{xobjects-0.5.7 → xobjects-0.5.9}/tests/test_shared_memory.py

@@ -53,7 +53,7 @@ def test_shared_memory():
 
                   // sum s[0] += s[1]
                   if (tid == 0){
-                    sdata[tid] += sdata[tid + 1]; 
+                    sdata[tid] += sdata[tid + 1];
 
                     // write sum from shared to global mem
                     atomicAdd(&result[tid], sdata[tid]);

xobjects-0.5.9/xobjects/_version.py

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

{xobjects-0.5.7 → xobjects-0.5.9}/xobjects/capi.py

@@ -5,7 +5,7 @@
 
 from .context import Kernel, Arg
 
-from .scalar import Int64, Void, Int8, is_scalar
+from .scalar import UInt32, Int64, Void, is_scalar
 from .struct import is_field, is_struct
 from .array import is_index, is_array
 from .ref import is_unionref, is_ref
@@ -103,9 +103,10 @@ def get_layers(parts):
 
 
 def int_from_obj(offset, conf):
-    inttype = gen_pointer(conf.get("inttype", "int64_t") + "*", conf)
-    chartype = gen_pointer(conf.get("chartype", "char") + "*", conf)
-    return f"*({inttype})(({chartype}) obj+{offset})"
+    """Generate code to read the integer at location obj + offset (bytes)."""
+    int_pointer_type = gen_pointer(conf.get("inttype", "int64_t") + "*", conf)
+    char_pointer_type = gen_pointer(conf.get("chartype", "char") + "*", conf)
+    return f"*({int_pointer_type})(({char_pointer_type}) obj+{offset})"
 
 
 def Field_get_c_offset(self, conf):
@@ -146,7 +147,7 @@ def Index_get_c_offset(part, conf, icount):
         out.append(f"  offset+={soffset};")
     else:
         lookup_field_offset = f"offset+{soffset}"
-        out.append(f"  offset={int_from_obj(lookup_field_offset, conf)};")
+        out.append(f"  offset+={int_from_obj(lookup_field_offset, conf)};")
     return out
 
 
@@ -206,16 +207,17 @@ def gen_fun_kernel(cls, path, action, const, extra, ret, add_nindex=False):
 def gen_c_pointed(target: Arg, conf):
     size = gen_c_size_from_arg(target, conf)
     ret = gen_c_type_from_arg(target, conf)
+
     if target.pointer or is_compound(target.atype) or is_string(target.atype):
         chartype = gen_pointer(conf.get("chartype", "char") + "*", conf)
         return f"({ret})(({chartype}) obj+offset)"
-    else:
-        rettype = gen_pointer(ret + "*", conf)
-        if size == 1:
-            return f"*(({rettype}) obj+offset)"
-        else:
-            chartype = gen_pointer(conf.get("chartype", "char") + "*", conf)
-            return f"*({rettype})(({chartype}) obj+offset)"
+
+    rettype = gen_pointer(ret + "*", conf)
+    if size == 1:
+        return f"*(({rettype}) obj+offset)"
+
+    chartype = gen_pointer(conf.get("chartype", "char") + "*", conf)
+    return f"*({rettype})(({chartype}) obj+offset)"
 
 
 def gen_method_get(cls, path, conf):
@@ -364,13 +366,74 @@ def gen_method_size(cls, path, conf):
 
 
 def gen_method_shape(cls, path, conf):
-    "return shape in an array"
-    return None, None
+    array_type = path[-1]
+
+    out_shape_arg = Arg(Int64, name="out_shape", pointer=True)
+
+    kernel = gen_fun_kernel(
+        cls,
+        path,
+        const=False,
+        action="shape",
+        extra=[out_shape_arg],
+        ret=None,
+        add_nindex=True,
+    )
+    decl = gen_c_decl_from_kernel(kernel, conf)
+
+    lst = [decl + "{"]
+
+    nd = len(array_type._shape)
+
+    if array_type._is_static_shape:
+        terms = [str(dim) for dim in array_type._shape]
+    else:
+        lst.append(gen_method_offset(path, conf))
+        arrarg = Arg(Int64, pointer=True)
+        pointed = gen_c_pointed(arrarg, conf)
+        typearr = gen_pointer("int64_t*", conf)
+        lst.append(f"  {typearr} arr = {pointed};")
+        dim_len_idx = 1
+        terms = []
+        for dim_len in array_type._shape:
+            if dim_len:
+                terms.append(str(dim_len))
+            else:
+                terms.append(f"arr[{dim_len_idx}]")
+                dim_len_idx += 1
+
+    for dim_idx in range(nd):
+        lst.append(f"  out_shape[{dim_idx}] = {terms[dim_idx]};")
+
+    lst.append("}")
+
+    return "\n".join(lst), kernel
 
 
 def gen_method_nd(cls, path, conf):
     "return length of shape"
-    return None, None
+    array_type = path[-1]
+
+    retarg = Arg(UInt32)
+
+    kernel = gen_fun_kernel(
+        cls,
+        path,
+        const=False,
+        action="nd",
+        extra=[],
+        ret=retarg,
+        add_nindex=True,
+    )
+    decl = gen_c_decl_from_kernel(kernel, conf)
+
+    lst = [decl + "{"]
+
+    nd = len(array_type._shape)
+    lst.append(f"  return {nd};")
+    lst.append("}")
+
+    return "\n".join(lst), kernel
 
 
 def gen_method_strides(cls, path, conf):
@@ -507,8 +570,8 @@ def methods_from_path(cls, path, conf):
 
     if is_array(lasttype):
         out.append(gen_method_len(cls, path, conf))
-    #    out.append(gen_method_shape(cls, path, conf))
-    #    out.append(gen_method_nd(cls, path, conf))
+        out.append(gen_method_shape(cls, path, conf))
+        out.append(gen_method_nd(cls, path, conf))
     #    out.append(gen_method_strides(cls, path, conf))
     #    out.append(gen_method_getpos(cls, path, conf))
 

{xobjects-0.5.7 → xobjects-0.5.9}/xobjects/context.py

@@ -375,12 +375,30 @@ class XContext(ABC):
         return sources
 
     @abstractmethod
-    def nparray_to_context_array(self, arr):
-        return arr
+    def nparray_to_context_array(self, arr, copy=False):
+        """Obtain an array on the context, given a numpy array.
+
+        Args:
+            arr: numpy array
+            copy: if True, always create a copy in the context. If False,
+                try to avoid copy if possible (not guaranteed).
+
+        Returns:
+            array on the context
+        """
 
     @abstractmethod
-    def nparray_from_context_array(self, dev_arr):
-        return dev_arr
+    def nparray_from_context_array(self, dev_arr, copy=False):
+        """Obtain a numpy array, given an array on the context.
+
+        Args:
+            arr: array on the context
+            copy: if True, always create a copy in the context. If False,
+                try to avoid copy if possible (not guaranteed).
+
+        Returns:
+            Numpy array
+        """
 
     @property
     @abstractmethod
@@ -650,14 +668,15 @@ def get_context_from_string(ctxstr):
 
     if ctxstr is None:
         return xo.ContextCpu()
+
+    ll = ctxstr.split(":")
+    if len(ll) <= 1:
+        ctxtype = ll[0]
+        option = []
     else:
-        ll = ctxstr.split(":")
-        if len(ll) <= 1:
-            ctxtype = ll[0]
-            option = []
-        else:
-            ctxtype, options = ctxstr.split(":")
-            option = options.split(",")
+        ctxtype, options = ctxstr.split(":")
+        option = options.split(",")
+
     if ctxtype == "ContextCpu":
         if len(option) == 0:
             return xo.ContextCpu()

{xobjects-0.5.7 → xobjects-0.5.9}/xobjects/context_cpu.py

@@ -543,7 +543,7 @@ class ContextCpu(XContext):
 
         return None
 
-    def nparray_to_context_array(self, arr):
+    def nparray_to_context_array(self, arr, copy=False):
         """
         Moves a numpy array to the device memory. No action is performed by
         this function in the CPU context. The method is provided
@@ -551,14 +551,16 @@ class ContextCpu(XContext):
 
         Args:
             arr (numpy.ndarray): Array to be transferred
+            copy (bool): If True, a copy of the array is made.
 
         Returns:
-            numpy.ndarray: The same array (no copy!).
-
+            numpy.ndarray: Numpy array with the same data, original or a copy.
         """
+        if copy:
+            arr = np.copy(arr)
         return arr
 
-    def nparray_from_context_array(self, dev_arr):
+    def nparray_from_context_array(self, dev_arr, copy=False):
         """
         Moves an array to the device to a numpy array. No action is performed by
         this function in the CPU context. The method is provided so that the CPU
@@ -566,10 +568,13 @@ class ContextCpu(XContext):
 
         Args:
             dev_arr (numpy.ndarray): Array to be transferred
-        Returns:
-            numpy.ndarray: The same array (no copy!)
+            copy (bool): If True, a copy of the array is made.
 
+        Returns:
+            numpy.ndarray: Numpy array with the same data, original or a copy.
         """
+        if copy:
+            dev_arr = np.copy(dev_arr)
         return dev_arr
 
     @property
@@ -579,7 +584,6 @@ class ContextCpu(XContext):
         through ``nplike_lib`` to keep compatibility with the other contexts.
 
         """
-
         return np
 
     @property
@@ -589,7 +593,6 @@ class ContextCpu(XContext):
         through ``splike_lib`` to keep compatibility with the other contexts.
 
         """
-
         return sp
 
     def synchronize(self):

{xobjects-0.5.7 → xobjects-0.5.9}/xobjects/context_cupy.py

@@ -487,29 +487,32 @@ class ContextCupy(XContext):
     def __str__(self):
         return f"{type(self).__name__}:{cupy.cuda.get_device_id()}"
 
-    def nparray_to_context_array(self, arr):
+    def nparray_to_context_array(self, arr, copy=False):
         """
         Copies a numpy array to the device memory.
 
         Args:
             arr (numpy.ndarray): Array to be transferred
+            copy (bool): This parameter is ignored for CUDA, as the data lives
+                on a different device.
 
         Returns:
             cupy.ndarray:The same array copied to the device.
-
         """
         dev_arr = cupy.array(arr)
         return dev_arr
 
-    def nparray_from_context_array(self, dev_arr):
+    def nparray_from_context_array(self, dev_arr, copy=False):
         """
         Copies an array to the device to a numpy array.
 
         Args:
             dev_arr (cupy.ndarray): Array to be transferred.
+            copy (bool): This parameter is ignored for CUDA, as the data lives
+                on a different device.
+
         Returns:
             numpy.ndarray: The same data copied to a numpy array.
-
         """
         return dev_arr.get()
 

{xobjects-0.5.7 → xobjects-0.5.9}/xobjects/context_pyopencl.py

@@ -81,6 +81,8 @@ if _enabled:
 
 
 class ContextPyopencl(XContext):
+    context_cache = {}
+
     @property
     def nplike_array_type(self):
         return cla.Array
@@ -128,20 +130,29 @@ class ContextPyopencl(XContext):
         super().__init__()
 
         # TODO assume one device only
-        if device is None:
+        if device in self.context_cache:
+            self.platform, self.device, self.context = self.context_cache[
+                device
+            ]
+        elif device is None:
             self.context = cl.create_some_context(interactive=False)
             self.device = self.context.devices[0]
             self.platform = self.device.platform
         else:
             if isinstance(device, str):
-                platform, device = map(int, device.split("."))
+                platform, device_ = map(int, device.split("."))
                 self.platform = cl.get_platforms()[platform]
-                self.device = self.platform.get_devices()[device]
+                self.device = self.platform.get_devices()[device_]
             else:
                 self.device = device
                 self.platform = device.platform
 
             self.context = cl.Context([self.device])
+            self.context_cache[device] = (
+                self.platform,
+                self.device,
+                self.context,
+            )
 
         self.queue = cl.CommandQueue(self.context)
 
@@ -253,28 +264,30 @@ class ContextPyopencl(XContext):
         device_id = self.platform.get_devices().index(self.device)
         return f"{type(self).__name__}:{platform_id}.{device_id}"
 
-    def nparray_to_context_array(self, arr):
-        """
-        Copies a numpy array to the device memory.
+    def nparray_to_context_array(self, arr, copy=False):
+        """Copies a numpy array to the device memory.
+
         Args:
             arr (numpy.ndarray): Array to be transferred
+            copy (bool): This parameter is ignored for OpenCL, as the data lives
+                on a different device.
 
         Returns:
             pyopencl.array.Array:The same array copied to the device.
-
         """
         dev_arr = cla.to_device(self.queue, arr)
         return dev_arr
 
-    def nparray_from_context_array(self, dev_arr):
-        """
-        Copies an array to the device to a numpy array.
+    def nparray_from_context_array(self, dev_arr, copy=False):
+        """Copies an array to the device to a numpy array.
 
         Args:
             dev_arr (pyopencl.array.Array): Array to be transferred.
+            copy (bool): This parameter is ignored for OpenCL, as the data lives
+                on a different device.
+
         Returns:
             numpy.ndarray: The same data copied to a numpy array.
-
         """
         return dev_arr.get()
 

{xobjects-0.5.7 → xobjects-0.5.9}/xobjects/general.py

@@ -18,9 +18,8 @@ _print = Print()
 
 
 def assert_allclose(a, b, rtol=0, atol=0, max_outliers=0):
-
     try:
-        if a == b: # if this passes we return without raising
+        if a == b:  # if this passes we return without raising
             return
     except:
         pass

xobjects-0.5.9/xobjects.egg-info/PKG-INFO

@@ -0,0 +1,39 @@
+Metadata-Version: 2.4
+Name: xobjects
+Version: 0.5.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
+Author: Riccardo De Maria
+Author-email: Riccardo De Maria <riccardo.de.maria@cern.ch>
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://xsuite.readthedocs.io/
+Project-URL: Bug Tracker, https://github.com/xsuite/xsuite/issues
+Project-URL: Documentation, https://xsuite.readthedocs.io/
+Project-URL: Source Code, https://github.com/xsuite/xobjects
+Project-URL: Download, https://pypi.python.org/pypi/xobjects
+Requires-Python: >=3.7
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: cffi
+Requires-Dist: scipy
+Provides-Extra: tests
+Requires-Dist: pytest; extra == "tests"
+Requires-Dist: pytest-mock; extra == "tests"
+Dynamic: author
+Dynamic: download-url
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: requires-python
+
+# XObjects
+
+Library to create and manipulate serialized object in CPU and GPU memory efficiently.
+
+The library defines:
+-  Contexts: create buffers and manage computations e.g. BytearrayContext, CLContext, ...
+-  Buffers: reference to internal memory buffer with an allocator e.g.
+   BytearrayContext, CLBuffer, ...
+-  Data types: define users define types to create and manipulate objects on
+   buffers and generate C api, e.g. Float64, Int64, String, Struct, Array, Union

{xobjects-0.5.7 → xobjects-0.5.9}/xobjects.egg-info/SOURCES.txt

@@ -1,4 +1,5 @@
 LICENSE
+README.md
 pyproject.toml
 setup.cfg
 setup.py
@@ -7,6 +8,7 @@ tests/test_array.py
 tests/test_buffer.py
 tests/test_capi.py
 tests/test_chunk.py
+tests/test_common.py
 tests/test_context_opencl.py
 tests/test_hybrid_class.py
 tests/test_kernel.py
@@ -44,5 +46,6 @@ xobjects/union.py
 xobjects.egg-info/PKG-INFO
 xobjects.egg-info/SOURCES.txt
 xobjects.egg-info/dependency_links.txt
+xobjects.egg-info/entry_points.txt
 xobjects.egg-info/requires.txt
 xobjects.egg-info/top_level.txt

xobjects-0.5.9/xobjects.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[xobjects]
+include = xobjects

xobjects-0.5.7/PKG-INFO

@@ -1,22 +0,0 @@
-Metadata-Version: 2.1
-Name: xobjects
-Version: 0.5.7
-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
-Author: Riccardo De Maria
-Author-email: riccardo.de.maria@cern.ch
-License: Apache 2.0
-Project-URL: Bug Tracker, https://github.com/xsuite/xsuite/issues
-Project-URL: Documentation, https://xsuite.readthedocs.io/
-Project-URL: Source Code, https://github.com/xsuite/xobjects
-Requires-Python: >=3.7
-License-File: LICENSE
-Requires-Dist: numpy
-Requires-Dist: cffi
-Requires-Dist: scipy
-Provides-Extra: tests
-Requires-Dist: pytest; extra == "tests"
-Requires-Dist: pytest-mock; extra == "tests"
-
-In-memory serialization and code generator for CPU and GPU

xobjects-0.5.7/pyproject.toml

@@ -1,4 +0,0 @@
-[tool.black]
-line-length = 79
-target-version = ['py36', 'py37', 'py38']
-include = '\.pyi?$'

xobjects-0.5.7/xobjects/_version.py

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

xobjects-0.5.7/xobjects.egg-info/PKG-INFO

@@ -1,22 +0,0 @@
-Metadata-Version: 2.1
-Name: xobjects
-Version: 0.5.7
-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
-Author: Riccardo De Maria
-Author-email: riccardo.de.maria@cern.ch
-License: Apache 2.0
-Project-URL: Bug Tracker, https://github.com/xsuite/xsuite/issues
-Project-URL: Documentation, https://xsuite.readthedocs.io/
-Project-URL: Source Code, https://github.com/xsuite/xobjects
-Requires-Python: >=3.7
-License-File: LICENSE
-Requires-Dist: numpy
-Requires-Dist: cffi
-Requires-Dist: scipy
-Provides-Extra: tests
-Requires-Dist: pytest; extra == "tests"
-Requires-Dist: pytest-mock; extra == "tests"
-
-In-memory serialization and code generator for CPU and GPU