Source code for gpe.gpu.bec2

"""GPU versions of the code.

Requires cupy and the various NVIDIA dependencies.
"""

from . import cupy as cp

from pytimeode import mixins_gpu

from .bases import PeriodicBasisGPU

from .. import bec2



[docs]
u = bec2.u






[docs]
class State(mixins_gpu.GPUArrayStateMixin, bec2.State):


[docs]
    xp = cp


    if cp:


[docs]
        asnumpy = staticmethod(cp.asnumpy)



    def __init__(
        self,
        basis=None,
        # Specify either basis or the following
        Nxyz=(2**5, 2**5, 2**5),
        Lxyz=(30 * u.micron, 50 * u.micron, 50 * u.micron),
        symmetric_grid=False,
        twist=None,
        **kw,
    ):
        args = dict(Nxyz=Nxyz, Lxyz=Lxyz, symmetric_lattice=symmetric_grid)
        if basis is None:
            basis = PeriodicBasisGPU(**args)
        else:
            kw.update(args)
        super().__init__(basis=basis, **kw)