Skip to content

Kernels

CPU

CPU kernels in the NLSE package are responsible for solving the nonlinear Schrödinger equation using CPU resources.

We use the popular Numba library to just-in-time compile array functions.

These functions are not meant to be called directly and instead work as the backend specific implementations for the main classes methods.

Due to the manual indexing involved in the implementations, these kernels do not support broadcasting (for now).

CPU kernels are suitable for small to medium-sized problems or when GPU resources are not available. As it uses multithreading internally, they will benefit from higher core counts.

nl_prop(A, A_sq, dz, alpha, V, g, Isat)

A compiled parallel implementation to apply real space terms

Parameters:

Name Type Description Default
A ndarray

The field to propagate

 required
A_sq ndarray

The field modulus squared

 required
dz float

Propagation step in m

 required
alpha float

Losses

 required
V ndarray

Potential

 required
g float

Interactions

 required
Isat float

Saturation

 required
Source code in NLSE/kernels_cpu.py 
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
@numba.njit(parallel=True, fastmath=True, cache=True, boundscheck=False)
def nl_prop(
    A: np.ndarray,
    A_sq: np.ndarray,
    dz: float,
    alpha: float,
    V: np.ndarray,
    g: float,
    Isat: float,
) -> None:
    """A compiled parallel implementation to apply real space terms

    Args:
        A (np.ndarray): The field to propagate
        A_sq (np.ndarray): The field modulus squared
        dz (float): Propagation step in m
        alpha (float): Losses
        V (np.ndarray): Potential
        g (float): Interactions
        Isat (float): Saturation
    """
    A = A.ravel()
    A_sq = A_sq.ravel()
    V = V.ravel()
    for i in numba.prange(A.size):
        # saturation
        sat = 1 / (1 + A_sq[i] / Isat)
        # Losses and interactions
        arg = -alpha + 1j * g * A_sq[i] * sat + 1j * V[i]
        A[i] *= np.exp(dz * arg)

nl_prop_c(A1, A_sq_1, A_sq_2, dz, alpha, V, g11, g12, Isat1, Isat2)

A fused kernel to apply real space terms Args: A1 (cp.ndarray): The field to propagate (1st component) A_sq_1 (cp.ndarray): The field modulus squared (1st component) A_sq_2 (cp.ndarray): The field modulus squared (2nd component) dz (float): Propagation step in m alpha (float): Losses V (cp.ndarray): Potential g11 (float): Intra-component interactions g12 (float): Inter-component interactions Isat1 (float): Saturation parameter of first component Isat2 (float): Saturation parameter of second component

Source code in NLSE/kernels_cpu.py 
 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
@numba.njit(parallel=True, fastmath=True, cache=True, boundscheck=False)
def nl_prop_c(
    A1: np.ndarray,
    A_sq_1: np.ndarray,
    A_sq_2: np.ndarray,
    dz: float,
    alpha: float,
    V: np.ndarray,
    g11: float,
    g12: float,
    Isat1: float,
    Isat2: float,
) -> None:
    """A fused kernel to apply real space terms
    Args:
        A1 (cp.ndarray): The field to propagate (1st component)
        A_sq_1 (cp.ndarray): The field modulus squared (1st component)
        A_sq_2 (cp.ndarray): The field modulus squared (2nd component)
        dz (float): Propagation step in m
        alpha (float): Losses
        V (cp.ndarray): Potential
        g11 (float): Intra-component interactions
        g12 (float): Inter-component interactions
        Isat1 (float): Saturation parameter of first component
        Isat2 (float): Saturation parameter of second component
    """
    A1 = A1.ravel()
    A_sq_1 = A_sq_1.ravel()
    A_sq_2 = A_sq_2.ravel()
    for i in numba.prange(A1.size):
        # Saturation parameter
        sat = 1 / (1 + A_sq_1[i] * 1 / Isat1 + A_sq_2[i] * 1 / Isat2)
        # Losses
        arg = -alpha * sat
        # Interactions
        arg += 1j * (g11 * A_sq_1[i] * sat + g12 * A_sq_2[i] * sat)
        # Potential
        arg += 1j * V[i]
        A1[i] *= np.exp(dz * arg)

nl_prop_without_V(A, A_sq, dz, alpha, g, Isat)

A fused kernel to apply real space terms

Parameters:

Name Type Description Default
A ndarray

The field to propagate

 required
A_sq ndarray

The field modulus squared

 required
dz float

Propagation step in m

 required
alpha float

Losses

 required
g float

Interactions

 required
Isat float

Saturation

 required
Source code in NLSE/kernels_cpu.py 
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
@numba.njit(parallel=True, fastmath=True, cache=True, boundscheck=False)
def nl_prop_without_V(
    A: np.ndarray,
    A_sq: np.ndarray,
    dz: float,
    alpha: float,
    g: float,
    Isat: float,
) -> None:
    """A fused kernel to apply real space terms

    Args:
        A (cp.ndarray): The field to propagate
        A_sq (cp.ndarray): The field modulus squared
        dz (float): Propagation step in m
        alpha (float): Losses
        g (float): Interactions
        Isat (float): Saturation
    """
    A = A.ravel()
    A_sq = A_sq.ravel()
    for i in numba.prange(A.size):
        # saturation
        sat = 1 / (1 + A_sq[i] / Isat)
        # Losses and interactions
        arg = -alpha + 1j * g * A_sq[i] * sat
        A[i] *= np.exp(dz * arg)

nl_prop_without_V_c(A1, A_sq_1, A_sq_2, dz, alpha, g11, g12, Isat1, Isat2)

A fused kernel to apply real space terms Args: A1 (cp.ndarray): The field to propagate (1st component) A_sq_1 (cp.ndarray): The field modulus squared (1st component) A_sq_2 (cp.ndarray): The field modulus squared (2nd component) dz (float): Propagation step in m alpha (float): Losses g11 (float): Intra-component interactions g12 (float): Inter-component interactions Isat1 (float): Saturation parameter of first component Isat2 (float): Saturation parameter of second component

Source code in NLSE/kernels_cpu.py 
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
@numba.njit(parallel=True, fastmath=True, cache=True, boundscheck=False)
def nl_prop_without_V_c(
    A1: np.ndarray,
    A_sq_1: np.ndarray,
    A_sq_2: np.ndarray,
    dz: float,
    alpha: float,
    g11: float,
    g12: float,
    Isat1: float,
    Isat2: float,
) -> None:
    """A fused kernel to apply real space terms
    Args:
        A1 (cp.ndarray): The field to propagate (1st component)
        A_sq_1 (cp.ndarray): The field modulus squared (1st component)
        A_sq_2 (cp.ndarray): The field modulus squared (2nd component)
        dz (float): Propagation step in m
        alpha (float): Losses
        g11 (float): Intra-component interactions
        g12 (float): Inter-component interactions
        Isat1 (float): Saturation parameter of first component
        Isat2 (float): Saturation parameter of second component
    """
    A1 = A1.ravel()
    A_sq_1 = A_sq_1.ravel()
    A_sq_2 = A_sq_2.ravel()
    for i in numba.prange(A1.size):
        # Saturation parameter
        sat = 1 / (1 + A_sq_1[i] * 1 / Isat1 + A_sq_2[i] * 1 / Isat2)
        # Losses
        arg = -alpha * sat
        # Interactions
        arg += 1j * (g11 * A_sq_1[i] * sat + g12 * A_sq_2[i] * sat)
        A1[i] *= np.exp(dz * arg)

rabi_coupling(A1, A2, dz, omega)

Apply a Rabi coupling term. This function implements the Rabi hopping term. It exchanges density between the two components.

Parameters:

Name Type Description Default
A1 ndarray

First field / component

 required
A2 ndarray

Second field / component

 required
dz float

Solver step

 required
omega float

Rabi coupling strength

 required
Source code in NLSE/kernels_cpu.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
@numba.njit(parallel=True, fastmath=True, cache=True, boundscheck=False)
def rabi_coupling(A1: np.ndarray, A2: np.ndarray, dz: float, omega: float) -> None:
    """Apply a Rabi coupling term.
    This function implements the Rabi hopping term.
    It exchanges density between the two components.

    Args:
        A1 (np.ndarray): First field / component
        A2 (np.ndarray): Second field / component
        dz (float): Solver step
        omega (float): Rabi coupling strength
    """
    A1 = A1.ravel()
    A2 = A2.ravel()
    A1_old = A1.copy()
    for i in numba.prange(A1.size):
        A1[i] = np.cos(omega * dz) * A1[i] - 1j * np.sin(omega * dz) * A2[i]
        A2[i] = np.cos(omega * dz) * A2[i] - 1j * np.sin(omega * dz) * A1_old[i]

square_mod(A, A_sq)

Compute the square modulus of the field

Parameters:

Name Type Description Default
A ndarray

The field

 required
A_sq ndarray

The modulus squared of the field

 required

Returns:

Type Description
None

None
Source code in NLSE/kernels_cpu.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
@numba.njit(parallel=True, fastmath=True, cache=True, boundscheck=False)
def square_mod(A: np.ndarray, A_sq: np.ndarray) -> None:
    """Compute the square modulus of the field

    Args:
        A (np.ndarray): The field
        A_sq (np.ndarray): The modulus squared of the field

    Returns:
        None
    """
    A = A.ravel()
    A_sq = A_sq.ravel()
    for i in numba.prange(A.size):
        A_sq[i] = A[i].real * A[i].real + A[i].imag * A[i].imag

vortex(im, i, j, ii, jj, ll)

Generates a vortex of charge l at a position (i,j) on the image im.

Parameters:

Name Type Description Default
im ndarray

Image

 required
i int

position row of the vortex

 required
j int

position column of the vortex

 required
ii int

meshgrid position row (coordinates of the image)

 required
jj int

meshgrid position column (coordinates of the image)

 required
ll int

vortex charge

 required

Returns:

Type Description
None

None
Source code in NLSE/kernels_cpu.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
@numba.njit(parallel=True, fastmath=True, cache=True, boundscheck=False)
def vortex(
    im: np.ndarray, i: int, j: int, ii: np.ndarray, jj: np.ndarray, ll: int
) -> None:
    """Generates a vortex of charge l at a position (i,j) on the image im.

    Args:
        im (np.ndarray): Image
        i (int): position row of the vortex
        j (int): position column of the vortex
        ii (int): meshgrid position row (coordinates of the image)
        jj (int): meshgrid position column (coordinates of the image)
        ll (int): vortex charge

    Returns:
        None
    """
    for i in numba.prange(im.shape[0]):
        for j in numba.prange(im.shape[1]):
            im[i, j] += np.angle(((ii[i, j] - i) + 1j * (jj[i, j] - j)) ** ll)

GPU (CUDA)

GPU kernels in the NLSE package are responsible for solving the nonlinear Schrödinger equation using GPU resources. They utilize the computational power of the graphics processing unit (GPU) to perform the necessary calculations.

These kernels use the cupy.fuse API to just-in-time compile the array operations to a single kernel.

The strategy here is to maximize GPU occupancy by grouping operations to a complex enough task such that we are not limited by memory bandwidth. As with the CPU kernels, the paradigm is to try and mutate arrays in place as much as possible to avoid costly memory transfers.

To this end, most of these kernels have the following signature:

@cp.fuse
def kernel(A: cp.ndarray, *args):
    # do something to A
    A += args[0]
    A *= args[1]

nl_prop(A, A_sq, dz, alpha, V, g, Isat)

A fused kernel to apply real space terms

Parameters:

Name Type Description Default
A ndarray

The field to propagate

 required
A_sq ndarray

The field modulus squared

 required
dz float

Propagation step in m

 required
alpha float

Losses

 required
V ndarray

Potential

 required
g float

Interactions

 required
Isat float

Saturation

 required
Source code in NLSE/kernels_gpu.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
@cp.fuse(kernel_name="nl_prop")
def nl_prop(
    A: cp.ndarray,
    A_sq: cp.ndarray,
    dz: float,
    alpha: float,
    V: cp.ndarray,
    g: float,
    Isat: float,
) -> None:
    """A fused kernel to apply real space terms

    Args:
        A (cp.ndarray): The field to propagate
        A_sq (cp.ndarray): The field modulus squared
        dz (float): Propagation step in m
        alpha (float): Losses
        V (cp.ndarray): Potential
        g (float): Interactions
        Isat (float): Saturation
    """
    # saturation
    sat = 1 / (1 + A_sq / Isat)
    # Interactions
    arg = 1j * g * A_sq * sat
    # Losses
    arg += -alpha * sat
    # Potential
    arg += 1j * V
    arg *= dz
    cp.exp(arg, out=arg)
    A *= arg

nl_prop_c(A1, A_sq_1, A_sq_2, dz, alpha, V, g11, g12, Isat1, Isat2)

A fused kernel to apply real space terms Args: A1 (cp.ndarray): The field to propagate (1st component) A_sq_1 (cp.ndarray): The field modulus squared (1st component) A_sq_2 (cp.ndarray): The field modulus squared (2nd component) dz (float): Propagation step in m alpha (float): Losses V (cp.ndarray): Potential g11 (float): Intra-component interactions g12 (float): Inter-component interactions Isat1 (float): Saturation parameter of first component Isat2 (float): Saturation parameter of second component

Source code in NLSE/kernels_gpu.py 
 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
@cp.fuse(kernel_name="nl_prop_c")
def nl_prop_c(
    A1: cp.ndarray,
    A_sq_1: cp.ndarray,
    A_sq_2: cp.ndarray,
    dz: float,
    alpha: float,
    V: cp.ndarray,
    g11: float,
    g12: float,
    Isat1: float,
    Isat2: float,
) -> None:
    """A fused kernel to apply real space terms
    Args:
        A1 (cp.ndarray): The field to propagate (1st component)
        A_sq_1 (cp.ndarray): The field modulus squared (1st component)
        A_sq_2 (cp.ndarray): The field modulus squared (2nd component)
        dz (float): Propagation step in m
        alpha (float): Losses
        V (cp.ndarray): Potential
        g11 (float): Intra-component interactions
        g12 (float): Inter-component interactions
        Isat1 (float): Saturation parameter of first component
        Isat2 (float): Saturation parameter of second component
    """
    # Saturation parameter
    sat = 1 / (1 + A_sq_1 * 1 / Isat1 + A_sq_2 * 1 / Isat2)
    # Interactions
    arg = 1j * (g11 * A_sq_1 * sat + g12 * A_sq_2 * sat)
    # Losses
    arg += -alpha * sat
    # Potential
    arg += 1j * V
    arg *= dz
    cp.exp(arg, out=arg)
    A1 *= arg

nl_prop_without_V(A, A_sq, dz, alpha, g, Isat)

A fused kernel to apply real space terms

Parameters:

Name Type Description Default
A ndarray

The field to propagate

 required
A_sq ndarray

The field modulus squared

 required
dz float

Propagation step in m

 required
alpha float

Losses

 required
g float

Interactions

 required
Isat float

Saturation

 required
Source code in NLSE/kernels_gpu.py 
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
@cp.fuse(kernel_name="nl_prop_without_V")
def nl_prop_without_V(
    A: cp.ndarray,
    A_sq: cp.ndarray,
    dz: float,
    alpha: float,
    g: float,
    Isat: float,
) -> None:
    """A fused kernel to apply real space terms

    Args:
        A (cp.ndarray): The field to propagate
        A_sq (cp.ndarray): The field modulus squared
        dz (float): Propagation step in m
        alpha (float): Losses
        g (float): Interactions
        Isat (float): Saturation
    """
    # saturation
    sat = 1 / (1 + A_sq / Isat)
    # Interactions
    arg = 1j * g * A_sq * sat
    # Losses
    arg += -alpha * sat
    arg *= dz
    cp.exp(arg, out=arg)
    A *= arg

nl_prop_without_V_c(A1, A_sq_1, A_sq_2, dz, alpha, g11, g12, Isat1, Isat2)

A fused kernel to apply real space terms Args: A1 (cp.ndarray): The field to propagate (1st component) A_sq_1 (cp.ndarray): The field modulus squared (1st component) A_sq_2 (cp.ndarray): The field modulus squared (2nd component) dz (float): Propagation step in m alpha (float): Losses g11 (float): Intra-component interactions g12 (float): Inter-component interactions Isat1 (float): Saturation parameter of first component Isat2 (float): Saturation parameter of second component

Source code in NLSE/kernels_gpu.py 
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
@cp.fuse(kernel_name="nl_prop_without_V_c")
def nl_prop_without_V_c(
    A1: cp.ndarray,
    A_sq_1: cp.ndarray,
    A_sq_2: cp.ndarray,
    dz: float,
    alpha: float,
    g11: float,
    g12: float,
    Isat1: float,
    Isat2: float,
) -> None:
    """A fused kernel to apply real space terms
    Args:
        A1 (cp.ndarray): The field to propagate (1st component)
        A_sq_1 (cp.ndarray): The field modulus squared (1st component)
        A_sq_2 (cp.ndarray): The field modulus squared (2nd component)
        dz (float): Propagation step in m
        alpha (float): Losses
        g11 (float): Intra-component interactions
        g12 (float): Inter-component interactions
        Isat1 (float): Saturation parameter of first component
        Isat2 (float): Saturation parameter of second component
    """
    # Saturation parameter
    sat = 1 / (1 + A_sq_1 * 1 / Isat1 + A_sq_2 * 1 / Isat2)
    # Interactions
    arg = 1j * (g11 * A_sq_1 * sat + g12 * A_sq_2 * sat)
    # Losses
    arg += -alpha * sat
    arg *= dz
    cp.exp(arg, out=arg)
    A1 *= arg

rabi_coupling(A1, A2, dz, omega)

Apply a Rabi coupling term. This function implements the Rabi hopping term. It exchanges density between the two components.

Parameters:

Name Type Description Default
A1 ndarray

First field / component

 required
A2 ndarray

Second field / component

 required
dz float

Solver step

 required
omega float

Rabi coupling strength

 required
Source code in NLSE/kernels_gpu.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
@cp.fuse(kernel_name="rabi_coupling")
def rabi_coupling(A1: cp.array, A2: cp.array, dz: float, omega: float) -> None:
    """Apply a Rabi coupling term.
    This function implements the Rabi hopping term.
    It exchanges density between the two components.

    Args:
        A1 (cp.ndarray): First field / component
        A2 (cp.ndarray): Second field / component
        dz (float): Solver step
        omega (float): Rabi coupling strength
    """
    A1_old = A1.copy()
    A1[:] = cp.cos(omega * dz) * A1 - 1j * cp.sin(omega * dz) * A2
    A2[:] = cp.cos(omega * dz) * A2 - 1j * cp.sin(omega * dz) * A1_old

square_mod(A, A_sq)

Compute the square modulus of the field

Parameters:

Name Type Description Default
A ndarray

The field

 required
A_sq ndarray

The modulus squared of the field

 required

Returns:

Type Description
None

None
Source code in NLSE/kernels_gpu.py 
179
180
181
182
183
184
185
186
187
188
189
190
@cp.fuse(kernel_name="square_mod_cp")
def square_mod(A: cp.ndarray, A_sq: cp.ndarray) -> None:
    """Compute the square modulus of the field

    Args:
        A (cp.ndarray): The field
        A_sq (cp.ndarray): The modulus squared of the field

    Returns:
        None
    """
    A_sq[:] = cp.abs(A) ** 2

vortex_cp(im, i, j, ii, jj, ll)

Generates a vortex of charge l at a position (i,j) on the image im.

Parameters:

Name Type Description Default
im ndarray

Image

 required
i int

position row of the vortex

 required
j int

position column of the vortex

 required
ii int

meshgrid position row (coordinates of the image)

 required
jj int

meshgrid position column (coordinates of the image)

 required
ll int

vortex charge

 required

Returns:

Type Description
None

None
Source code in NLSE/kernels_gpu.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
@cp.fuse(kernel_name="vortex_cp")
def vortex_cp(
    im: cp.ndarray, i: int, j: int, ii: cp.ndarray, jj: cp.ndarray, ll: int
) -> None:
    """Generates a vortex of charge l at a position (i,j) on the image im.

    Args:
        im (np.ndarray): Image
        i (int): position row of the vortex
        j (int): position column of the vortex
        ii (int): meshgrid position row (coordinates of the image)
        jj (int): meshgrid position column (coordinates of the image)
        ll (int): vortex charge

    Returns:
        None
    """
    im += cp.angle(((ii - i) + 1j * (jj - j)) ** ll)

OpenCL (GPU or CPU)

In order to take advantage of the modern hardware, there is experimental support for OpenCL for the NLSE class.

The long-term goal is to use a unified interface for GPU and CPU (which is what OpenCL already does).

This uses PyOpenCL and its array interface to follow the same approach as the other two implementations.

Due to the limited support for OpenCL on newer hardware (mostly Macs), this might not be the winning strategy so use at your own risk ! 😇

nl_prop(A, A_sq, dz, alpha, V, g, Isat)

A fused kernel to apply real space terms

Parameters:

Name Type Description Default
A Array

The field to propagate

 required
A_sq Array

The field modulus squared

 required
dz float

Propagation step in m

 required
alpha float

Losses

 required
V Array

Potential

 required
g float

Interactions

 required
Isat float

Saturation

 required
Source code in NLSE/kernels_cl.py 
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def nl_prop(
    A: cla.Array,
    A_sq: cla.Array,
    dz: float,
    alpha: float,
    V: cla.Array,
    g: float,
    Isat: float,
) -> None:
    """A fused kernel to apply real space terms

    Args:
        A (cla.Array): The field to propagate
        A_sq (cla.Array): The field modulus squared
        dz (float): Propagation step in m
        alpha (float): Losses
        V (cla.Array): Potential
        g (float): Interactions
        Isat (float): Saturation
    """
    # saturation
    sat = 1 / (1 + A_sq / Isat)
    # Interactions
    arg = 1j * g * A_sq * sat
    # Losses
    arg += -alpha * sat
    # Potential
    arg += 1j * V
    arg *= dz
    arg = clmath.exp(arg)
    A *= arg

nl_prop_c(A1, A_sq_1, A_sq_2, dz, alpha, V, g11, g12, Isat)

A fused kernel to apply real space terms Args: A1 (cla.Array): The field to propagate (1st component) A_sq_1 (cla.Array): The field modulus squared (1st component) A_sq_2 (cla.Array): The field modulus squared (2nd component) dz (float): Propagation step in m alpha (float): Losses V (cla.Array): Potential g11 (float): Intra-component interactions g12 (float): Inter-component interactions Isat1 (float): Saturation parameter of first component Isat2 (float): Saturation parameter of second component

Source code in NLSE/kernels_cl.py 
 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
def nl_prop_c(
    A1: cla.Array,
    A_sq_1: cla.Array,
    A_sq_2: cla.Array,
    dz: float,
    alpha: float,
    V: cla.Array,
    g11: float,
    g12: float,
    Isat: float,
) -> None:
    """A fused kernel to apply real space terms
    Args:
        A1 (cla.Array): The field to propagate (1st component)
        A_sq_1 (cla.Array): The field modulus squared (1st component)
        A_sq_2 (cla.Array): The field modulus squared (2nd component)
        dz (float): Propagation step in m
        alpha (float): Losses
        V (cla.Array): Potential
        g11 (float): Intra-component interactions
        g12 (float): Inter-component interactions
        Isat1 (float): Saturation parameter of first component
        Isat2 (float): Saturation parameter of second component
    """
    # Saturation parameter
    sat = 1 / (1 + (A_sq_1 + A_sq_2) / Isat)
    # Interactions
    arg = 1j * (g11 * A_sq_1 * sat + g12 * A_sq_2 * sat)
    # Losses
    arg += -alpha * sat
    # Potential
    arg += 1j * V
    arg *= dz
    arg = clmath.exp(arg)
    A1 *= arg

nl_prop_without_V(A, A_sq, dz, alpha, g, Isat)

A fused kernel to apply real space terms

Parameters:

Name Type Description Default
A Array

The field to propagate

 required
A_sq Array

The field modulus squared

 required
dz float

Propagation step in m

 required
alpha float

Losses

 required
g float

Interactions

 required
Isat float

Saturation

 required
Source code in NLSE/kernels_cl.py 
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
def nl_prop_without_V(
    A: cla.Array,
    A_sq: cla.Array,
    dz: float,
    alpha: float,
    g: float,
    Isat: float,
) -> None:
    """A fused kernel to apply real space terms

    Args:
        A (cla.Array): The field to propagate
        A_sq (cla.Array): The field modulus squared
        dz (float): Propagation step in m
        alpha (float): Losses
        g (float): Interactions
        Isat (float): Saturation
    """
    # saturation
    sat = 1 / (1 + A_sq / Isat)
    # Interactions
    arg = 1j * g * A_sq * sat
    # Losses
    arg += -alpha * sat
    arg *= dz
    arg = clmath.exp(arg)
    A *= arg

nl_prop_without_V_c(A1, A_sq_1, A_sq_2, dz, alpha, g11, g12, Isat)

A fused kernel to apply real space terms Args: A1 (cla.Array): The field to propagate (1st component) A_sq_1 (cla.Array): The field modulus squared (1st component) A_sq_2 (cla.Array): The field modulus squared (2nd component) dz (float): Propagation step in m alpha (float): Losses g11 (float): Intra-component interactions g12 (float): Inter-component interactions Isat1 (float): Saturation parameter of first component Isat2 (float): Saturation parameter of second component

Source code in NLSE/kernels_cl.py 
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
def nl_prop_without_V_c(
    A1: cla.Array,
    A_sq_1: cla.Array,
    A_sq_2: cla.Array,
    dz: float,
    alpha: float,
    g11: float,
    g12: float,
    Isat: float,
) -> None:
    """A fused kernel to apply real space terms
    Args:
        A1 (cla.Array): The field to propagate (1st component)
        A_sq_1 (cla.Array): The field modulus squared (1st component)
        A_sq_2 (cla.Array): The field modulus squared (2nd component)
        dz (float): Propagation step in m
        alpha (float): Losses
        g11 (float): Intra-component interactions
        g12 (float): Inter-component interactions
        Isat1 (float): Saturation parameter of first component
        Isat2 (float): Saturation parameter of second component
    """
    # Saturation parameter
    sat = 1 / (1 + A_sq_1 / Isat)
    # Interactions
    arg = 1j * (g11 * A_sq_1 * sat + g12 * A_sq_2 * sat)
    # Losses
    arg += -alpha * sat
    arg *= dz
    arg = clmath.exp(arg)
    A1 *= arg

rabi_coupling(A, dz, omega)

Apply a Rabi coupling term. This function implements the Rabi hopping term. It exchanges density between the two components.

Parameters:

Name Type Description Default
A Array

First field / component

 required
dz float

Solver step

 required
omega float

Rabi coupling strength

 required
Source code in NLSE/kernels_cl.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
def rabi_coupling(A, dz: float, omega: float) -> None:
    """Apply a Rabi coupling term.
    This function implements the Rabi hopping term.
    It exchanges density between the two components.

    Args:
        A (cla.Array): First field / component
        dz (float): Solver step
        omega (float): Rabi coupling strength
    """
    A1 = A[..., 0, :, :]
    A2 = A[..., 1, :, :]
    A1_old = A1.copy()
    A1[:] = clmath.cos(omega * dz) * A1 - 1j * clmath.sin(omega * dz) * A2
    A2[:] = clmath.cos(omega * dz) * A2 - 1j * clmath.sin(omega * dz) * A1_old

square_mod(A, A_sq)

Compute the square modulus of the field

Parameters:

Name Type Description Default
A Array

The field

 required
A_sq Array

The modulus squared of the field

 required

Returns:

Type Description
None

None
Source code in NLSE/kernels_cl.py 
173
174
175
176
177
178
179
180
181
182
183
def square_mod(A: cla.Array, A_sq: cla.Array) -> None:
    """Compute the square modulus of the field

    Args:
        A (cla.Array): The field
        A_sq (cla.Array): The modulus squared of the field

    Returns:
        None
    """
    A_sq[:] = A.real * A.real + A.imag * A.imag

vortex_cp(im, i, j, ii, jj, ll)

Generates a vortex of charge l at a position (i,j) on the image im.

Parameters:

Name Type Description Default
im ndarray

Image

 required
i int

position row of the vortex

 required
j int

position column of the vortex

 required
ii int

meshgrid position row (coordinates of the image)

 required
jj int

meshgrid position column (coordinates of the image)

 required
ll int

vortex charge

 required

Returns:

Type Description
None

None
Source code in NLSE/kernels_cl.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
def vortex_cp(
    im: cla.Array, i: int, j: int, ii: cla.Array, jj: cla.Array, ll: int
) -> None:
    """Generates a vortex of charge l at a position (i,j) on the image im.

    Args:
        im (np.ndarray): Image
        i (int): position row of the vortex
        j (int): position column of the vortex
        ii (int): meshgrid position row (coordinates of the image)
        jj (int): meshgrid position column (coordinates of the image)
        ll (int): vortex charge

    Returns:
        None
    """
    im += clmath.atan(((ii - i) + 1j * (jj - j)) ** ll)