Skip to content

Commit 46a67bb

Browse files
bleykaufbleykauf
authored
Merge pull request #84 from bleykauf/feature/intensity-profile
Refactor IntensityProfile initialization to use keyword arguments only
2 parents 5f68127 + 5bcc9f2 commit 46a67bb

File tree

7 files changed

+69
-28
lines changed

7 files changed

+69
-28
lines changed

‎aisim/beam.py‎

Lines changed: 39 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
"""Classes and functions related to the interferometry lasers."""
22

3+
from functools import partial
4+
35
import matplotlib.pyplot as plt
46
import numpy as np
57
from matplotlib.colors import Colormap
@@ -64,38 +66,62 @@ class IntensityProfile:
6466
6567
Attributes
6668
----------
67-
r_profile : float
68-
1/e² radius of the Gaussian intensity profile in m
69-
center_rabi_freq : float
70-
Rabi frequency at center of intensity profile in rad/s
7169
r_beam : float or None
7270
Beam radius in m. If set, Rabi frequency will be set to 0 outside of the beam.
71+
profile_func : function
72+
Function that calculates the Rabi frequency at a position of a Gaussian beam.
7373
"""
7474

75-
def __init__(self, r_profile, center_rabi_freq, r_beam=None):
76-
self.r_profile = r_profile
77-
self.center_rabi_freq = center_rabi_freq
75+
def __init__(self, *, r_profile, center_rabi_freq, r_beam=None):
7876
self.r_beam = r_beam
77+
self.profile_func = partial(
78+
self._gaussian_profile,
79+
r_profile=r_profile,
80+
center_rabi_freq=center_rabi_freq,
81+
)
7982

8083
def get_rabi_freq(self, pos):
8184
"""
8285
Rabi frequency at a position of a Gaussian beam.
8386
8487
Parameters
8588
----------
86-
pos : (n × 2) array or (n × 3) array
89+
pos : (n x 2) array or (n x 3) array
8790
positions of the n atoms in two or three dimensions (x, y, [z]).
8891
8992
Returns
9093
-------
9194
rabi_freqs : array of float
92-
the Rabi frequencies for the n positions. If ``r_beam`` is set, the Rabi
95+
the Rabi frequencies for the n positions. If `r_beam` is set, the Rabi
9396
frequency will be set to 0 outside of the beam.
9497
"""
95-
r_sq = pos[:, 0] ** 2 + pos[:, 1] ** 2
96-
rabi_freqs = self.center_rabi_freq * np.exp(-2 * r_sq / (self.r_profile**2))
98+
rabi_freqs = self.profile_func(pos)
9799
if self.r_beam is not None:
98-
rabi_freqs[r_sq > self.r_beam] = 0.0
100+
rabi_freqs[pos[:, 0] ** 2 + pos[:, 1] ** 2 > self.r_beam] = 0.0
101+
return rabi_freqs
102+
103+
@staticmethod
104+
def _gaussian_profile(pos, r_profile: float, center_rabi_freq: float):
105+
"""
106+
Rabi frequency at a position of a Gaussian beam.
107+
108+
Parameters
109+
----------
110+
pos : (n x 2) array or (n x 3) array
111+
positions of the n atoms in two or three dimensions (x, y, [z]).
112+
r_profile : float
113+
1/e² radius of the Gaussian intensity profile in m
114+
center_rabi_freq : float
115+
Rabi frequency at center of intensity profile in rad/s
116+
117+
Returns
118+
-------
119+
rabi_freqs : array of float
120+
the Rabi frequencies for the n positions. If `r_beam` is set, the Rabi
121+
frequency will be set to 0 outside of the beam.
122+
"""
123+
r_sq = pos[:, 0] ** 2 + pos[:, 1] ** 2
124+
rabi_freqs = center_rabi_freq * np.exp(-2 * r_sq / (r_profile**2))
99125
return rabi_freqs
100126

101127

@@ -157,7 +183,7 @@ def get_value(self, pos: np.ndarray) -> np.ndarray:
157183
158184
Parameters
159185
----------
160-
pos : n × 3 array
186+
pos : n x 3 array
161187
array of position vectors (x, y, z) where the wavefront is probed
162188
163189
Returns

‎docs/examples/basic-usage.ipynb‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,13 @@
4848
")"
4949
]
5050
},
51+
{
52+
"cell_type": "code",
53+
"execution_count": null,
54+
"metadata": {},
55+
"outputs": [],
56+
"source": []
57+
},
5158
{
5259
"cell_type": "markdown",
5360
"metadata": {},

‎docs/examples/multiport-ai.ipynb‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,9 @@
4747
"\n",
4848
"center_rabi_freq = 2 * np.pi * 12.5e3\n",
4949
"r_profile = 29.5e-3 / 2 # 1/e^2 beam radius in m\n",
50-
"intensity_profile = ais.IntensityProfile(r_profile, center_rabi_freq)\n",
50+
"intensity_profile = ais.IntensityProfile(\n",
51+
" r_profile=r_profile, center_rabi_freq=center_rabi_freq\n",
52+
")\n",
5153
"wave_vectors = ais.Wavevectors(k1=8.052945514e6, k2=-8.052802263e6)\n",
5254
"\n",
5355
"# Creating an atomic ensemble, initially in the ground state |g, v_0, x_0 + v_0*t>\n",
@@ -161,7 +163,7 @@
161163
{
162164
"data": {
163165
"text/plain": [
164-
"[0.4617680981617182, 0.0, 0.0, 0.0, 0.0, 0.5382319018382818]"
166+
"[0.4541289929081983, 0.0, 0.0, 0.0, 0.0, 0.5458710070918017]"
165167
]
166168
},
167169
"execution_count": 6,
@@ -198,12 +200,12 @@
198200
{
199201
"data": {
200202
"text/plain": [
201-
"[0.1188964472677993,\n",
203+
"[0.11464363907969147,\n",
202204
" 0.0,\n",
203-
" 0.3746771445733455,\n",
204-
" 0.34287165089391897,\n",
205+
" 0.37881323142669276,\n",
206+
" 0.3394853538285068,\n",
205207
" 0.0,\n",
206-
" 0.16355475726493632]"
208+
" 0.16705777566510893]"
207209
]
208210
},
209211
"execution_count": 8,
@@ -247,7 +249,7 @@
247249
{
248250
"data": {
249251
"text/plain": [
250-
"[0.011844398649651894, 0.7057043968176124]"
252+
"[0.014352736559798305, 0.7039458486954012]"
251253
]
252254
},
253255
"execution_count": 10,

‎docs/examples/rabi-oscillations.ipynb‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,9 @@
176176
"source": [
177177
"center_rabi_freq = 2 * np.pi * 12.5e3 # center Rabi frequency in Hz\n",
178178
"r_profile = 29.5e-3 / 2 # 1/e^2 beam radius in m\n",
179-
"intensity_profile = ais.IntensityProfile(r_profile, center_rabi_freq)"
179+
"intensity_profile = ais.IntensityProfile(\n",
180+
" r_profile=r_profile, center_rabi_freq=center_rabi_freq\n",
181+
")"
180182
]
181183
},
182184
{
@@ -301,7 +303,7 @@
301303
{
302304
"data": {
303305
"text/plain": [
304-
"<matplotlib.legend.Legend at 0x10fcb4ad0>"
306+
"<matplotlib.legend.Legend at 0x10e778d70>"
305307
]
306308
},
307309
"execution_count": 11,

‎docs/examples/wavefront-aberrations.ipynb‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -274,7 +274,7 @@
274274
{
275275
"data": {
276276
"text/plain": [
277-
"<matplotlib.legend.Legend at 0x112ba74d0>"
277+
"<matplotlib.legend.Legend at 0x1238a7770>"
278278
]
279279
},
280280
"execution_count": 11,

‎tests/beam_test.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ def test_intensity_profile():
5454
r_profile = 1
5555
pos1 = np.array([[0, 0, 0]])
5656
pos2 = np.array([[0, 1, 0]])
57-
intensity_profile = ais.IntensityProfile(r_profile, 1)
57+
intensity_profile = ais.IntensityProfile(r_profile=r_profile, center_rabi_freq=1)
5858

5959
rabi1 = intensity_profile.get_rabi_freq(pos1)
6060
rabi2 = intensity_profile.get_rabi_freq(pos2)
@@ -63,7 +63,9 @@ def test_intensity_profile():
6363
np.testing.assert_almost_equal(ratio[0], 1 / np.e**2)
6464

6565
# test that Rabi frequency is zero outside of the beam
66-
intensity_profile2 = ais.IntensityProfile(r_profile, 1, r_beam=0.5)
66+
intensity_profile2 = ais.IntensityProfile(
67+
r_profile=r_profile, center_rabi_freq=1, r_beam=0.5
68+
)
6769
rabi3 = intensity_profile2.get_rabi_freq(pos1)
6870
assert rabi3[0] == 1
6971
rabi4 = intensity_profile2.get_rabi_freq(pos2)

‎tests/prop_test.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ def test_free_propagator():
4747
def test_two_level_transition_propagator_unitarity():
4848
atoms = create_random_thermal_atoms(100)
4949

50-
intensity_profile = ais.IntensityProfile(1, 1)
50+
intensity_profile = ais.IntensityProfile(r_profile=1, center_rabi_freq=1)
5151
wf = ais.gen_wavefront(10, seed=1)
5252
wave_vectors = ais.Wavevectors()
5353

@@ -81,7 +81,7 @@ def prop_unitarity_tester(n_pulses, pi_half_time):
8181
init_state.append(0)
8282
atoms = create_random_thermal_atoms(100, state_kets=init_state)
8383

84-
intensity_profile = ais.IntensityProfile(1, 1)
84+
intensity_profile = ais.IntensityProfile(r_profile=1, center_rabi_freq=1)
8585
wf = ais.gen_wavefront(10, seed=1)
8686
wave_vectors = ais.Wavevectors()
8787

@@ -114,7 +114,9 @@ def prop_time_reversal_tester(n_pulses, pi_half_time):
114114
r_beam = 29.5e-3 / 2 # 1/e^2 beam radius in m
115115
wave_vectors = ais.Wavevectors(k1=8.052945514e6, k2=-8.052802263e6)
116116
center_rabi_freq = 2 * np.pi / 4 / pi_half_time
117-
intensity_profile = ais.IntensityProfile(r_beam, center_rabi_freq)
117+
intensity_profile = ais.IntensityProfile(
118+
r_profile=r_beam, center_rabi_freq=center_rabi_freq
119+
)
118120

119121
for n_pulse in range(0, n_pulses):
120122
propagator = ais.SpatialSuperpositionTransitionPropagator(

0 commit comments

Comments
 (0)