Skip to content

GPE

In this page you will find the reference documentation for the CNLSE class.

Bases: NLSE

A class to solve GPE.

Source code in NLSE/gpe.py 
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
class GPE(NLSE):
    """A class to solve GPE."""

    def __init__(
        self,
        gamma: float,
        N: float,
        window: float,
        g: float,
        V: Union[np.ndarray, None],
        m: float = 87 * atomic_mass,
        NX: int = 1024,
        NY: int = 1024,
        sat: float = np.inf,
        nl_length: float = 0,
        backend: str = __BACKEND__,
    ) -> object:
        """Instantiate the simulation.

        Solves an equation : d/dt psi = -1/2m(d2/dx2 + d2/dy2) psi + V psi +
          g psi**2 psi

        Args:
            gamma (float): Losses in Hz
            N (float): Total number of atoms
            window (float): Window size in m
            g (float): Interaction energy in Hz*m^2
            V (np.ndarray): Potential in Hz
            m (float, optionnal): mass of one atom in kg.
                Defaults to 87*atomic_mass for Rubidium 87.
            NX (int, optional): Number of points in x.
                Defaults to 1024.
            NY (int, optional): Number of points in y.
                Defaults to 1024.
            sat (float): Saturation parameter in Hz/m^2.
            nl_length (float, optional): Non local length scale in m.
                Defaults to 0.
            backend (str, optional): "GPU" or "CPU". Defaults to __BACKEND__.
        """
        super().__init__(
            alpha=gamma,
            power=N,
            window=window,
            n2=g,
            V=V,
            L=0,
            NX=NX,
            NY=NY,
            Isat=sat,
            nl_length=nl_length,
            wvl=2 * np.pi / m,
            backend=backend,
        )
        # listof physical parameters
        self.g = g
        self.V = V
        self.gamma = self.alpha
        self.N = self.power
        self.m = self.k
        self.sat = sat
        self.delta_t = self.delta_z
        # do some conversion for the units
        self.I_sat *= epsilon_0 * c / 2

    def _build_propagator(self) -> np.ndarray:
        """Build the linear propagation matrix.

        Returns:
            propagator (np.ndarray): the propagator matrix
        """
        propagator = np.exp(
            -1j
            * 0.5
            * hbar
            * (self.Kxx**2 + self.Kyy**2)
            / self.m
            * self.delta_t
        ).astype(np.complex64)
        return propagator

    def _prepare_output_array(
        self, E_in: np.ndarray, normalize: bool
    ) -> np.ndarray:
        """Prepare the output array depending on __BACKEND__.

        Args:
            E_in (np.ndarray): Input array
            normalize (bool): Normalize the field to the total power.
        Returns:
            np.ndarray: Output array
        """
        if self.backend == "GPU" and self.__CUPY_AVAILABLE__:
            A = cp.empty_like(E_in)
            A_sq = cp.empty_like(A, dtype=A.real.dtype)
            E_in = cp.asarray(E_in)
        else:
            A = pyfftw.empty_aligned(
                E_in.shape, dtype=E_in.dtype, n=pyfftw.simd_alignment
            )
            A_sq = np.empty_like(A, dtype=A.real.dtype)
        if normalize:
            # normalization of the field
            integral = (
                (E_in.real * E_in.real + E_in.imag * E_in.imag)
                * self.delta_X
                * self.delta_Y
            ).sum(axis=self._last_axes)
            E_00 = (self.N / integral) ** 0.5
            A[:] = (E_00.T * E_in.T).T
        return A, A_sq

    def plot_field(self, A_plot: np.ndarray, z: float) -> None:
        """Plot a field for monitoring.

        Args:
            A_plot (np.ndarray): Field to plot
            z (float): Propagation distance in m.
        """
        # if array is multi-dimensional, drop dims until the shape is 2D
        if A_plot.ndim > 2:
            while len(A_plot.shape) > 2:
                A_plot = A_plot[0]
        if self.__CUPY_AVAILABLE__ and isinstance(A_plot, cp.ndarray):
            A_plot = A_plot.get()
        fig, ax = plt.subplots(1, 3, layout="constrained", figsize=(15, 5))
        fig.suptitle(rf"Field at $z$ = {z:.2e} m")
        ext_real = [
            np.min(self.X) * 1e3,
            np.max(self.X) * 1e3,
            np.min(self.Y) * 1e3,
            np.max(self.Y) * 1e3,
        ]
        ext_fourier = [
            np.min(self.Kx) * 1e-3,
            np.max(self.Kx) * 1e-3,
            np.min(self.Ky) * 1e-3,
            np.max(self.Ky) * 1e-3,
        ]
        rho = np.abs(A_plot) ** 2
        phi = np.angle(A_plot)
        im_fft = np.abs(np.fft.fftshift(np.fft.fft2(A_plot)))
        im0 = ax[0].imshow(rho, extent=ext_real)
        ax[0].set_title("Intensity")
        ax[0].set_xlabel("x (mm)")
        ax[0].set_ylabel("y (mm)")
        fig.colorbar(im0, ax=ax[0], shrink=0.6, label="Density (at/m^2)")
        im1 = ax[1].imshow(
            phi,
            extent=ext_real,
            cmap="twilight_shifted",
            vmin=-np.pi,
            vmax=np.pi,
        )
        ax[1].set_title("Phase")
        ax[1].set_xlabel("x (mm)")
        ax[1].set_ylabel("y (mm)")
        fig.colorbar(im1, ax=ax[1], shrink=0.6, label="Phase (rad)")
        im2 = ax[2].imshow(
            im_fft,
            extent=ext_fourier,
            cmap="nipy_spectral",
        )
        ax[2].set_title("Fourier space")
        ax[2].set_xlabel(r"$k_x$ ($mm^{-1}$)")
        ax[2].set_ylabel(r"$k_y$ ($mm^{-1}$)")
        fig.colorbar(im2, ax=ax[2], shrink=0.6, label="Intensity (a.u.)")
        plt.show()

__init__(gamma, N, window, g, V, m=87 * atomic_mass, NX=1024, NY=1024, sat=np.inf, nl_length=0, backend=__BACKEND__)

Instantiate the simulation.

d/dt psi = -1/2m(d2/dx2 + d2/dy2) psi + V psi +

g psi**2 psi

Parameters:

Name Type Description Default
gamma float

Losses in Hz

 required
N float

Total number of atoms

 required
window float

Window size in m

 required
g float

Interaction energy in Hz*m^2

 required
V ndarray

Potential in Hz

 required
m (float, optionnal)

mass of one atom in kg. Defaults to 87*atomic_mass for Rubidium 87.

 87 * atomic_mass
NX int

Number of points in x. Defaults to 1024.

 1024
NY int

Number of points in y. Defaults to 1024.

 1024
sat float

Saturation parameter in Hz/m^2.

 inf
nl_length float

Non local length scale in m. Defaults to 0.

 0
backend str

"GPU" or "CPU". Defaults to BACKEND.

 __BACKEND__
Source code in NLSE/gpe.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def __init__(
    self,
    gamma: float,
    N: float,
    window: float,
    g: float,
    V: Union[np.ndarray, None],
    m: float = 87 * atomic_mass,
    NX: int = 1024,
    NY: int = 1024,
    sat: float = np.inf,
    nl_length: float = 0,
    backend: str = __BACKEND__,
) -> object:
    """Instantiate the simulation.

    Solves an equation : d/dt psi = -1/2m(d2/dx2 + d2/dy2) psi + V psi +
      g psi**2 psi

    Args:
        gamma (float): Losses in Hz
        N (float): Total number of atoms
        window (float): Window size in m
        g (float): Interaction energy in Hz*m^2
        V (np.ndarray): Potential in Hz
        m (float, optionnal): mass of one atom in kg.
            Defaults to 87*atomic_mass for Rubidium 87.
        NX (int, optional): Number of points in x.
            Defaults to 1024.
        NY (int, optional): Number of points in y.
            Defaults to 1024.
        sat (float): Saturation parameter in Hz/m^2.
        nl_length (float, optional): Non local length scale in m.
            Defaults to 0.
        backend (str, optional): "GPU" or "CPU". Defaults to __BACKEND__.
    """
    super().__init__(
        alpha=gamma,
        power=N,
        window=window,
        n2=g,
        V=V,
        L=0,
        NX=NX,
        NY=NY,
        Isat=sat,
        nl_length=nl_length,
        wvl=2 * np.pi / m,
        backend=backend,
    )
    # listof physical parameters
    self.g = g
    self.V = V
    self.gamma = self.alpha
    self.N = self.power
    self.m = self.k
    self.sat = sat
    self.delta_t = self.delta_z
    # do some conversion for the units
    self.I_sat *= epsilon_0 * c / 2

plot_field(A_plot, z)

Plot a field for monitoring.

Parameters:

Name Type Description Default
A_plot ndarray

Field to plot

 required
z float

Propagation distance in m.

 required
Source code in NLSE/gpe.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
def plot_field(self, A_plot: np.ndarray, z: float) -> None:
    """Plot a field for monitoring.

    Args:
        A_plot (np.ndarray): Field to plot
        z (float): Propagation distance in m.
    """
    # if array is multi-dimensional, drop dims until the shape is 2D
    if A_plot.ndim > 2:
        while len(A_plot.shape) > 2:
            A_plot = A_plot[0]
    if self.__CUPY_AVAILABLE__ and isinstance(A_plot, cp.ndarray):
        A_plot = A_plot.get()
    fig, ax = plt.subplots(1, 3, layout="constrained", figsize=(15, 5))
    fig.suptitle(rf"Field at $z$ = {z:.2e} m")
    ext_real = [
        np.min(self.X) * 1e3,
        np.max(self.X) * 1e3,
        np.min(self.Y) * 1e3,
        np.max(self.Y) * 1e3,
    ]
    ext_fourier = [
        np.min(self.Kx) * 1e-3,
        np.max(self.Kx) * 1e-3,
        np.min(self.Ky) * 1e-3,
        np.max(self.Ky) * 1e-3,
    ]
    rho = np.abs(A_plot) ** 2
    phi = np.angle(A_plot)
    im_fft = np.abs(np.fft.fftshift(np.fft.fft2(A_plot)))
    im0 = ax[0].imshow(rho, extent=ext_real)
    ax[0].set_title("Intensity")
    ax[0].set_xlabel("x (mm)")
    ax[0].set_ylabel("y (mm)")
    fig.colorbar(im0, ax=ax[0], shrink=0.6, label="Density (at/m^2)")
    im1 = ax[1].imshow(
        phi,
        extent=ext_real,
        cmap="twilight_shifted",
        vmin=-np.pi,
        vmax=np.pi,
    )
    ax[1].set_title("Phase")
    ax[1].set_xlabel("x (mm)")
    ax[1].set_ylabel("y (mm)")
    fig.colorbar(im1, ax=ax[1], shrink=0.6, label="Phase (rad)")
    im2 = ax[2].imshow(
        im_fft,
        extent=ext_fourier,
        cmap="nipy_spectral",
    )
    ax[2].set_title("Fourier space")
    ax[2].set_xlabel(r"$k_x$ ($mm^{-1}$)")
    ax[2].set_ylabel(r"$k_y$ ($mm^{-1}$)")
    fig.colorbar(im2, ax=ax[2], shrink=0.6, label="Intensity (a.u.)")
    plt.show()