|
| 1 | +""" |
| 2 | +This file is part of pyAMReX |
| 3 | +
|
| 4 | +Copyright 2026 AMReX community |
| 5 | +Authors: Axel Huebl |
| 6 | +License: BSD-3-Clause-LBNL |
| 7 | +""" |
| 8 | + |
| 9 | +from .extensions.Array4 import register_Array4_extension |
| 10 | +from .extensions.ArrayOfStructs import register_AoS_extension |
| 11 | +from .extensions.MultiFab import register_MultiFab_extension |
| 12 | +from .extensions.ParticleContainer import ( |
| 13 | + list_particle_species, |
| 14 | + read_particles, |
| 15 | + register_ParticleContainer_extension, |
| 16 | +) |
| 17 | +from .extensions.PODVector import register_PODVector_extension |
| 18 | +from .extensions.SmallMatrix import register_SmallMatrix_extension |
| 19 | +from .extensions.StructOfArrays import register_SoA_extension |
| 20 | + |
| 21 | + |
| 22 | +def setup_module(ns, amr): |
| 23 | + """Populate an ``amrex.space{1,2,3}d`` namespace. |
| 24 | +
|
| 25 | + Those three packages are identical apart from the compiled pybind module |
| 26 | + they wrap and their ``d_decl()``, so everything else is defined once here |
| 27 | + and installed into their namespace. |
| 28 | +
|
| 29 | + Class-level additions could equally be done from the ``register_*`` |
| 30 | + functions, because a class object is shared. Module-level names cannot: |
| 31 | + ``from .amrex_?d_pybind import *`` has already run by then, so a name added |
| 32 | + to the pybind module afterwards would not appear in the package namespace. |
| 33 | + They have to be written into ``ns`` instead, which is what this does. |
| 34 | +
|
| 35 | + Injected callables get their ``__module__`` set to the target module, so |
| 36 | + that Sphinx ``autofunction`` and the CI stub generator attribute them to |
| 37 | + ``amrex.space3d`` rather than to this helper. |
| 38 | +
|
| 39 | + Parameters |
| 40 | + ---------- |
| 41 | + ns : dict |
| 42 | + The calling module's ``globals()``. |
| 43 | + amr : module |
| 44 | + That module's compiled bindings, e.g. ``amrex_3d_pybind``. |
| 45 | + """ |
| 46 | + name = ns["__name__"] |
| 47 | + |
| 48 | + ns["__version__"] = amr.__version__ |
| 49 | + ns["__doc__"] = amr.__doc__ |
| 50 | + ns["__license__"] = amr.__license__ |
| 51 | + ns["__author__"] = amr.__author__ |
| 52 | + |
| 53 | + # enhance the C++ classes with methods written in pure Python |
| 54 | + register_Array4_extension(amr) |
| 55 | + register_MultiFab_extension(amr) |
| 56 | + register_PODVector_extension(amr) |
| 57 | + register_SmallMatrix_extension(amr) |
| 58 | + register_SoA_extension(amr) |
| 59 | + register_AoS_extension(amr) |
| 60 | + register_ParticleContainer_extension(amr) |
| 61 | + |
| 62 | + def Print(*args, **kwargs): |
| 63 | + """Wrap amrex::Print() - only the IO processor writes""" |
| 64 | + if not amr.initialized(): |
| 65 | + print("warning: Print all - AMReX not initialized") |
| 66 | + print(*args, **kwargs) |
| 67 | + elif amr.ParallelDescriptor.IOProcessor(): |
| 68 | + print(*args, **kwargs) |
| 69 | + |
| 70 | + def read_particles_( |
| 71 | + plotfile, particle_dir="particles", communicate=True, container=None |
| 72 | + ): |
| 73 | + """Read AMReX particle data from a plotfile/checkpoint into a container. |
| 74 | +
|
| 75 | + See :py:func:`amrex.extensions.ParticleContainer.read_particles` for details. |
| 76 | + """ |
| 77 | + return read_particles(amr, plotfile, particle_dir, communicate, container) |
| 78 | + |
| 79 | + read_particles_.__name__ = "read_particles" |
| 80 | + read_particles_.__qualname__ = "read_particles" |
| 81 | + |
| 82 | + def module_getattr(attr): |
| 83 | + """Resolve ``xp`` lazily (PEP 562). |
| 84 | +
|
| 85 | + ``amr.xp`` is the array namespace matching this build: NumPy on CPU, |
| 86 | + CuPy for CUDA/HIP, dpnp for SYCL. It is the module counterpart of the |
| 87 | + ``to_xp`` methods, for code that needs to call into the array library |
| 88 | + itself, e.g. ``amr.xp.sin(...)``. |
| 89 | +
|
| 90 | + Like every other CuPy/dpnp use in pyAMReX, those are optional |
| 91 | + dependencies: they are imported here on first access, never at import |
| 92 | + time, so ``import amrex`` works on a GPU build without them. Only |
| 93 | + touching ``amr.xp`` (or a ``to_cupy``/``to_dpnp``/``to_xp`` call) |
| 94 | + requires one to be installed. |
| 95 | +
|
| 96 | + Raises |
| 97 | + ------ |
| 98 | + ImportError |
| 99 | + On a GPU build whose array library (CuPy or dpnp) is not installed. |
| 100 | + """ |
| 101 | + if attr == "xp": |
| 102 | + import importlib |
| 103 | + |
| 104 | + from .extensions.dlpack_helpers import xp_module_name |
| 105 | + |
| 106 | + module_name = xp_module_name(amr) |
| 107 | + try: |
| 108 | + xp = importlib.import_module(module_name) |
| 109 | + except ImportError as e: |
| 110 | + raise ImportError( |
| 111 | + f"amrex.xp needs {module_name!r}, which is an optional " |
| 112 | + f"dependency of pyAMReX and is not installed. Install it, " |
| 113 | + f"or use the to_numpy()/to_cupy()/to_dpnp() methods " |
| 114 | + f"directly." |
| 115 | + ) from e |
| 116 | + ns["xp"] = xp # subsequent lookups skip __getattr__ |
| 117 | + return xp |
| 118 | + raise AttributeError(f"module {name!r} has no attribute {attr!r}") |
| 119 | + |
| 120 | + module_getattr.__name__ = "__getattr__" |
| 121 | + module_getattr.__qualname__ = "__getattr__" |
| 122 | + |
| 123 | + ns["Print"] = Print |
| 124 | + ns["read_particles"] = read_particles_ |
| 125 | + ns["list_particle_species"] = list_particle_species |
| 126 | + ns["__getattr__"] = module_getattr |
| 127 | + |
| 128 | + for injected in (Print, read_particles_, module_getattr): |
| 129 | + injected.__module__ = name |
0 commit comments