PodVector_*.from_* NumPy/CuPy

Helpers to copy data directly from `NumPy`/`CuPy` to simplify
user code, e.g., to add particles in AMReX codes.

Author: ax3l

src/amrex/extensions/PODVector.py

@@ -98,6 +98,130 @@ def podvector_to_xp(self, copy=False):
     return self.to_cupy(copy) if amr.Config.have_gpu else self.to_numpy(copy)
 
 
+def _is_host_accessible(cls):
+    """Check if a PODVector type's allocator provides host-accessible memory.
+
+    On CPU builds all allocators are host-accessible.  On GPU builds only
+    ``std``, ``pinned`` and ``managed`` allocators live in host memory;
+    the rest (``arena``, ``device``, ``async``, ...) are device-only.
+    """
+    import inspect
+
+    amr = inspect.getmodule(cls)
+    if not amr.Config.have_gpu:
+        return True
+    suffix = cls.__name__.rsplit("_", 1)[-1]
+    return suffix in ("std", "pinned", "managed")
+
+
+def podvector_from_numpy(cls, arr):
+    """
+    Create a new PODVector from a NumPy array (or array-like).
+
+    Always copies the data into a newly allocated PODVector.
+    Only works for allocator types with host-accessible memory (e.g.,
+    ``std``, ``pinned``).  For device-only allocators, use
+    :meth:`from_cupy` or :meth:`from_xp` instead.
+
+    Parameters
+    ----------
+    cls : type
+        The PODVector type to construct.
+    arr : array_like
+        Input data, convertible to a NumPy array.
+
+    Returns
+    -------
+    PODVector
+        A new PODVector with a copy of the data.
+
+    Raises
+    ------
+    TypeError
+        If the allocator is not host-accessible.
+    """
+    import numpy as np
+
+    n = len(arr)
+    if n == 0:
+        return cls()
+
+    if not _is_host_accessible(cls):
+        raise TypeError(
+            f"{cls.__name__} is not host-accessible. "
+            "Use from_cupy() or from_xp() instead."
+        )
+
+    pv = cls(n)
+    np.array(pv, copy=False)[:] = arr
+    return pv
+
+
+def podvector_from_cupy(cls, arr):
+    """
+    Create a new PODVector from a CuPy array (or array-like).
+
+    Always copies the data into a newly allocated PODVector.
+    Works for every allocator type: for host-only allocators the
+    data is staged to the host through NumPy automatically.
+
+    Parameters
+    ----------
+    cls : type
+        The PODVector type to construct.
+    arr : array_like
+        Input data, convertible to a CuPy array.
+
+    Returns
+    -------
+    PODVector
+        A new PODVector with a copy of the data.
+    """
+    import cupy as cp
+
+    n = len(arr)
+    if n == 0:
+        return cls()
+    pv = cls(n)
+    if _is_host_accessible(cls):
+        import numpy as np
+
+        np.array(pv, copy=False)[:] = cp.asnumpy(arr)
+    else:
+        cp.asarray(pv)[:] = arr
+    return pv
+
+
+def podvector_from_xp(cls, arr):
+    """
+    Create a new PODVector from a NumPy or CuPy array,
+    depending on amr.Config.have_gpu .
+
+    Always copies the data into a newly allocated PODVector.
+    Unlike :meth:`to_xp`, a zero-copy view is not possible here because
+    PODVector always owns its memory through its allocator.
+
+    This function is similar to CuPy's xp naming suggestion for CPU/GPU agnostic code:
+    https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code
+
+    Parameters
+    ----------
+    cls : type
+        The PODVector type to construct.
+    arr : array_like
+        Input data (NumPy or CuPy array).
+
+    Returns
+    -------
+    PODVector
+        A new PODVector with a copy of the data.
+    """
+    import inspect
+
+    amr = inspect.getmodule(cls)
+    return cls.from_cupy(arr) if amr.Config.have_gpu else cls.from_numpy(arr)
+
+
 def register_PODVector_extension(amr):
     """PODVector helper methods"""
     import inspect
@@ -112,6 +236,12 @@ def register_PODVector_extension(amr):
             and member.__name__.startswith("PODVector_")
         ),
     ):
+        # instance methods: PODVector -> array
         POD_type.to_numpy = podvector_to_numpy
         POD_type.to_cupy = podvector_to_cupy
         POD_type.to_xp = podvector_to_xp
+
+        # class methods: array -> PODVector
+        POD_type.from_numpy = classmethod(podvector_from_numpy)
+        POD_type.from_cupy = classmethod(podvector_from_cupy)
+        POD_type.from_xp = classmethod(podvector_from_xp)

tests/test_podvector.py

@@ -41,3 +41,49 @@ def test_array_interface():
 
     podv[1] = 5
     assert arr[1] == podv[1] == 5
+
+
+def test_from_numpy():
+    import numpy as np
+
+    # basic roundtrip
+    arr = np.array([1.0, 2.5, 3.7, 4.0], dtype=np.float64)
+    podv = amr.PODVector_real_std.from_numpy(arr)
+    assert podv.size() == 4
+    result = podv.to_numpy()
+    np.testing.assert_array_equal(result, arr)
+
+    # from_numpy creates a copy, not a view
+    arr[0] = 999.0
+    assert podv[0] != 999.0
+
+    # empty array
+    empty = np.array([], dtype=np.float64)
+    podv_empty = amr.PODVector_real_std.from_numpy(empty)
+    assert podv_empty.size() == 0
+
+    # from list (array-like)
+    podv_list = amr.PODVector_real_std.from_numpy([10.0, 20.0])
+    assert podv_list.size() == 2
+    assert podv_list[1] == 20.0
+
+
+def test_from_xp():
+    import numpy as np
+
+    arr = np.array([1.0, 2.0, 3.0])
+    podv = amr.PODVector_real_std.from_xp(arr)
+    assert podv.size() == 3
+    result = podv.to_numpy()
+    np.testing.assert_array_equal(result, arr)
+
+
+def test_from_xp_default():
+    """Test from_xp on the default (platform-adaptive) PODVector type alias."""
+    import numpy as np
+
+    arr = np.array([5.0, 6.0, 7.0])
+    podv = amr.PODVector_real_default.from_xp(arr)
+    assert podv.size() == 3
+    result = podv.to_numpy(copy=True)
+    np.testing.assert_array_equal(result, arr)