```python
from matplotlib.pylab import *
from matplotlib.patches import Rectangle
from matplotlib.collections import PatchCollection
from matplotlib.lines import Line2D
π = pi
```

```python
style.use(['dark_background', 'bmh'])
%matplotlib notebook
```

Car-trailer diagram (inverted image `res/car-trainer-k.png` available as well):

Car-trailer equation:
\begin{align}
\dot x &= s \cos \theta_0 \\
\dot y &= s \sin \theta_0 \\
\dot \theta_0 &= \frac{s}{L} \tan \phi \\
\dot \theta_1 &= \frac{s}{d_1} \sin(\theta_1 - \theta_0)
\end{align}
where $s$: signed speed, $\phi$: negative steering angle,

```python
class Truck:
    def __init__(self, display=False):

        self.W = 1  # car and trailer width, for drawing only
        self.L = 1 * self.W  # car length
        self.d = 4 * self.L  # d_1
        self.s = -0.1  # speed
        self.display = display

        self.box = [0, 40, -10, 10]
        if self.display:
            self.f = figure(figsize=(10, 5), num='The truck backer-upper', facecolor='none')
            self.ax = self.f.add_axes([0.01, 0.01, 0.98, 0.98], facecolor='black')
            self.patches = list()

            self.ax.axis('equal')
            b = self.box
            self.ax.axis([b[0] - 1, b[1], b[2], b[3]])
            self.ax.set_xticks([]); self.ax.set_yticks([])
            self.ax.axhline(); self.ax.axvline()

        self.reset()

    def reset(self, ϕ=0):
        self.ϕ = ϕ  # car initial steering angle

        # self.θ0 = deg2rad(30)  # car initial direction
        # self.θ1 = deg2rad(-30)  # trailer initial direction
        # self.x, self.y = 20, -5  # initial car coordinates

        self.θ0 = random() * 2 * π  # 0 <= ϑ₀ < 2π
        self.θ1 = (random() - 0.5) * π / 2 + self.θ0  # -π/4 <= ϑ₁ - ϑ₀ < π/4
        self.x = (random() * .75 + 0.25) * self.box[1]
        self.y = (random() - 0.5) * (self.box[3] - self.box[2])

        # If poorly initialise, then re-initialise
        if not self.valid():
            self.reset(ϕ)

        # Draw, if display is True
        if self.display: self.draw()

    def step(self, ϕ=0, dt=1):

        # Check for illegal conditions
        if self.is_jackknifed():
            print('The truck is jackknifed!')
            return

        if self.is_offscreen():
            print('The car or trailer is off screen')
            return

        self.ϕ = ϕ
        x, y, W, L, d, s, θ0, θ1, ϕ = self._get_atributes()

        # Perform state update
        self.x += s * cos(θ0) * dt
        self.y += s * sin(θ0) * dt
        self.θ0 += s / L * tan(ϕ) * dt
        self.θ1 += s / d * sin(θ0 - θ1) * dt

        return (self.x, self.y, self.θ0, *self._traler_xy(), self.θ1)

    def state(self):
        return (self.x, self.y, self.θ0, *self._traler_xy(), self.θ1)

    def _get_atributes(self):
        return (
            self.x, self.y, self.W, self.L, self.d, self.s,
            self.θ0, self.θ1, self.ϕ
        )

    def _traler_xy(self):
        x, y, W, L, d, s, θ0, θ1, ϕ = self._get_atributes()
        return x - d * cos(θ1), y - d * sin(θ1)

    def is_jackknifed(self):
        x, y, W, L, d, s, θ0, θ1, ϕ = self._get_atributes()
        return abs(θ0 - θ1) * 180 / π > 90

    def is_offscreen(self):
        x, y, W, L, d, s, θ0, θ1, ϕ = self._get_atributes()

        x1, y1 = x + 1.5 * L * cos(θ0), y + 1.5 * L * sin(θ0)
        x2, y2 = self._traler_xy()

        b = self.box
        return not (
            b[0] <= x1 <= b[1] and b[2] <= y1 <= b[3] and
            b[0] <= x2 <= b[1] and b[2] <= y2 <= b[3]
        )

    def valid(self):
        return not self.is_jackknifed() and not self.is_offscreen()

    def draw(self):
        if not self.display: return
        if self.patches: self.clear()
        self._draw_car()
        self._draw_trailer()
        self.f.canvas.draw()

    def clear(self):
        for p in self.patches:
            p.remove()
        self.patches = list()

    def _draw_car(self):
        x, y, W, L, d, s, θ0, θ1, ϕ = self._get_atributes()
        ax = self.ax

        x1, y1 = x + L / 2 * cos(θ0), y + L / 2 * sin(θ0)
        bar = Line2D((x, x1), (y, y1), lw=5, color='C2', alpha=0.8)
        ax.add_line(bar)

        car = Rectangle(
            (x1, y1 - W / 2), L, W, 0, color='C2', alpha=0.8, transform=
            matplotlib.transforms.Affine2D().rotate_deg_around(x1, y1, θ0 * 180 / π) +
            ax.transData
        )
        ax.add_patch(car)

        x2, y2 = x1 + L / 2 ** 0.5 * cos(θ0 + π / 4), y1 + L / 2 ** 0.5 * sin(θ0 + π / 4)
        left_wheel = Line2D(
            (x2 - L / 4 * cos(θ0 + ϕ), x2 + L / 4 * cos(θ0 + ϕ)),
            (y2 - L / 4 * sin(θ0 + ϕ), y2 + L / 4 * sin(θ0 + ϕ)),
            lw=3, color='C5', alpha=1)
        ax.add_line(left_wheel)

        x3, y3 = x1 + L / 2 ** 0.5 * cos(π / 4 - θ0), y1 - L / 2 ** 0.5 * sin(π / 4 - θ0)
        right_wheel = Line2D(
            (x3 - L / 4 * cos(θ0 + ϕ), x3 + L / 4 * cos(θ0 + ϕ)),
            (y3 - L / 4 * sin(θ0 + ϕ), y3 + L / 4 * sin(θ0 + ϕ)),
            lw=3, color='C5', alpha=1)
        ax.add_line(right_wheel)

        self.patches += [car, bar, left_wheel, right_wheel]

    def _draw_trailer(self):
        x, y, W, L, d, s, θ0, θ1, ϕ = self._get_atributes()
        ax = self.ax

        x, y = x - d * cos(θ1), y - d * sin(θ1) - W / 2
        trailer = Rectangle(
            (x, y), d, W, 0, color='C0', alpha=0.8, transform=
            matplotlib.transforms.Affine2D().rotate_deg_around(x, y + W/2, θ1 * 180 / π) +
            ax.transData
        )
        ax.add_patch(trailer)

        self.patches += [trailer]
```

```python
truck = Truck(display=True)
```

    <IPython.core.display.Javascript object>

```python
ϕ = deg2rad(-35)  # positive left, negative right
truck.step(ϕ)
truck.draw()
```

```python
truck.reset()
```

```python
import torch
import torch.nn as nn
from torch.optim import SGD
from tqdm import tqdm
```

```python
# Build expert data set

episodes = 20
inputs = list()
outputs = list()
truck = Truck(); episodes = 10_000  # uncooment for creating the data set

for episode in tqdm(range(episodes)):

    truck.reset()

    while truck.valid():
        initial_state = truck.state()
        ϕ = (random() - 0.5) * π / 2
        inputs.append((ϕ, *initial_state))
        outputs.append(truck.step(ϕ))
        truck.draw()
```

    100%|███████████████████████████████████████████████████████████████████████████| 10000/10000 [00:30<00:00, 331.27it/s]

```python
len(inputs), len(outputs)
```

    (598462, 598462)

```python
state_size = 6
steering_size = 1
hidden_units_e = 45

emulator = nn.Sequential(
    nn.Linear(steering_size + state_size, hidden_units_e),
    nn.ReLU(),
    nn.Linear(hidden_units_e, state_size)
)

optimiser_e = SGD(emulator.parameters(), lr=0.005)
criterion = nn.MSELoss()
```

```python
tensor_inputs = torch.Tensor(inputs)
tensor_outputs = torch.Tensor(outputs)
```

```python
mean = tensor_inputs.mean(0)
std = tensor_inputs.std(0)
tensor_inputs = (tensor_inputs - mean) / std
tensor_outputs = (tensor_outputs - mean[1:]) / std[1:]
```

```python
# Split the data into 80:20 for test:train.
test_size = int(len(tensor_inputs) * 0.8)
print(len(tensor_inputs), test_size)

train_inputs = tensor_inputs[:test_size]
train_outputs = tensor_outputs[:test_size]
test_inputs = tensor_inputs[test_size:]
test_outputs = tensor_outputs[test_size:]
```

    598462 478769

```python
len(train_inputs)
```

    478769

```python
# Emulator training
cnt = 0
for i in torch.randperm(len(train_inputs)):
    ϕ_state = train_inputs[i]
    next_state_prediction = emulator(ϕ_state)

    next_state = train_outputs[i]
    loss = criterion(next_state_prediction, next_state)

    optimiser_e.zero_grad()
    loss.backward()
    optimiser_e.step()

    if cnt == 0 or (cnt + 1) % 1000 == 0:
        print(f'{cnt + 1:4d} / {len(train_inputs)}, {loss.item():.10f}')
    cnt += 1
```

       1 / 478769, 0.9950614572
    1000 / 478769, 0.0804964378
    2000 / 478769, 0.0428197682
    3000 / 478769, 0.0141127231
    4000 / 478769, 0.0033202693
    5000 / 478769, 0.0044318531
    6000 / 478769, 0.0051610484
    7000 / 478769, 0.0228864122
    8000 / 478769, 0.0106072584
    9000 / 478769, 0.0128876530
    10000 / 478769, 0.0070054666
    11000 / 478769, 0.0059467456
    12000 / 478769, 0.0052833580
    13000 / 478769, 0.0047106510
    14000 / 478769, 0.0028358276
    15000 / 478769, 0.0077827293
    16000 / 478769, 0.0030465918
    17000 / 478769, 0.0035736961
    18000 / 478769, 0.0074268803
    19000 / 478769, 0.0066001457
    20000 / 478769, 0.0007502480
    21000 / 478769, 0.0039543477
    22000 / 478769, 0.0010553007
    23000 / 478769, 0.0006903820
    24000 / 478769, 0.0028175602
    25000 / 478769, 0.0024409688
    26000 / 478769, 0.0037122539
    27000 / 478769, 0.0019317693
    28000 / 478769, 0.0036818748
    29000 / 478769, 0.0021453791
    30000 / 478769, 0.0030969512
    31000 / 478769, 0.0002536585
    32000 / 478769, 0.0005552660
    33000 / 478769, 0.0047557242
    34000 / 478769, 0.0071815737
    35000 / 478769, 0.0029095730
    36000 / 478769, 0.0016798958
    37000 / 478769, 0.0012545474
    38000 / 478769, 0.0033800413
    39000 / 478769, 0.0019660164
    40000 / 478769, 0.0019252558
    41000 / 478769, 0.0022069914
    42000 / 478769, 0.0005663931
    43000 / 478769, 0.0012183717
    44000 / 478769, 0.0029595662
    45000 / 478769, 0.0006360257
    46000 / 478769, 0.0010558913
    47000 / 478769, 0.0009985096
    48000 / 478769, 0.0018096274
    49000 / 478769, 0.0007918300
    50000 / 478769, 0.0002374086
    51000 / 478769, 0.0005819083
    52000 / 478769, 0.0024469870
    53000 / 478769, 0.0017623119
    54000 / 478769, 0.0001458147
    55000 / 478769, 0.0009807196
    56000 / 478769, 0.0003426303
    57000 / 478769, 0.0012112396
    58000 / 478769, 0.0006564126
    59000 / 478769, 0.0006200691
    60000 / 478769, 0.0009668985
    61000 / 478769, 0.0009315140
    62000 / 478769, 0.0004245004
    63000 / 478769, 0.0005843976
    64000 / 478769, 0.0002264343
    65000 / 478769, 0.0009123204
    66000 / 478769, 0.0010148169
    67000 / 478769, 0.0002241186
    68000 / 478769, 0.0006592170
    69000 / 478769, 0.0004526471
    70000 / 478769, 0.0007023583
    71000 / 478769, 0.0023639239
    72000 / 478769, 0.0001560459
    73000 / 478769, 0.0009863733
    74000 / 478769, 0.0013332741
    75000 / 478769, 0.0012098072
    76000 / 478769, 0.0004894566
    77000 / 478769, 0.0017910395
    78000 / 478769, 0.0004137233
    79000 / 478769, 0.0002172645
    80000 / 478769, 0.0012061695
    81000 / 478769, 0.0011628958
    82000 / 478769, 0.0011113064
    83000 / 478769, 0.0006120533
    84000 / 478769, 0.0001533306
    85000 / 478769, 0.0005629210
    86000 / 478769, 0.0009137733
    87000 / 478769, 0.0016682744
    88000 / 478769, 0.0012799081
    89000 / 478769, 0.0005020430
    90000 / 478769, 0.0001229989
    91000 / 478769, 0.0005099811
    92000 / 478769, 0.0002611939
    93000 / 478769, 0.0009367279
    94000 / 478769, 0.0002065285
    95000 / 478769, 0.0002529499
    96000 / 478769, 0.0011993489
    97000 / 478769, 0.0003351772
    98000 / 478769, 0.0006693659
    99000 / 478769, 0.0003812376
    100000 / 478769, 0.0010871300
    101000 / 478769, 0.0005891433
    102000 / 478769, 0.0003245359
    103000 / 478769, 0.0007093695
    104000 / 478769, 0.0011793880
    105000 / 478769, 0.0009577349
    106000 / 478769, 0.0002015376
    107000 / 478769, 0.0003224985
    108000 / 478769, 0.0009496468
    109000 / 478769, 0.0003111093
    110000 / 478769, 0.0008396524
    111000 / 478769, 0.0004818033
    112000 / 478769, 0.0001938330
    113000 / 478769, 0.0011346542
    114000 / 478769, 0.0018436868
    115000 / 478769, 0.0007785479
    116000 / 478769, 0.0005620422
    117000 / 478769, 0.0006480485
    118000 / 478769, 0.0008514260
    119000 / 478769, 0.0004951090
    120000 / 478769, 0.0005901508
    121000 / 478769, 0.0003961015
    122000 / 478769, 0.0011176075
    123000 / 478769, 0.0008564061
    124000 / 478769, 0.0008206949
    125000 / 478769, 0.0002678942
    126000 / 478769, 0.0004150719
    127000 / 478769, 0.0001757417
    128000 / 478769, 0.0010822020
    129000 / 478769, 0.0002383867
    130000 / 478769, 0.0011134652
    131000 / 478769, 0.0007586555
    132000 / 478769, 0.0009424029
    133000 / 478769, 0.0002274624
    134000 / 478769, 0.0004314528
    135000 / 478769, 0.0000890405
    136000 / 478769, 0.0035272930
    137000 / 478769, 0.0009572839
    138000 / 478769, 0.0001938106
    139000 / 478769, 0.0004489179
    140000 / 478769, 0.0005064711
    141000 / 478769, 0.0001686048
    142000 / 478769, 0.0007692210
    143000 / 478769, 0.0017530272
    144000 / 478769, 0.0001334869
    145000 / 478769, 0.0001317821
    146000 / 478769, 0.0003962379
    147000 / 478769, 0.0004330036
    148000 / 478769, 0.0002961184
    149000 / 478769, 0.0012331858
    150000 / 478769, 0.0004032670
    151000 / 478769, 0.0050755306
    152000 / 478769, 0.0003022535
    153000 / 478769, 0.0005021636
    154000 / 478769, 0.0003630172
    155000 / 478769, 0.0005797331
    156000 / 478769, 0.0008894119
    157000 / 478769, 0.0003820368
    158000 / 478769, 0.0005773376
    159000 / 478769, 0.0003472201
    160000 / 478769, 0.0001560339
    161000 / 478769, 0.0004325507
    162000 / 478769, 0.0006542958
    163000 / 478769, 0.0003554873
    164000 / 478769, 0.0001609899
    165000 / 478769, 0.0001525247
    166000 / 478769, 0.0004682516
    167000 / 478769, 0.0008461117
    168000 / 478769, 0.0010413878
    169000 / 478769, 0.0004442078
    170000 / 478769, 0.0001186525
    171000 / 478769, 0.0002936413
    172000 / 478769, 0.0001263322
    173000 / 478769, 0.0001638131
    174000 / 478769, 0.0004988777
    175000 / 478769, 0.0004285156
    176000 / 478769, 0.0004807032
    177000 / 478769, 0.0006099577
    178000 / 478769, 0.0005521689
    179000 / 478769, 0.0006773847
    180000 / 478769, 0.0001495597
    181000 / 478769, 0.0002629830
    182000 / 478769, 0.0003887035
    183000 / 478769, 0.0008322782
    184000 / 478769, 0.0003076661
    185000 / 478769, 0.0004622064
    186000 / 478769, 0.0001474292
    187000 / 478769, 0.0002729353
    188000 / 478769, 0.0000870756
    189000 / 478769, 0.0002063987
    190000 / 478769, 0.0003386727
    191000 / 478769, 0.0000484285
    192000 / 478769, 0.0004513926
    193000 / 478769, 0.0008016932
    194000 / 478769, 0.0003567701
    195000 / 478769, 0.0000622364
    196000 / 478769, 0.0004522736
    197000 / 478769, 0.0002234644
    198000 / 478769, 0.0004175944
    199000 / 478769, 0.0003637099
    200000 / 478769, 0.0003848148
    201000 / 478769, 0.0002521462
    202000 / 478769, 0.0009894043
    203000 / 478769, 0.0000661020
    204000 / 478769, 0.0004480809
    205000 / 478769, 0.0001264342
    206000 / 478769, 0.0000716306
    207000 / 478769, 0.0003006925
    208000 / 478769, 0.0003661555
    209000 / 478769, 0.0003107784
    210000 / 478769, 0.0013076331
    211000 / 478769, 0.0001443364
    212000 / 478769, 0.0002980012
    213000 / 478769, 0.0008795884
    214000 / 478769, 0.0001604464
    215000 / 478769, 0.0002015812
    216000 / 478769, 0.0001481988
    217000 / 478769, 0.0000958917
    218000 / 478769, 0.0002497236
    219000 / 478769, 0.0003180612
    220000 / 478769, 0.0006477754
    221000 / 478769, 0.0006816951
    222000 / 478769, 0.0003031361
    223000 / 478769, 0.0003644585
    224000 / 478769, 0.0004111501
    225000 / 478769, 0.0006001596
    226000 / 478769, 0.0003331308
    227000 / 478769, 0.0004091925
    228000 / 478769, 0.0002160502
    229000 / 478769, 0.0000962787
    230000 / 478769, 0.0000556897
    231000 / 478769, 0.0003453483
    232000 / 478769, 0.0003957007
    233000 / 478769, 0.0001065692
    234000 / 478769, 0.0004295753
    235000 / 478769, 0.0001929458
    236000 / 478769, 0.0001027593
    237000 / 478769, 0.0002077750
    238000 / 478769, 0.0002964026
    239000 / 478769, 0.0002601243
    240000 / 478769, 0.0000779160
    241000 / 478769, 0.0004668287
    242000 / 478769, 0.0001751910
    243000 / 478769, 0.0002516678
    244000 / 478769, 0.0007430706
    245000 / 478769, 0.0002185611
    246000 / 478769, 0.0001535818
    247000 / 478769, 0.0001705773
    248000 / 478769, 0.0000731046
    249000 / 478769, 0.0004051411
    250000 / 478769, 0.0000388975
    251000 / 478769, 0.0001087418
    252000 / 478769, 0.0005217825
    253000 / 478769, 0.0001090987
    254000 / 478769, 0.0001108959
    255000 / 478769, 0.0001304024
    256000 / 478769, 0.0002687192
    257000 / 478769, 0.0002374350
    258000 / 478769, 0.0001188405
    259000 / 478769, 0.0003181596
    260000 / 478769, 0.0002098954
    261000 / 478769, 0.0003115956
    262000 / 478769, 0.0000713611
    263000 / 478769, 0.0000658669
    264000 / 478769, 0.0000792838
    265000 / 478769, 0.0004621624
    266000 / 478769, 0.0004280486
    267000 / 478769, 0.0003321384
    268000 / 478769, 0.0001673132
    269000 / 478769, 0.0000887993
    270000 / 478769, 0.0000845830
    271000 / 478769, 0.0000693241
    272000 / 478769, 0.0001379017
    273000 / 478769, 0.0002903245
    274000 / 478769, 0.0002103866
    275000 / 478769, 0.0001685253
    276000 / 478769, 0.0000471871
    277000 / 478769, 0.0001910802
    278000 / 478769, 0.0001545368
    279000 / 478769, 0.0000726421
    280000 / 478769, 0.0003706776
    281000 / 478769, 0.0002097775
    282000 / 478769, 0.0001204970
    283000 / 478769, 0.0001253122
    284000 / 478769, 0.0001048192
    285000 / 478769, 0.0001240431
    286000 / 478769, 0.0003320249
    287000 / 478769, 0.0000604929
    288000 / 478769, 0.0003665326
    289000 / 478769, 0.0003017480
    290000 / 478769, 0.0000735220
    291000 / 478769, 0.0002039460
    292000 / 478769, 0.0000723469
    293000 / 478769, 0.0001043268
    294000 / 478769, 0.0002647403
    295000 / 478769, 0.0001176578
    296000 / 478769, 0.0004005697
    297000 / 478769, 0.0000926692
    298000 / 478769, 0.0001768936
    299000 / 478769, 0.0000885737
    300000 / 478769, 0.0003588063
    301000 / 478769, 0.0000864405
    302000 / 478769, 0.0002341192
    303000 / 478769, 0.0000812490
    304000 / 478769, 0.0002080883
    305000 / 478769, 0.0000730130
    306000 / 478769, 0.0005355164
    307000 / 478769, 0.0000381832
    308000 / 478769, 0.0000631933
    309000 / 478769, 0.0002684426
    310000 / 478769, 0.0001108966
    311000 / 478769, 0.0000143062
    312000 / 478769, 0.0000837733
    313000 / 478769, 0.0001165400
    314000 / 478769, 0.0003737039
    315000 / 478769, 0.0000715345
    316000 / 478769, 0.0001101542
    317000 / 478769, 0.0000274832
    318000 / 478769, 0.0005832550
    319000 / 478769, 0.0002243998
    320000 / 478769, 0.0001659056
    321000 / 478769, 0.0001140977
    322000 / 478769, 0.0000696813
    323000 / 478769, 0.0000978014
    324000 / 478769, 0.0001306933
    325000 / 478769, 0.0002003123
    326000 / 478769, 0.0001197310
    327000 / 478769, 0.0000545317
    328000 / 478769, 0.0004195267
    329000 / 478769, 0.0003477322
    330000 / 478769, 0.0000619723
    331000 / 478769, 0.0006023692
    332000 / 478769, 0.0001249482
    333000 / 478769, 0.0002678133
    334000 / 478769, 0.0002012259
    335000 / 478769, 0.0001316315
    336000 / 478769, 0.0000695630
    337000 / 478769, 0.0001096378
    338000 / 478769, 0.0000708267
    339000 / 478769, 0.0000199817
    340000 / 478769, 0.0003340940
    341000 / 478769, 0.0001211530
    342000 / 478769, 0.0006738254
    343000 / 478769, 0.0000599353
    344000 / 478769, 0.0006511465
    345000 / 478769, 0.0003617066
    346000 / 478769, 0.0000434869
    347000 / 478769, 0.0000716250
    348000 / 478769, 0.0003606776
    349000 / 478769, 0.0004025252
    350000 / 478769, 0.0001117167
    351000 / 478769, 0.0000417683
    352000 / 478769, 0.0000188909
    353000 / 478769, 0.0000786730
    354000 / 478769, 0.0000983291
    355000 / 478769, 0.0000555389
    356000 / 478769, 0.0001201445
    357000 / 478769, 0.0001637083
    358000 / 478769, 0.0000702233
    359000 / 478769, 0.0000572274
    360000 / 478769, 0.0001693493
    361000 / 478769, 0.0000446440
    362000 / 478769, 0.0003010733
    363000 / 478769, 0.0000562239
    364000 / 478769, 0.0002073496
    365000 / 478769, 0.0000203700
    366000 / 478769, 0.0002346559
    367000 / 478769, 0.0000479669
    368000 / 478769, 0.0001153019
    369000 / 478769, 0.0001122367
    370000 / 478769, 0.0001278491
    371000 / 478769, 0.0002443215
    372000 / 478769, 0.0001630644
    373000 / 478769, 0.0001163957
    374000 / 478769, 0.0000848839
    375000 / 478769, 0.0000630332
    376000 / 478769, 0.0000458114
    377000 / 478769, 0.0000437738
    378000 / 478769, 0.0000073608
    379000 / 478769, 0.0005259587
    380000 / 478769, 0.0002021564
    381000 / 478769, 0.0003388440
    382000 / 478769, 0.0000727262
    383000 / 478769, 0.0000649415
    384000 / 478769, 0.0000240307
    385000 / 478769, 0.0000322553
    386000 / 478769, 0.0003604684
    387000 / 478769, 0.0000646129
    388000 / 478769, 0.0002190641
    389000 / 478769, 0.0000863954
    390000 / 478769, 0.0000832419
    391000 / 478769, 0.0001009488
    392000 / 478769, 0.0005569173
    393000 / 478769, 0.0001104695
    394000 / 478769, 0.0000932012
    395000 / 478769, 0.0004150065
    396000 / 478769, 0.0001139548
    397000 / 478769, 0.0001214280
    398000 / 478769, 0.0003697828
    399000 / 478769, 0.0000256053
    400000 / 478769, 0.0002000824
    401000 / 478769, 0.0001051067
    402000 / 478769, 0.0001679548
    403000 / 478769, 0.0001243517
    404000 / 478769, 0.0002421602
    405000 / 478769, 0.0001518641
    406000 / 478769, 0.0002384409
    407000 / 478769, 0.0001376766
    408000 / 478769, 0.0002424939
    409000 / 478769, 0.0000783248
    410000 / 478769, 0.0000554960
    411000 / 478769, 0.0001503048
    412000 / 478769, 0.0000440569
    413000 / 478769, 0.0002122000
    414000 / 478769, 0.0000592621
    415000 / 478769, 0.0000709168
    416000 / 478769, 0.0000235208
    417000 / 478769, 0.0001028531
    418000 / 478769, 0.0000977050
    419000 / 478769, 0.0000373011
    420000 / 478769, 0.0000723999
    421000 / 478769, 0.0001588263
    422000 / 478769, 0.0000303736
    423000 / 478769, 0.0000750340
    424000 / 478769, 0.0004104636
    425000 / 478769, 0.0001548859
    426000 / 478769, 0.0000539983
    427000 / 478769, 0.0001250313
    428000 / 478769, 0.0001153923
    429000 / 478769, 0.0002255330
    430000 / 478769, 0.0003895083
    431000 / 478769, 0.0000292332
    432000 / 478769, 0.0000458322
    433000 / 478769, 0.0000476175
    434000 / 478769, 0.0000622866
    435000 / 478769, 0.0006813942
    436000 / 478769, 0.0000289125
    437000 / 478769, 0.0001208999
    438000 / 478769, 0.0001246577
    439000 / 478769, 0.0001501882
    440000 / 478769, 0.0004132458
    441000 / 478769, 0.0001395835
    442000 / 478769, 0.0000690258
    443000 / 478769, 0.0008519820
    444000 / 478769, 0.0000363401
    445000 / 478769, 0.0000805624
    446000 / 478769, 0.0001430681
    447000 / 478769, 0.0000801698
    448000 / 478769, 0.0000612512
    449000 / 478769, 0.0000493287
    450000 / 478769, 0.0000588115
    451000 / 478769, 0.0000525372
    452000 / 478769, 0.0000547081
    453000 / 478769, 0.0001086641
    454000 / 478769, 0.0000476390
    455000 / 478769, 0.0001057836
    456000 / 478769, 0.0001041464
    457000 / 478769, 0.0000102033
    458000 / 478769, 0.0000209338
    459000 / 478769, 0.0001127455
    460000 / 478769, 0.0005546486
    461000 / 478769, 0.0000503670
    462000 / 478769, 0.0001241465
    463000 / 478769, 0.0000650835
    464000 / 478769, 0.0000255689
    465000 / 478769, 0.0000487211
    466000 / 478769, 0.0000404158
    467000 / 478769, 0.0000403452
    468000 / 478769, 0.0000831553
    469000 / 478769, 0.0001609566
    470000 / 478769, 0.0000474199
    471000 / 478769, 0.0002879203
    472000 / 478769, 0.0000969278
    473000 / 478769, 0.0000083798
    474000 / 478769, 0.0002241727
    475000 / 478769, 0.0000858267
    476000 / 478769, 0.0001033287
    477000 / 478769, 0.0005197478
    478000 / 478769, 0.0001578679

```python
# Test
total_loss = 0
with torch.no_grad():
    for idx, ϕ_state in enumerate(test_inputs):
        next_state_prediction = emulator(ϕ_state)

        next_state = test_outputs[idx]
        total_loss += criterion(next_state_prediction, next_state).item()

ave_test_loss = total_loss/test_size
print(f'Test loss: {ave_test_loss:.10f}')
```

    Test loss: 0.0000271195

```python
# Here you need to insert the code for training the controller
# by using the emulator for backpropagation

# If you succeed, feel free to send a PR
```

```python

```

    Test loss: 0.0639088841

```python

```

Using RNN(GRU) learning a car model.ipynb


# ARC Additions: Notebook of new functions and testing / benchmarking to verify parameters and sources

Goal: Transfer hyperfine atomic physics functions from internal Matlab toolbox to ARC.

Jonathan Pritchard, University of Strathclyde (2020)

```python
#Configure Environment

from arc import *

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np

cs = Cesium()
rb87 = Rubidium87()
```

## Check ARC Parameters for D1/2 and inverted (n+1)P_{1/2,3/2}

 + New function `A, B = self.getHFSCoefficients(self, n, l, j, s=None)` added to return HFS constants in Hz. Refernce [5] and [6] have useful reference info

### Hyperfine Structure Functions

 + `A, B = atom.getHFSCoefficients(n,l,j)`

Returns HFS constants in Hz. This is currently using hard-coded values for Cs and Rb, but has been written as a place-holder for adding in a table. Refernce [5] and [6] have useful reference info.

 + `getHFSEnergyShift(j, f, A, B)`

Returns HFS Energy shift from centre of mass.

\begin{equation}
\Delta E_\mathrm{hfs} = \frac{A}{2}K+B\frac{\frac{3}{2}K(K+1)-2I(I+1)J(J+1)}{2I(2I-1)2J(2J-1)},
\end{equation}
where $K=F(F+1)-I(I+1)-J(J+1)$.

### References
[1] Steck

[2] M. Ortiz and J. Campos, J. Quant. Spectrosc. Radiat. Transfer 26, 107 (1981)

[3] Vasilyev, A. A. and Savukov, I. M. and Safronova, M. S. and Berry, H. G., PRA 66, 020101 (2002)

[4] M. S. Safronova, Carl J. Williams, and Charles W. Clark, PRA 69, 022509 (2004)

[5] D. Feiertag, A. Sahm, and G. zu Putlitz, Z. Phys. 255, 93 (1972)

[6] E. Arrimondo, M. Inguscio and P. Violino, Rev. Mod. Phys. 49, 31 (1977)

```python
print("Cs Numbers")
print("==========")
print("State Lifetimes:")
# Lifetime Comparison -->
print("\t[1] : Tau6P1/2 = 34.894(44) ns, ARC = %2.3f ns (%2.2f %% error)"
      % (cs.getStateLifetime(6, 1, 0.5) * 1e9,
         (cs.getStateLifetime(6, 1, 0.5) * 1e9 - 34.894) / 34.894 * 100)
         )
print("\t[1] : Tau6P3/2 = 30.473(44) ns, ARC = %2.3f ns (%2.2f %% error)"
      % (cs.getStateLifetime(6, 1, 1.5) * 1e9,
         (cs.getStateLifetime(6, 1, 1.5) * 1e9 - 30.473) / 30.473 * 100)
        )
print("\t[2] : Tau7P1/2 = 155 ns, ARC = %2.3f ns (%2.2f %% error)"
      % (cs.getStateLifetime(7, 1, 0.5) * 1e9,
         (cs.getStateLifetime(7, 1, 0.5) * 1e9 - 155) / 155 * 100)
        )
print("\t[2] : Tau7P3/2 = 133 ns, ARC = %2.3f ns (%2.2f %% error)"
      % (cs.getStateLifetime(7, 1, 1.5) * 1e9,
         (cs.getStateLifetime(7, 1, 1.5) * 1e9 - 133) / 133 * 100)
         )
# Matrix Elements
print("Matrix Elements")
r = 4.5003; rm = cs.getReducedMatrixElementJ(6, 0, 0.5, 6, 1, 0.5)
print("\t[1] : <6S1/2||er||6P1/2> = %1.4f, ARC = %2.4f ns (%2.2f %% error)"
      % (r, rm, (rm - r) / r * 100)
     )
r = 6.3338; rm = cs.getReducedMatrixElementJ(6, 0, 0.5, 6, 1, 1.5)
print("\t[1] : <6S1/2||er||6P3/2> = %1.4f, ARC = %2.4f ns (%2.2f %% error)"
      % (r, rm, (rm - r) / r * 100)
      )
r = 0.2757; rm = cs.getReducedMatrixElementJ(6, 0, 0.5, 7,1,0.5)
print("\t[3] : <6S1/2||er||7P1/2> = %1.4f, ARC = %2.4f ns (%2.2f %% error)"
      % (r, rm, (rm - r) / r * 100)
      )
r = 0.5856; rm = cs.getReducedMatrixElementJ(6, 0, 0.5, 7, 1, 1.5)
print("\t[3] : <6S1/2||er||7P3/2> = %1.4f, ARC = %2.4f ns (%2.2f %% error)"
      % (r, rm, (rm - r) / r * 100)
      )
# Hyperfine Splitting Coefficients
print("Hyperfine Constants")
A, B = cs.getHFSCoefficients(6, 0, 0.5)
omegaQ = 4 * A
print("\t[1] : 6S1/2 A = %1.4f MHz,B = %1.4f MHz" % (A * 1e-6, B * 1e-6))
A, B = cs.getHFSCoefficients(6, 1, 0.5)
print("\t[1] : 6P1/2 A = %1.4f MHz,B = %1.4f MHz" % (A * 1e-6, B * 1e-6))
A, B = cs.getHFSCoefficients(6, 1, 1.5)
print("\t[1] : 6P3/2 A = %1.4f MHz,B = %1.4f MHz" % (A * 1e-6, B * 1e-6))
A, B = cs.getHFSCoefficients(7, 1, 0.5)
print("\t[5] : 7P1/2 A = %1.4f MHz,B = %1.4f MHz" % (A * 1e-6, B * 1e-6))
A, B = cs.getHFSCoefficients(7, 1, 1.5)
print("\t[6] : 7P3/2 A = %1.4f MHz,B = %1.4f MHz" % (A * 1e-6, B * 1e-6))

print("\nVerify getHFSCoefficients(j,f,A,B) using Cs 6P3/2")
A, B = cs.getHFSCoefficients(6, 1, 1.5)
Ehfs = cs.getHFSEnergyShift(1.5, 5, A, B)
print("\t[1] : 6P3/2 F=5: Ehfs = +263.81 MHz --> ARC Ehfs = %f MHz" % (Ehfs * 1e-6))
Ehfs = cs.getHFSEnergyShift(1.5, 4, A, B)
print("\t[1] : 6P3/2 F=4: Ehfs =  +12.815 MHz --> ARC Ehfs = %f MHz" % (Ehfs * 1e-6))
Ehfs = cs.getHFSEnergyShift(1.5, 3, A, B)
print("\t[1] : 6P3/2 F=3: Ehfs = -188.44 MHz --> ARC Ehfs = %f MHz" % (Ehfs * 1e-6))
Ehfs = cs.getHFSEnergyShift(1.5, 2, A, B)
print("\t[1] : 6P3/2 F=2: Ehfs = -339.64 MHz --> ARC Ehfs = %f MHz" % (Ehfs * 1e-6))

# Clock Frequency
print("\n\t[1] : 6S1/2 Clock Frequency = 9.192631770 GHz --> ARC = %f GHz (%2.2f %% error)"
      % (omegaQ * 1e-9, (omegaQ - 9.192631770e9) / 9.192631770e9 * 100)
     )

print("\nRb Numbers")
print("==========")
print("State Lifetimes:")

# Lifetime Comparison -->
print("\t[1] : Tau5P1/2 = 27.70(4) ns, ARC = %2.3f ns (%2.2f %% error)"
      % (rb87.getStateLifetime(5, 1, 0.5) * 1e9,
         (rb87.getStateLifetime(5, 1, 0.5) * 1e9 - 27.70) / 27.70 * 100)
      )
print("\t[1] : Tau5P3/2 = 26.24(4) ns, ARC = %2.3f ns (%2.2f %% error)"
      % (rb87.getStateLifetime(5, 1, 1.5) * 1e9,(rb87.getStateLifetime(5, 1, 1.5) * 1e9 - 26.24) / 26.24 * 100))
print("\t[4] : Tau6P1/2 = 125 ns, ARC = %2.3f ns (%2.2f %% error)"
       % (rb87.getStateLifetime(6, 1, 0.5) * 1e9,
         (rb87.getStateLifetime(6, 1, 0.5) * 1e9 - 125) / 125 * 100)
      )
print("\t[4] : Tau6P3/2 = 112 ns, ARC = %2.3f ns (%2.2f %% error)" %
      (rb87.getStateLifetime(6, 1, 1.5) * 1e9,
       (rb87.getStateLifetime(6, 1, 1.5) * 1e9 - 112) / 112 * 100)
      )

# Matrix Elements
print("Matrix Elements")
r = 4.2313
rm = rb87.getReducedMatrixElementJ(5, 0, 0.5, 5, 1, 0.5)
print("\t[4] : <5S1/2||er||5P1/2> = %1.4f, ARC = %2.4f ns (%2.2f %% error)"
      % (r, rm, (rm - r) / r * 100)
     )
r = 5.9780
rm = rb87.getReducedMatrixElementJ(5, 0, 0.5, 5, 1, 1.5)
print("\t[4] : <5S1/2||er||5P3/2> = %1.4f, ARC = %2.4f ns (%2.2f %% error)"
      % (r, rm, (rm - r) / r * 100)
     )
r = 0.333
rm = rb87.getReducedMatrixElementJ(5, 0, 0.5, 6, 1, 0.5)
print("\t[4] : <5S1/2||er||6P1/2> = %1.4f, ARC = %2.4f ns (%2.2f %% error)"
      % (r, rm, (rm - r) / r * 100)
     )
r = 0.541
rm = rb87.getReducedMatrixElementJ(5, 0, 0.5, 6, 1, 1.5)
print("\t[4] : <5S1/2||er||6P3/2> = %1.4f, ARC = %2.4f ns (%2.2f %% error)"
      % (r, rm, (rm - r) / r * 100)
     )
print("Hyperfine Constants")
A, B = rb87.getHFSCoefficients(5, 0, 0.5)
omegaQ = (0.5 + rb87.I) * A
print("\t[1] : 5S1/2 A = %1.4f MHz,B = %1.4f MHz"
      % (A * 1e-6, B * 1e-6))
A, B = rb87.getHFSCoefficients(5, 1, 0.5)
print("\t[1] : 5P1/2 A = %1.4f MHz,B = %1.4f MHz"
      % (A * 1e-6, B * 1e-6))
A, B = rb87.getHFSCoefficients(5, 1, 1.5)
print("\t[1] : 5P3/2 A = %1.4f MHz,B = %1.4f MHz"
      % (A * 1e-6, B * 1e-6)
     )
A, B = rb87.getHFSCoefficients(6, 1, 0.5)
print("\t[5] : 6P1/2 A = %1.4f MHz,B = %1.4f MHz"
      % (A * 1e-6, B * 1e-6)
     )
A, B = rb87.getHFSCoefficients(6, 1, 1.5)
print("\t[6] : 6P3/2 A = %1.4f MHz,B = %1.4f MHz"
      % (A * 1e-6, B * 1e-6)
     )

# Clock Frequency
print("\n\t[1] : 5S1/2 Clock Frequency =  6.83468261090429 GHz --> ARC = %f GHz (%2.2f %% error)"
      % (omegaQ * 1e-9, (omegaQ - 6.83468261090429e9) / 6.83468261090429e9 * 100)
      )
```

    Cs Numbers
    ==========
    State Lifetimes:
    	[1] : Tau6P1/2 = 34.894(44) ns, ARC = 34.894 ns (0.00 % error)
    	[1] : Tau6P3/2 = 30.473(44) ns, ARC = 30.474 ns (0.00 % error)
    	[2] : Tau7P1/2 = 155 ns, ARC = 154.368 ns (-0.41 % error)
    	[2] : Tau7P3/2 = 133 ns, ARC = 129.775 ns (-2.42 % error)
    Matrix Elements
    	[1] : <6S1/2||er||6P1/2> = 4.5003, ARC = 4.5003 ns (0.00 % error)
    	[1] : <6S1/2||er||6P3/2> = 6.3338, ARC = 6.3337 ns (-0.00 % error)
    	[3] : <6S1/2||er||7P1/2> = 0.2757, ARC = 0.2757 ns (0.00 % error)
    	[3] : <6S1/2||er||7P3/2> = 0.5856, ARC = 0.5856 ns (0.00 % error)
    Hyperfine Constants
    	[1] : 6S1/2 A = 2298.1579 MHz,B = 0.0000 MHz
    	[1] : 6P1/2 A = 291.9201 MHz,B = 0.0000 MHz
    	[1] : 6P3/2 A = 50.2750 MHz,B = -0.5300 MHz
    	[5] : 7P1/2 A = 94.3500 MHz,B = 0.0000 MHz
    	[6] : 7P3/2 A = 16.6050 MHz,B = -0.1500 MHz

    Verify getHFSCoefficients(j,f,A,B) using Cs 6P3/2
    	[1] : 6P3/2 F=5: Ehfs = +263.81 MHz --> ARC Ehfs = 263.811250 MHz
    	[1] : 6P3/2 F=4: Ehfs =  +12.815 MHz --> ARC Ehfs = 12.814821 MHz
    	[1] : 6P3/2 F=3: Ehfs = -188.44 MHz --> ARC Ehfs = -188.436607 MHz
    	[1] : 6P3/2 F=2: Ehfs = -339.64 MHz --> ARC Ehfs = -339.640179 MHz

    	[1] : 6S1/2 Clock Frequency = 9.192631770 GHz --> ARC = 9.192632 GHz (0.00 % error)

    Rb Numbers
    ==========
    State Lifetimes:
    	[1] : Tau5P1/2 = 27.70(4) ns, ARC = 27.690 ns (-0.04 % error)
    	[1] : Tau5P3/2 = 26.24(4) ns, ARC = 26.238 ns (-0.01 % error)
    	[4] : Tau6P1/2 = 125 ns, ARC = 129.295 ns (3.44 % error)
    	[4] : Tau6P3/2 = 112 ns, ARC = 118.143 ns (5.48 % error)
    Matrix Elements
    	[4] : <5S1/2||er||5P1/2> = 4.2313, ARC = 4.2321 ns (0.02 % error)
    	[4] : <5S1/2||er||5P3/2> = 5.9780, ARC = 5.9783 ns (0.01 % error)
    	[4] : <5S1/2||er||6P1/2> = 0.3330, ARC = 0.3330 ns (0.00 % error)
    	[4] : <5S1/2||er||6P3/2> = 0.5410, ARC = 0.5410 ns (0.00 % error)
    Hyperfine Constants
    	[1] : 5S1/2 A = 3417.3413 MHz,B = 0.0000 MHz
    	[1] : 5P1/2 A = 408.3280 MHz,B = 0.0000 MHz
    	[1] : 5P3/2 A = 84.7185 MHz,B = 12.4965 MHz
    	[5] : 6P1/2 A = 132.5600 MHz,B = 0.0000 MHz
    	[6] : 6P3/2 A = 27.6740 MHz,B = 3.9450 MHz

    	[1] : 5S1/2 Clock Frequency =  6.83468261090429 GHz --> ARC = 6.834683 GHz (-0.00 % error)

### Angular Momentum Functions

 + `getSphericalDipoleMatrixElement(j1, mj1, j2, mj2, q)`

    Angular part of spherical matrix element $A$ in units of reduced matrix element $\langle j\vert\vert\mu\vert\vert j'\rangle$ (Wigner-Eckart decomposition)

\begin{equation}
\langle j m_j\vert \mu_q \vert j' m_j' \rangle = (-1)^{j-m_j}\begin{pmatrix}j&1&j'\\-m_j&-q&m_j'\end{pmatrix}\langle j\vert\vert\mu\vert\vert j'\rangle = A\langle j\vert\vert\mu\vert\vert j'\rangle,
\end{equation}

 + `_reducedMatrixElementFJ(j1, f1, j2, f2)`

 Reduced dipole matrix element $\langle f \vert\vert\mu\vert\vert f'\rangle$ in units of $\langle j\vert\vert\mu\vert\vert j'\rangle$

 \begin{equation}
\langle f \vert\vert\mu\vert\vert f'\rangle = (-1)^{j+I+f'+1}\sqrt{(2f+1)(2f'+1)}
\begin{Bmatrix}f&1&f'\\j'&I&j\end{Bmatrix}\langle j\vert\vert\mu\vert\vert j'\rangle, \label{eq:frf}
\end{equation}

 + `getSphericalMatrixElementHFStoFS(j1, f1, mf1, j2, mj2, q)`

 Angular coefficient of matrix element $\langle jfm_f \vert er_q\vert j'f'm_f'\rangle$ for coupling to state without resolved HFS, in units of $\langle j\vert\vert\mu\vert\vert j'\rangle$

\begin{align}
\langle f,m_f \vert\mu_q\vert j',m_j'\rangle &=\displaystyle\sum_{f'}C_{j'm_j'I(m_f+q-m_j')}^{f'(m_f+q)}(-1)^{f-m_f}\begin{pmatrix}f&1&f'\\-m_f&-q&m_f+q\end{pmatrix}\\&
\times(-1)^{j+I+f'+1}\sqrt{(2f+1)(2f'+1)}
\begin{Bmatrix}f&1&f'\\j'&I&j\end{Bmatrix}\langle j\vert\vert\mu\vert\vert j'\rangle
\end{align}

+ `getMagneticDipoleMatrixElementHFS(l, j, f1, mf1, m2, mf2, q)`

Returns spherical component of magnetic dipole transition from $\vert f_1,m_{f_1}\rangle\rightarrow\vert f_2,m_{f_2}\rangle$ in units of $\mu_B B_q$ corresponding to relative strength of microwave transitions between hyperfine states of the same $n\ell j$ state driven by magnetic field with polarisation $q$.

\begin{align}
\langle f_1,m_{f_1}\vert \mu_q \vert f_2,m_{f_2}\rangle = g_J \mu_B B_q (-1)^{f_2+j+I+1+f_1-m_{f_1}}\sqrt{(2f_1+1)(2f_2+1)j(j+1)(2j+1)} \begin{pmatrix}f_1&1&f_2\\-m_{f_1} & -q & m_{f_2}\end{pmatrix}\begin{Bmatrix}f_1&1&f_2\\j & I & j\end{Bmatrix}
\end{align}

 + `getDipoleMatrixElementHFStoFS(n1, l1, j1, f1, mf1, n2, l2, j2, mj2, q)`

 Full dipole matrix element $\langle jfm_f \vert er_q\vert j'f'm_f'\rangle$ for coupling to state without resolved HFS, in units of $ea_0$

\begin{align}
\langle f,m_f \vert\mu_q\vert j',m_j'\rangle &=\displaystyle\sum_{f'}C_{j'm_j'I(m_f+q-m_j')}^{f'(m_f+q)}(-1)^{f-m_f}\begin{pmatrix}f&1&f'\\-m_f&-q&m_f+q\end{pmatrix}\\&
\times(-1)^{j+I+f'+1}\sqrt{(2f+1)(2f'+1)}
\begin{Bmatrix}f&1&f'\\j'&I&j\end{Bmatrix}\langle j\vert\vert\mu\vert\vert j'\rangle
\end{align}

 + `getBranchingRatio(jg, fg, mfg, je, fe, mfe)`

 Branching ratio $b$ for fraction of spontaneous decay from $\vert{j_e,f_e,m_{f_e}}\rangle\rightarrow\vert{j_g,f_g,m_{f_g}}\rangle$.

\begin{align}
b = \displaystyle\sum_q (2j_e+1)\left(\begin{matrix}f_g & 1 & f_e \\ -m_{f_g} & q & m_{f_e}\end{matrix}\right)^2 \vert \langle j_e,f_e\vert \vert er \vert\vert j_g,f_g\rangle\vert^2
\end{align}

 + `getSaturationIntensity(ng, lg, jg, fg, mfg, ne, fe, je, fe, mfe)`

 Returns saturation intenstiy $I_\mathrm{sat}$ for transition from $\vert{j_g,f_g,m_{f_g}}\rangle\rightarrow\vert{j_e,f_e,m_{f_e}}\rangle$ in units of W/m$^2$.

\begin{align}
I_\mathrm{sat} = \frac{c\epsilon_0\Gamma^2\hbar^2}{4\vert \epsilon_q\cdot\mathrm{d}\vert^2}
\end{align}

 + `getSaturationIntensityIsotropic(ng, lg, jg, fg, ne, fe, je, fe)`

 Returns saturation intenstiy $I_\mathrm{sat}$ averaged over transiton from $f_g\rightarrow f_e$ for isotropic polarisation in units of W/m$^2$.

\begin{align}
I_\mathrm{sat,iso} = \frac{c\epsilon_0\Gamma^2\hbar^2}{4\vert \mathrm{d_{iso}}\vert^2},
\end{align}
where $\vert \mathrm{d_{iso}}\vert^2$ is the isotropic dipole moment obtained by averaging over all possible transitions and polarisations.
\begin{align}
\vert d_{iso} \vert^2 = \frac{1}{3(2fg+1)}\displaystyle\sum_{q,m_{f_g}}\vert\langle ng,lg,jg,fg,mfg|er|ne,le,je,fe,mfg+q\rangle\vert^2
\end{align}

```python
# Test sph=atm.getSphericalDipoleMatrixElement(j1,mj1,j2,mj2,q)
# ============================================================
print("\nTest sph=atm.getSphericalDipoleMatrixElement(j1,mj1,j2,mj2,q):")
print("==============================================================")
print("Spherical matrix element in units of reduced matrix element <j'||er||j> for j and  <f'||er||f> for f")
sph = cs.getSphericalDipoleMatrixElement(0.5, 0.5, 1.5, 1.5, 1)
print("|1/2,1/2> to |3/2,3/2> for q=1 (sigma+) : sph = 1/2 = %2.5f, sph = %2.5f"
      % (1 / 2, sph)
     )
sph = cs.getSphericalDipoleMatrixElement(4, 4, 5, 5, 1)
print("|4,4> to |5,5> for q=1 (sigma+) : sph = sqrt{1/11} = %2.5f, sph = %2.5f"
      % (np.sqrt(1 / 11), sph)
     )
sph = cs.getSphericalDipoleMatrixElement(4, 4, 4, 4, 0)
print("|4,4> to |4,4> for q=0 (pi) : sph = sqrt{4/45} = %2.5f, sph = %2.5f"
      % (np.sqrt(4 / 45), sph))
sph = cs.getSphericalDipoleMatrixElement(4, 4, 3, 3, -1)
print("|4,4> to |3,3> for q=-1 (sigma-) : sph = 1/3 = %2.5f, sph = %2.5f"
      % (1 / 3, sph)
     )

# Test sph=atm.getSphericalMatrixElementHFStoFS(j1,f1,mf1,j2,mj2,q)
# =================================================================
print("\nTest sph=atm.getSphericalMatrixElementHFStoFS(j1,f1,mf1,j2,mj2,q):")
print("==================================================================")
print("Spherical matrix element in units of reduced matrix element <j'||er||j'>")
sph = cs.getSphericalMatrixElementHFStoFS(1.5, 5, 5, 0.5, 0.5, -1)
print("Cs P3/2 |5,5> to S1/2 |1/2,1/2> for q=-1 (sigma-) : sph = 1/2 = %2.5f, sph = %2.5f"
       % (1 / 2, sph)
     )
sph = cs.getSphericalMatrixElementHFStoFS(1.5, 5, 1, 0.5, 0.5, 1)
print("Cs P3/2 |5,1> to S1/2 |1/2,1/2> for q=1 (sigma+) : sph = sqrt(1/40) = %2.5f, sph = %2.5f"
       % (np.sqrt(1 / 40), sph)
     )
sph = cs.getSphericalMatrixElementHFStoFS(1.5, 5, 1, 0.5, 0.5, 0)
print("Cs P3/2 |5,1> to S1/2 |1/2,1/2> for q=0 (pi) : sph = sqrt(1/12) = %2.5f, sph = %2.5f"
       % (np.sqrt(1 / 12), sph)
     )
sph = cs.getSphericalMatrixElementHFStoFS(1.5, 5, 1, 0.5, 0.5, -1)
print("Cs P3/2 |5,1> to S1/2 |1/2,1/2> for q=-1 (sigma-) : sph = sqrt(1/24) = %2.5f, sph = %2.5f"
      % (np.sqrt(1 / 24), sph))

# Test sph=atm.getDipoleMatrixElementHFStoFS(j1,f1,mf1,j2,mj2,q)
# ==============================================================
print("\nTest rme=atm.getDipoleMatrixElementHFStoFS(n1,l1,j1,f1,mf1,n2,l2,j2,mj2,q):")
print("=============================================================================")
print("Spherical matrix element in ea0")
rmeJ = cs.getReducedMatrixElementJ(6, 1, 1.5, 80, 0, 0.5)
print("Cs <6P3/2||er||80S1/2> = 0.0086 = %f" % (rmeJ))
rme = cs.getDipoleMatrixElementHFStoFS(6, 1, 1.5, 5, 5, 80, 0, 0.5, 0.5, -1)
print("Cs 6P3/2 |5,5> to 80S1/2 |1/2,1/2> for q=-1 (sigma-) : 0.0043 ea0, rme = %f ea0"
      % (rme))
rme = cs.getDipoleMatrixElementHFStoFS(6, 1, 1.5, 5, 1, 80, 0, 0.5, 0.5, -1)
print("Cs 6P3/2 |5,1> to 80S1/2 |1/2,1/2> for q=-1 (sigma-) : 0.0018 ea0, rme = %f ea0"
      % (rme))

rmeJ = cs.getReducedMatrixElementJ(6, 1, 1.5, 80, 2, 2.5)
print("Cs <6P3/2||er||80D5/2> = -0.0259 = %f"
      % (rmeJ))
rme = cs.getDipoleMatrixElementHFStoFS(6, 1, 1.5, 5, 5, 80, 2, 2.5, 2.5, 1)
print("Cs 6P3/2 |5,5> to 80D5/2 |5/2,5/2> for q=+1 (sigma+) : -0.0106 ea0, rme = %f ea0"
      % (rme))

sph = cs.getSphericalMatrixElementHFStoFS(1.5, 5, 5, 2.5, 2.5, 1)
print("Cs P3/2 |5,5> to D5/2 |5/2,5/2> for q=+1 (sigma+) : sph = sqrt(1/6) = %2.5f, sph = %2.5f"
       % (np.sqrt(1 / 6), sph))

# Test sph=atm.getMagneticDipoleMatrixElementHFS(l,j,f1,mf1,f2,mf2,q)
#==================================================================
print("\nTest sphB=atm.getMagneticDipoleMatrixElementHFS(l,j,f1,mf1,f2,mf2,q):")
print("=====================================================================")
print("Spherical matrix element in uB*B")
sphB = cs.getMagneticDipoleMatrixElementHFS(0, 0.5, 4, 0, 3, 0, 0)
print("Cs 6S1/2 <4,0|B_0|3,0> = 1.000 = %f muB"
      % (sphB))
sphB = cs.getMagneticDipoleMatrixElementHFS(0, 0.5, 4, 3, 3, 3, 0)
print("Cs 6S1/2 <4,3|B_0|3,3> = sqrt(7/16) = %f = %f muB"
      % (np.sqrt(7 / 16), sphB))
sphB = cs.getMagneticDipoleMatrixElementHFS(0, 0.5, 4, 4, 3, 3, -1)
print("Cs 6S1/2 <4,4|B_1|3,3> = sqrt(7/4) = %f = %f muB"
      % (np.sqrt(7 / 4), sphB))
sphB = rb87.getMagneticDipoleMatrixElementHFS(0, 0.5, 2, 0, 1, 0, 0)
print("Rb 5S1/2 <2,0|B_0|1,0> = 1.000 = %f muB"
       % (sphB))
sphB = rb87.getMagneticDipoleMatrixElementHFS(0, 0.5, 2, 1, 1, 1, 0)
print("Rb 5S1/2 <2,1|B_0|1,1> = sqrt(3/4) = %f = %f muB"
      % (np.sqrt(3 / 4), sphB))
sphB = rb87.getMagneticDipoleMatrixElementHFS(0, 0.5, 2, 2, 1, 1, -1)
print("Rb 5S1/2 <2,2|B_1|1,1> = sqrt(3/2) = %f = %f muB"
      % (np.sqrt(3 / 2), sphB))

# Test getBranchingRatio(jg,fg,mfg,je,fe,mfe)
# ===========================================
print("\nTest br=atm.getBranchingRatio(jg,fg,mfg,je,fe,mfe):")
print("===================================================")
# Cs P3/2 |5,5> to S1/2|4,4> = 1
br = cs.getBranchingRatio(0.5, 4, 4, 1.5, 5, 5)
print("Cs P3/2 |5,5> to S1/2|4,4> = 1, br = %2.5f"
      % (br))
# Cs P3/2 |5,4> to S1/2|4,4> = 1/5
br = cs.getBranchingRatio(0.5, 4, 4, 1.5, 5, 4)
print("Cs P3/2 |5,4> to S1/2|4,4> = 1/5 = %2.5f, br = %2.5f"
      % (1 / 5, br))
# Cs P3/2 |5,4> to S1/2|4,3> = 4/5
br = cs.getBranchingRatio(0.5, 4, 3, 1.5, 5, 4)
print("Cs P3/2 |5,4> to S1/2|4,3> = 4/5 = %2.5f, br = %2.5f"
      % (4 / 5, br))
# Cs P3/2 |4,4> to S1/2|4,4> = 7/15
br = cs.getBranchingRatio(0.5, 4, 4, 1.5, 4, 4)
print("Cs P3/2 |4,4> to S1/2|4,4> = 7/15 = %2.5f, br = %2.5f"
      % (7 / 15, br))
# Cs P3/2 |4,4> to S1/2|4,3> = 7/60
br = cs.getBranchingRatio(0.5, 4, 3, 1.5, 4, 4)
print("Cs P3/2 |4,4> to S1/2|4,3> = 7/60 = %2.5f, br = %2.5f"
      % (7 / 60, br))
# Cs P3/2 |4,4> to S1/2|3,3> = 7/60
br = cs.getBranchingRatio(0.5, 3, 3, 1.5, 4, 4)
print("Cs P3/2 |4,4> to S1/2|3,3> = 5/12 = %2.5f, br = %2.5f"
      % (5 / 12, br))

# Test getSaturationIntensity(ng,lg,jg,fg,mfg,ne,le,je,fe,mfe)
#============================================================
print("\nTest Is=atm.getSaturationIntensity(ng,lg,jg,fg,mfg,ne,le,je,fe,mfe):")
print("====================================================================")
# Cs P3/2 |5,5> to S1/2|4,4> = 11 W/m^2
Is = cs.getSaturationIntensity(6, 0, 0.5, 4, 4, 6, 1, 1.5, 5, 5)
print("Cs 6S1/2|4,4> to 6P3/2 |5,5> Is = 11 W/m^2, br = %2.5f W/m^2"
      %(Is))
Is = cs.getSaturationIntensity(6, 0, 0.5, 4, 0, 6, 1, 1.5, 5, 0)
print("Cs 6S1/2|4,0> to 6P3/2 |5,0> Is = 19.8 W/m^2, br = %2.5f W/m^2"
      %(Is))

# Test getSaturationIntensityIsotropic(ng,lg,jg,fg,ne,le,je,fe)
# ============================================================
print("\nTest Is=atm.getSaturationIntensityIsotropic(ng,lg,jg,fg,ne,le,je,fe):")
print("======================================================================")
# Cs P3/2 |5,5> to S1/2|4,4> = 11 W/m^2
Is = cs.getSaturationIntensityIsotropic(6, 0, 0.5, 4, 6, 1, 1.5, 5)
print("Cs 6S1/2 F=4 to 6P3/2 F=5 Is_ISO = 27.059 W/m^2 [1], Is_ISO = %2.5f W/m^2"
      %(Is))
Is = cs.getSaturationIntensityIsotropic(6, 0, 0.5, 4, 6, 1, 1.5, 4)
print("Cs 6S1/2 F=4 to 6P3/2 F=4 Is_ISO = 56.69 W/m^2, Is_ISO = %2.5f W/m^2"
      %(Is))

```

    Test sph=atm.getSphericalDipoleMatrixElement(j1,mj1,j2,mj2,q):
    ==============================================================
    Spherical matrix element in units of reduced matrix element <j'||er||j> for j and  <f'||er||f> for f
    |1/2,1/2> to |3/2,3/2> for q=1 (sigma+) : sph = 1/2 = 0.50000, sph = 0.50000
    |4,4> to |5,5> for q=1 (sigma+) : sph = sqrt{1/11} = 0.30151, sph = 0.30151
    |4,4> to |4,4> for q=0 (pi) : sph = sqrt{4/45} = 0.29814, sph = 0.29814
    |4,4> to |3,3> for q=-1 (sigma-) : sph = 1/3 = 0.33333, sph = 0.33333

    Test sph=atm.getSphericalMatrixElementHFStoFS(j1,f1,mf1,j2,mj2,q):
    ==================================================================
    Spherical matrix element in units of reduced matrix element <j'||er||j'>
    Cs P3/2 |5,5> to S1/2 |1/2,1/2> for q=-1 (sigma-) : sph = 1/2 = 0.50000, sph = 0.50000
    Cs P3/2 |5,1> to S1/2 |1/2,1/2> for q=1 (sigma+) : sph = sqrt(1/40) = 0.15811, sph = 0.15811
    Cs P3/2 |5,1> to S1/2 |1/2,1/2> for q=0 (pi) : sph = sqrt(1/12) = 0.28868, sph = 0.28868
    Cs P3/2 |5,1> to S1/2 |1/2,1/2> for q=-1 (sigma-) : sph = sqrt(1/24) = 0.20412, sph = 0.20412

    Test rme=atm.getDipoleMatrixElementHFStoFS(n1,l1,j1,f1,mf1,n2,l2,j2,mj2,q):
    =============================================================================
    Spherical matrix element in ea0
    Cs <6P3/2||er||80S1/2> = 0.0086 = -0.008715
    Cs 6P3/2 |5,5> to 80S1/2 |1/2,1/2> for q=-1 (sigma-) : 0.0043 ea0, rme = -0.004358 ea0
    Cs 6P3/2 |5,1> to 80S1/2 |1/2,1/2> for q=-1 (sigma-) : 0.0018 ea0, rme = -0.001779 ea0
    Cs <6P3/2||er||80D5/2> = -0.0259 = -0.025645
    Cs 6P3/2 |5,5> to 80D5/2 |5/2,5/2> for q=+1 (sigma+) : -0.0106 ea0, rme = -0.010469 ea0
    Cs P3/2 |5,5> to D5/2 |5/2,5/2> for q=+1 (sigma+) : sph = sqrt(1/6) = 0.40825, sph = 0.40825

    Test sphB=atm.getMagneticDipoleMatrixElementHFS(l,j,f1,mf1,f2,mf2,q):
    =====================================================================
    Spherical matrix element in uB*B
    Cs 6S1/2 <4,0|B_0|3,0> = 1.000 = 1.000000 muB
    Cs 6S1/2 <4,3|B_0|3,3> = sqrt(7/16) = 0.661438 = 0.661438 muB
    Cs 6S1/2 <4,4|B_1|3,3> = sqrt(7/4) = 1.322876 = 1.322876 muB
    Rb 5S1/2 <2,0|B_0|1,0> = 1.000 = 1.000000 muB
    Rb 5S1/2 <2,1|B_0|1,1> = sqrt(3/4) = 0.866025 = 0.866025 muB
    Rb 5S1/2 <2,2|B_1|1,1> = sqrt(3/2) = 1.224745 = 1.224745 muB

    Test br=atm.getBranchingRatio(jg,fg,mfg,je,fe,mfe):
    ===================================================
    Cs P3/2 |5,5> to S1/2|4,4> = 1, br = 1.00000
    Cs P3/2 |5,4> to S1/2|4,4> = 1/5 = 0.20000, br = 0.20000
    Cs P3/2 |5,4> to S1/2|4,3> = 4/5 = 0.80000, br = 0.80000
    Cs P3/2 |4,4> to S1/2|4,4> = 7/15 = 0.46667, br = 0.46667
    Cs P3/2 |4,4> to S1/2|4,3> = 7/60 = 0.11667, br = 0.11667
    Cs P3/2 |4,4> to S1/2|3,3> = 5/12 = 0.41667, br = 0.41667

    Test Is=atm.getSaturationIntensity(ng,lg,jg,fg,mfg,ne,le,je,fe,mfe):
    ====================================================================
    Cs 6S1/2|4,4> to 6P3/2 |5,5> Is = 11 W/m^2, br = 11.02370 W/m^2
    Cs 6S1/2|4,0> to 6P3/2 |5,0> Is = 19.8 W/m^2, br = 19.84266 W/m^2

    Test Is=atm.getSaturationIntensityIsotropic(ng,lg,jg,fg,ne,le,je,fe):
    ======================================================================
    Cs 6S1/2 F=4 to 6P3/2 F=5 Is_ISO = 27.059 W/m^2 [1], Is_ISO = 27.05817 W/m^2
    Cs 6S1/2 F=4 to 6P3/2 F=4 Is_ISO = 56.69 W/m^2, Is_ISO = 56.69332 W/m^2

### Magnetic Dipole Matrix Element

+ `getMagneticDipoleMatrixElementHFS(l, j, f1, mf1, m2, mf2, q)`

Returns spherical component of magnetic dipole transition from $\vert f_1,m_{f_1}\rangle\rightarrow\vert f_2,m_{f_2}\rangle$ in units of $\mu_B B_q$ corresponding to relative strength of microwave transitions between hyperfine states of the same $n\ell j$ state driven by magnetic field with polarisation $q$.

\begin{align}
\langle f_1,m_{f_1}\vert \mu_q \vert f_2,m_{f_2}\rangle = g_J \mu_B B_q (-1)^{f_2+j+I+1+f_1-m_{f_1}}\sqrt{(2f_1+1)(2f_2+1)j(j+1)(2j+1)} \begin{pmatrix}f_1&1&f_2\\-m_{f_1} & -q & m_{f_2}\end{pmatrix}\begin{Bmatrix}f_1&1&f_2\\j & I & j\end{Bmatrix}
\end{align}

## Ground-State Raman Transitions

We consider here the case of two-photon ground-state Raman transitions from $\vert 0 \rangle = \vert f_0,m_{f_0}\rangle$ to $\vert 1 \rangle = \vert f_1,m_{f_1}\rangle$ via an intermediate excited state $\vert e\rangle = \vert f_e,m_{f_e}\rangle$. In the rotating wave approximation, the Hamiltonian can be written as

\begin{equation}
\mathcal{H} = \hbar\begin{pmatrix} 0 & \Omega^a_{0\rightarrow f_e}/2 & 0 \\\Omega^a_{0\rightarrow f_e}/2 & -\Delta & \Omega^b_{1\rightarrow f_e}/2 \\ 0 & \Omega^b_{1\rightarrow f_e}/2 & -\delta\end{pmatrix},
\end{equation}

where $\Omega^a_{0\rightarrow f_e}$ is the Rabi frequency from $\vert 0\rangle\rightarrow\vert e\rangle$ due to laser A and $\Omega^b_{1\rightarrow f_e}$ is the Rabi frequency from $\vert 1\rangle\rightarrow\vert e\rangle$ due to laser B, $\Delta=(\Delta_a+\Delta_a)/2$ is the average detuning from the intermediate state and $\delta=\Delta_a-\Delta_b$ is the two-photon detuning.

In the limit $\Delta \gg \Omega^a,\Omega^b,\Gamma_e$ the intermediate excited state can be adiabatically elliminiated leading to an effective two-photon Rabi frequency $\Omega_R = \Omega^a\Omega^b/2\Delta$ and an effective detuning
$\delta_\mathrm{eff}=\delta+\Delta_\mathrm{AC}$ where $\Delta_\mathrm{AC}$ is the differential Stark shift given by

\begin{equation}
\Delta_\mathrm{AC} = \frac{\vert\Omega^b_{0\rightarrow f_e}\vert^2}{4\Delta} - \frac{\vert\Omega^a_{1\rightarrow f_e}\vert^2}{4\Delta}.
\end{equation}

On the two-photon transition ($\delta_\mathrm{eff}=0$), the probability of spontanteous emission during a $\pi$-pulse can be evaluated from $P_\mathrm{sc}=\Gamma_e\rho_e\tau_\pi$, where $\tau_\pi=\pi/\Omega_R$ is the $\pi$-pulse duration and $\rho_e$ is the excited state population given by

\begin{equation}
\rho_e = \frac{1}{2}\left(\frac{\vert\Omega^a_{0\rightarrow f_e}\vert^2}{2\Delta^2}+\frac{\vert\Omega^b_{1\rightarrow f_e}\vert^2}{2\Delta^2}\right).
\end{equation}

For a multi-level atom, the terms describing Raman transition, diffrential AC Stark shift and scattering probability must be calculated by summing over the $\vert f_e,m_{f_e}\rangle$ states of the intermediate excited state. Additionally, for transitions between hyperfine ground states separated by frequency $\omega_{01}$ we must also account for the off-resonant AC Stark shift and scattering terms arising from laser A acting on $\vert 1 \rangle$ with detuning $\Delta'=\Delta+\omega_{01}$ and laser B acting on $\vert 0 \rangle$ with detuning $\Delta'=\Delta-\omega_{01}$.

\begin{align}
\Omega_R&=\displaystyle\sum_{f_e,m_{f_e}}\frac{\Omega^a_{0\rightarrow f_e}\Omega^b_{1\rightarrow f_e}}{2(\Delta-\Delta_{f_e})}\\
\Delta_{\mathrm{AC}} &= \displaystyle\sum_{f_e,m_{f_e}}\left[\frac{\vert\Omega^a_{0\rightarrow f_e}\vert^2-\vert\Omega^b_{1\rightarrow f_e}\vert^2}{4(\Delta-\Delta_{f_e})}+\frac{\vert\Omega^a_{1\rightarrow f_e}\vert^2}{4(\Delta+\omega_{01}-\Delta_{f_e})}-\frac{\vert\Omega^b_{0\rightarrow f_e}\vert^2}{4(\Delta-\omega_{01}-\Delta_{f_e})}\right],\\
P_\mathrm{sc} &= \frac{\Gamma_et_\pi}{2}\displaystyle\sum_{f_e,m_{f_e}}\left[\frac{\vert\Omega^a_{0\rightarrow f_e}\vert^2}{2(\Delta-\Delta_{f_e})^2}+\frac{\vert\Omega^b_{1\rightarrow f_e}\vert^2}{2(\Delta-\Delta_{f_e})^2}+\frac{\vert\Omega^a_{1\rightarrow f_e}\vert^2}{4(\Delta+\omega_{01}-\Delta_{f_e})^2}+\frac{\vert\Omega^b_{0\rightarrow f_e}\vert^2}{4(\Delta-\omega_{01}-\Delta_{f_e})^2}\right],
\end{align}

 + `OmegaR, AC, Psc = groundStateRamanTransition(Pa, wa, qa, Pb, wb, qb, Delta, f0, mf0, f1, mf1, ne, le, je)`

Returns two-photon Rabi frequency $\Omega_R$, differential AC Starkshift $\Delta_\mathrm{AC}$ and scattering probability $P_\mathrm{sc}$ for transition from $\vert f_0,m_{f_0}\rangle\rightarrow\vert f_1,m_{f_1}\rangle$ via intermediate excited state with quantum numbers $n_e,\ell_e,j_e$ with detuning $\Delta$ from excited state centre of mass using laser $a/b$ with power $P$, waist $w$ and polarisation $q$ driving transitions from $f_0,f_1$ respectively.

```python
# Test Raman Transition Calculation
#==================================

#Laser Parameters
Pa = 0.1e-3; wa = 5e-6; qa = -1;
Pb = 0.1e-3; wb = 5e-6; qb = -1;

#Detuning Array
Delta0 = -50e9 * 2 * np.pi
Delta = np.linspace(-150, 150, 10000) * 1e9 * 2 * np.pi

#D1 Line
ne = 6; le = 1; je = 0.5; f1 = 4; mf1 = 0; f0 = 3; mf0 = 0;
OmR, AC, Psc = cs.groundStateRamanTransition(Pa, wa, qa,
                                             Pb, wb, qb,
                                             Delta0,
                                             f0, mf0, f1, mf1,
                                             ne, le, je)
print('\nRaman Transitions |%d,%d> to |%d,%d> via %dP_%d/2'
      % (f1, mf1, f0, mf0, ne, je*2))
print('\tParameters:\tDelta/2pi = %2.f GHz, Pa = %1.2f mW, wa = %2.1f um, Pa = %1.2f mW, wa = %2.1f um' \
      % (Delta0 / 2.0 / np.pi * 1e-9, Pa * 1e3, wa * 1e6, Pb * 1e3, wb * 1e6))
print('\tResults:\tOmR/2pi = %2.5f MHz, AC/2pi = %2.5f MHz, Psc = %f'
      % (OmR / 2.0 / np.pi * 1e-6, AC / 2.0 / np.pi * 1e-6, Psc))

#D2 Line
ne = 6; le = 1; je = 1.5;
OmR, AC, Psc = cs.groundStateRamanTransition(Pa, wa, qa,
                                             Pb, wb, qb,
                                             Delta0,
                                             f0, mf0,
                                             f1, mf1,
                                             ne,le,je)
print('\nRaman Transitions |%d,%d> to |%d,%d> via %dP_%d/2' % (f1, mf1, f0, mf0, ne, je*2))
print('\tParameters:\tDelta/2pi = %2.f GHz, Pa = %1.2f mW, wa = %2.1f um, Pa = %1.2f mW, wa = %2.1f um' \
      % (Delta0 / 2.0 / np.pi * 1e-9, Pa * 1e3, wa * 1e6, Pb * 1e3, wb * 1e6))
print('\tResults:\tOmR/2pi = %2.5f MHz, AC/2pi = %2.5f MHz, Psc = %f'
      % (OmR / 2.0 / np.pi * 1e-6, AC / 2.0 / np.pi * 1e-6, Psc))

#7P1/2
ne = 7; le = 1; je = 0.5; Pa = 10e-3; Pb = 10e-3;
OmR, AC, Psc = cs.groundStateRamanTransition(Pa, wa, qa,
                                             Pb, wb, qb,
                                             Delta0,
                                             f0, mf0,
                                             f1, mf1,
                                             ne, le, je)
print('\nRaman Transitions |%d,%d> to |%d,%d> via %dP_%d/2'
       % (f1, mf1, f0, mf0, ne, je*2))
print('\tParameters:\tDelta/2pi = %2.f GHz, Pa = %1.2f mW, wa = %2.1f um, Pa = %1.2f mW, wa = %2.1f um' \
      % (Delta0 / 2.0 / np.pi * 1e-9, Pa * 1e3, wa * 1e6, Pb * 1e3, wb * 1e6)
     )
print('\tResults:\tOmR/2pi = %2.5f MHz, AC/2pi = %2.5f MHz, Psc = %f'
      % (OmR / 2.0 / np.pi * 1e-6, AC / 2.0 / np.pi * 1e-6, Psc))

#7P3/2
ne = 7; le = 1; je = 1.5; Pa = 10e-3; Pb = 10e-3;
OmR, AC, Psc = cs.groundStateRamanTransition(Pa, wa, qa,
                                             Pb, wb, qb,
                                             Delta0,
                                             f0, mf0, f1,
                                             mf1, ne, le, je)
print('\nRaman Transitions |%d,%d> to |%d,%d> via %dP_%d/2'
      % (f1, mf1, f0, mf0, ne, je * 2))
print('\tParameters:\tDelta/2pi = %2.f GHz, Pa = %1.2f mW, wa = %2.1f um, Pa = %1.2f mW, wa = %2.1f um' \
      % (Delta0 / 2.0 / np.pi * 1e-9, Pa * 1e3, wa * 1e6, Pb * 1e3, wb * 1e6))
print('\tResults:\tOmR/2pi = %2.5f MHz, AC/2pi = %2.5f MHz, Psc = %f'
      % (OmR / 2.0 / np.pi * 1e-6, AC / 2.0 / np.pi * 1e-6, Psc))

#Prepare Output Figure
f, ax = plt.subplots(1, 3, figsize=(12, 5))
ax[0].set_xlabel(r'$\Delta/2\pi$ (GHz)')
ax[0].set_ylabel(r'$\Omega/2\pi$ (MHz)')
ax[0].set_ylim([-20, 20])
ax[1].set_xlabel(r'$\Delta/2\pi$ (GHz)')
ax[1].set_ylabel(r'$\Delta_\mathrm{AC}/2\pi$ (MHz)')
ax[1].set_ylim([-20, 20])
ax[2].set_xlabel(r'$\Delta/2\pi$ (GHz)')
ax[2].set_ylabel(r'$P_\mathrm{sc}$')
ax[2].set_ylim([0, 0.01])

#Reset to D2 Line
ne = 6; le = 1; je = 1.5; f1 = 4; mf1 = 0; f0 = 3; mf0 = 0;
Pa = 0.1e-3; wa = 5e-6; qa = -1;
Pb = 0.1e-3; wb = 5e-6; qb = -1;

for mf in range(-3,4):
    OmR, AC, Psc = cs.groundStateRamanTransition(Pa, wa, qa,
                                                 Pb, wb, qb,
                                                 Delta,
                                                 f1, mf, f0, mf + qa - qb,
                                                 ne, le, je)
    ax[0].plot(Delta / (2e9 * np.pi), OmR / (2e6 * np.pi))
    ax[1].plot(Delta / (2e9 * np.pi), AC / (2e6 * np.pi))
    ax[2].plot(Delta / (2e9 * np.pi), Psc,
               label=format("$\\vert%d,%d\\rangle\\rightarrow\\vert%d,%d\\rangle$"
                            % (f1,mf,f0,mf)))
ax[2].legend(loc=4)
ax[1].set_title('D2 Raman Transitions')
plt.show()
```

## Two-photon Rydberg excitation

We consider here the case of two-photon ground-state Raman transitions from $\vert g \rangle = \vert f_g,m_{f_g}\rangle$ to $\vert r \rangle = \vert n_r,l_r,j_r,m_{j_r}\rangle$ via an intermediate excited state $\vert e\rangle = \vert f_e,m_{f_e}\rangle$. In the rotating wave approximation, the Hamiltonian can be written as

\begin{equation}
\mathcal{H} = \hbar\begin{pmatrix} 0 & \Omega_p/2 & 0 \\\Omega_p/2 & -\Delta_p & \Omega_c/2 \\ 0 & \Omega_c/2 & -\delta\end{pmatrix},
\end{equation}

where $\Omega_p$ is the Rabi frequency from $\vert 0\rangle\rightarrow\vert e\rangle$ due to the probe laser and $\Omega_c$ is the Rabi frequency from $\vert e\rangle\rightarrow\vert r\rangle$ due to the coupling laser, $\delta=(\Delta_p+\Delta_c)$ is the two-photon detuning.

In the limit $\Delta_p \gg \Omega_p,\Omega_c,\Gamma_e$ the intermediate excited state can be adiabatically elliminiated leading to an effective two-photon Rabi frequency $\Omega_R = \Omega_p\Omega_c/2\Delta$ and an effective detuning
$\delta_\mathrm{eff}=\delta+\Delta_\mathrm{AC}$ where $\Delta_\mathrm{AC}$ is the differential Stark shift given by

\begin{equation}
\Delta_\mathrm{AC} = \frac{\vert\Omega_p\vert^2}{4\Delta_p} - \frac{\vert\Omega_c\vert^2}{4\Delta_c}.
\end{equation}

On the two-photon transition ($\delta_\mathrm{eff}=0$), the probability of spontanteous emission during a $\pi$-pulse can be evaluated from $P_\mathrm{sc}=\Gamma_e\rho_e\tau_\pi$, where $\tau_\pi=\pi/\Omega_R$ is the $\pi$-pulse duration and $\rho_e$ is the excited state population given by

\begin{equation}
\rho_e = \frac{1}{2}\left(\frac{\vert\Omega_p\vert^2}{2\Delta_p^2}+\frac{\vert\Omega_c\vert^2}{2\Delta_c^2}\right).
\end{equation}

For a multi-level atom, the terms describing Raman transition, diffrential AC Stark shift and scattering probability must be calculated by summing over the $\vert f_e,m_{f_e}\rangle$ states of the intermediate excited state. Additionally, for transitions between hyperfine ground states separated by frequency $\omega_{01}$ we must also account for the off-resonant AC Stark shift and scattering terms arising from laser A acting on $\vert 1 \rangle$ with detuning $\Delta'=\Delta+\omega_{01}$ and laser B acting on $\vert 0 \rangle$ with detuning $\Delta'=\Delta-\omega_{01}$.

\begin{align}
\Omega_R&=\displaystyle\sum_{f_e,m_{f_e}}\frac{\Omega_p^{g\rightarrow f_e}\Omega_c^{f_e\rightarrow r}}{2(\Delta-\Delta_{f_e})}\\
\Delta_{\mathrm{AC}_g} &= \displaystyle\sum_{f_e,m_{f_e}}\frac{\vert\Omega_p^{g\rightarrow f_e}\vert^2}{4(\Delta-\Delta_{f_e})}\\
\Delta_{\mathrm{AC}_r} &= \displaystyle\sum_{f_e,m_{f_e}}\frac{\vert\Omega_c^{f_e\rightarrow r}\vert^2}{4(\Delta-\Delta_{f_e})}\\
P_\mathrm{sc} &= \frac{\Gamma_et_\pi}{2}\displaystyle\sum_{f_e,m_{f_e}}\left[\frac{\vert\Omega_p^{g\rightarrow f_e}\vert^2}{2(\Delta-\Delta_{f_e})^2}+\frac{\vert\Omega_c^{f_e\rightarrow r}\vert^2}{2(\Delta-\Delta_{f_e})^2}\right]
\end{align}

 + `OmegaR, ACg, ACr, Psc = twoPhotonRydbergExcitation(Pp, wp, qp, Pc, wc, qc, Delta, fg, mfg, ne, le, je, nr, lr, jr, mjr)`

Returns two-photon Rabi frequency $\Omega_R$, ground-state AC Starkshift $\Delta_\mathrm{AC_g}$, Rydberg-state AC Starkshift $\Delta_\mathrm{AC_r}$ and scattering probability $P_\mathrm{sc}$ for transition from $\vert f_g,m_{f_g}\rangle\rightarrow n_rL_{j_r}\vert j_r,m_{j_r}\rangle$ via intermediate excited state with quantum numbers $n_e,\ell_e,j_e$ with detuning $\Delta$ from excited state centre of mass driven by probe/coupling lasers with power $P$, waist $w$ and polarisation $q$ driving transitions from $\vert g\rangle,\vert r\rangle$ respectively.

```python
#Rydberg Excitation vs 6P3/2
Pp=500e-9;wp=3e-6;qp=1;
Pc=200e-3;wc=5e-6;qc=-1;
#Ground State
fg=4;mfg=0;
#Excited state (6P3/2)
ne=6;le=1;je=1.5;
#Rydberg state (82S1/2)
nr=80;lr=0;jr=0.5;mjr=-0.5;

SPD='SPD'

#Detuning Array
Delta0=1.1e9*2*np.pi
Delta=np.linspace(-2.5,2.5,1001)*1e9*2*np.pi

[OmR,ACg,ACr,Psc]=cs.twoPhotonRydbergExcitation(Pp,wp,qp,Pc,wc,qc,Delta0,fg,mfg,ne,le,je,nr,lr,jr,mjr)
print('\nRydberg Excitation 6S1/2 |%d,%d> to %d%c |%d/2,%d/2> via %dP_%d/2' % (fg,mfg,nr,SPD[lr],2*jr,2*mjr,ne,je*2))
print('\n\tParameters:\tDelta/2pi = %2.2f GHz, Pp = %1.2f uW, wp = %2.1f um, qp=%d, Pc = %1.2f mW, wc = %2.1f um, qc=%d' \
      % (Delta0/2.0/np.pi*1e-9,Pp*1e6,wp*1e6,qp,Pc*1e3,wc*1e6,qc))
print('\n\tResults:\tOmR/2pi = %2.5f MHz\n\t\t\tACg/2pi = %2.5f MHz\n\t\t\tACr/2pi = %2.5f MHz\n\t\t\tDelta_AC/2pi = %2.5f MHz\n\t\t\tPsc = %f' \
      % (OmR/2.0/np.pi*1e-6,ACg/2.0/np.pi*1e-6,ACr/2.0/np.pi*1e-6,(ACg-ACr)/2.0/np.pi*1e-6,Psc))

#Prepare Output Figure
[f,ax]=plt.subplots(1,3,figsize=(16, 5))
ax[0].set_xlabel(r'$\Delta/2\pi$ (GHz)')
ax[0].set_ylabel(r'$\Omega/2\pi$ (MHz)')
ax[0].set_ylim([-10,10])
ax[1].set_xlabel(r'$\Delta/2\pi$ (GHz)')
ax[1].set_ylabel(r'$\Delta_\mathrm{AC}/2\pi$ (MHz)')
ax[1].set_ylim([-20,20])
ax[2].set_xlabel(r'$\Delta/2\pi$ (GHz)')
ax[2].set_ylabel(r'$P_\mathrm{sc}$')
ax[2].set_ylim([0,0.25])
ax[0].plot(Delta/(2e9*np.pi),0*Delta,'k--')
ax[1].plot(Delta/(2e9*np.pi),0*Delta,'k--')

mj=-jr-1
for m in range(int(2*jr+1)):
    mj+=1
    for mfg in [-fg,0,fg]:
        [OmR,ACg,ACr,Psc]=cs.twoPhotonRydbergExcitation(Pp,wp,qp,Pc,wc,qc,Delta,fg,mfg,ne,le,je,nr,lr,jr,mj)
        if(np.sum(np.abs(OmR))>0.1):
            ax[0].plot(Delta/(2e9*np.pi),OmR/(2e6*np.pi))
            ax[1].plot(Delta/(2e9*np.pi),(ACg-ACr)/(2e6*np.pi))
            ax[2].plot(Delta/(2e9*np.pi),Psc,label=format("$\\vert%d,%d\\rangle\\rightarrow%d%c\\vert%1.0f/2,%1.0f/2\\rangle$" \
                                                      % (fg,mfg,nr,SPD[lr],2*jr,2*mj)))
ax[2].legend(loc=1)
ax[1].set_title('Rydberg excitation 6S_{1/2} -> %dP_{%1.f/2} -> %d%c_{%1.0f/2}'%(ne,2*je,nr,SPD[lr],2*jr))

plt.show()
```

```python
#Rydberg Excitation vs 7P1/2
Pp=1.535e-3;wp=5e-6;qp=1;
Pc=20e-3;wc=3e-6;qc=-1;
#Ground State
fg=4;mfg=0;
#Excited state (7P1/2)
ne=7;le=1;je=0.5;
#Rydberg state (82S1/2)
nr=80;lr=0;jr=0.5;mjr=-0.5;

#Detuning Array
Delta0=1e9*2*np.pi
Delta=np.linspace(-2.5,2.5,1001)*1e9*2*np.pi

[OmR,ACg,ACr,Psc]=cs.twoPhotonRydbergExcitation(Pp,wp,qp,Pc,wc,qc,Delta0,fg,mfg,ne,le,je,nr,lr,jr,mjr)
print('\nRydberg Excitation 6S1/2 |%d,%d> to %d%c |%d/2,%d/2> via %dP_%d/2' % (fg,mfg,nr,SPD[lr],2*jr,2*mjr,ne,je*2))
print('\n\tParameters:\tDelta/2pi = %2.2f GHz, Pp = %1.2f uW, wp = %2.1f um, qp=%d, Pc = %1.2f mW, wc = %2.1f um, qc=%d' \
      % (Delta0/2.0/np.pi*1e-9,Pp*1e6,wp*1e6,qp,Pc*1e3,wc*1e6,qc))
print('\n\tResults:\tOmR/2pi = %2.5f MHz\n\t\t\tACg/2pi = %2.5f MHz\n\t\t\tACr/2pi = %2.5f MHz\n\t\t\tDelta_AC/2pi = %2.5f MHz\n\t\t\tPsc = %f' \
      % (OmR/2.0/np.pi*1e-6,ACg/2.0/np.pi*1e-6,ACr/2.0/np.pi*1e-6,(ACg-ACr)/2.0/np.pi*1e-6,Psc))

#Prepare Output Figure
[f,ax]=plt.subplots(1,3,figsize=(16, 5))
ax[0].set_xlabel(r'$\Delta/2\pi$ (GHz)')
ax[0].set_ylabel(r'$\Omega/2\pi$ (MHz)')
ax[0].set_ylim([-10,10])
ax[1].set_xlabel(r'$\Delta/2\pi$ (GHz)')
ax[1].set_ylabel(r'$\Delta_\mathrm{AC}/2\pi$ (MHz)')
ax[1].set_ylim([-20,20])
ax[2].set_xlabel(r'$\Delta/2\pi$ (GHz)')
ax[2].set_ylabel(r'$P_\mathrm{sc}$')
ax[2].set_ylim([0,0.25])
ax[0].plot(Delta/(2e9*np.pi),0*Delta,'k--')
ax[1].plot(Delta/(2e9*np.pi),0*Delta,'k--')

SPD='SPD'

mj=-jr-1
for m in range(int(2*jr+1)):
    mj+=1
    for mfg in [-fg,0,fg]:
        [OmR,ACg,ACr,Psc]=cs.twoPhotonRydbergExcitation(Pp,wp,qp,Pc,wc,qc,Delta,fg,mfg,ne,le,je,nr,lr,jr,mj)
        if(np.sum(np.abs(OmR))>0.1):
            ax[0].plot(Delta/(2e9*np.pi),OmR/(2e6*np.pi))
            ax[1].plot(Delta/(2e9*np.pi),(ACg-ACr)/(2e6*np.pi))
            ax[2].plot(Delta/(2e9*np.pi),Psc,label=format("$\\vert%d,%d\\rangle\\rightarrow%d%c\\vert%1.0f/2,%1.0f/2\\rangle$" \
                                                      % (fg,mfg,nr,SPD[lr],2*jr,2*mj)))
ax[2].legend(loc=1)
ax[1].set_title('Rydberg excitation 6S_{1/2} -> %dP_{%1.f/2} -> %d%c_{%1.0f/2}'%(ne,2*je,nr,SPD[lr],2*jr))

plt.show()
```

## Magnetic Functions

 + `getLandegj(l, j, s=0.5)`

    Returns fine-structure Landé g-factor $g_J$

\begin{equation}
g_J\simeq 1+\frac{j(j+1)+s(s+1)-l(l+1)}{2j(j+1)}
\end{equation}

 + `getLandegjExact(l, j, s=0.5)`

    Returns fine-structure Landé g-factor $g_J$ using precise $g_L,$ $g_S$ values

\begin{equation}
g_J=g_L\frac{j(j+1)-s(s+1)+l(l+1)}{2j(j+1)}+g_S\frac{j(j+1)+s(s+1)-l(l+1)}{2j(j+1)}
\end{equation}

 + `getLandegf(l, j, f, s=0.5)`

    Returns hyperfine-structure Landé g-factor $g_F$

\begin{equation}
g_F\simeq g_J\frac{f(f+1)-I(I+1)+j(j+1)}{2f(f+1)}
\end{equation}

 + `getLandegfExact(l, j, f, s=0.5)`

    Returns hyperfine-structure Landé g-factor $g_F$ using precise $g_L,$ $g_S$, $g_I$ values

\begin{equation}
g_F=g_J\frac{f(f+1)-I(I+1)+j(j+1)}{2f(f+1)}+g_I\frac{f(f+1)+I(I+1)-j(j+1)}{2f(f+1)}
\end{equation}

 + `Ez = breitRabi(l, j, B, Ahfs, Bhfs)`

Return exact Zeman energy levels for the l,j manifold via diagonalisation of $\mathcal{H}=\mathcal{H}_\mathrm{hfs}+\mathcal{H}_\mathrm{Z}$ where $\mathcal{H}_\mathrm{hfs}$ is the Hamiltonian for the hyperfine interaction given by
\begin{equation}
\mathcal{H}_\mathrm{hfs}=A_\mathrm{hfs}I\cdot J + B_\mathrm{hfs}\frac{3(I\cdot J)^2+3/2 I\cdot J -I^2J^2}{2I(2I+1)2J(2J+1)},
\end{equation}
and $\mathcal{H}_Z$ is the Hamiltonian for the Zeeman interaction given by
\begin{equation}
\mathcal{H}_Z=\frac{\mu_B}{\hbar}(g_J J_z+g_I I_z)B_z
\end{equation}

For the $j=1/2$ ground-state, this can be compared against the breitRabi formula [1]

\begin{equation}
    E_{\vert j=1/2,m_j,I,m_I\rangle} = -\frac{A_\mathrm{hfs}(I+1/2)}{2(2I+1)}+g_I\mu_B m B_z \pm \frac{A_\mathrm{hfs}(I+1/2)}{2}\left(1+\frac{4mx}{2I+1}+x^2\right)^{1/2},
\end{equation}
where $m=m_I\pm m_j$ and
\begin{equation}
x=\frac{g_j-g_I)\mu_B B_z}{A_\mathrm{hfs}(I+1/2)}.
\end{equation}

```python
#Lande g-factors
gj = cs.getLandegj(0, 0.5)
gf3 = cs.getLandegf(0, 0.5, 4)
gf4 = cs.getLandegf(0, 0.5, 4)
print("Cs 6S1/2 gJ = %2.2f, F=3 gF = %2.2f, F=4 gF = %2.2f,"
      % (gj, gf3, gf4))
gj = cs.getLandegjExact(0, 0.5)
gf3 = cs.getLandegfExact(0, 0.5, 4)
gf4 = cs.getLandegfExact(0, 0.5, 4)
print("Cs 6S1/2 gJ = %f, F=3 gF = %f, F=4 gF = %f,"
      % (gj, gf3, gf4))

#Zeeman Interaction 6S1/2
#========================
#Define State of Interest
lg = 0; jg = 1 / 2;

#B-Field (G)
B = np.linspace(0, 5000, 101) * 1e-4
#Exact-Diagonalisation (Breit-Rabi)
en, f, mf = cs.breitRabi(6, lg, jg, B)

f, ax = plt.subplots(3, 1, figsize=(12, 10))
ax[0].plot(B * 1e4, en * 1e-9, 'k')
ax[0].set_xlabel(r'$B_z$ (G)')
ax[0].set_ylabel(r'$E$ (GHz)')
ax[0].set_title(r'Breit-Rabi Diagram $6S_{1/2}$')

#Plot Breit-Rabi Formula
#=======================
#   Steck Eq(25)
#                         dEhfs                dEhfs      4mx
#   E(j=1/2,mj;I,mI) = - ------- + gI*uB*m*B ± ----- (1+ ----- + x^2)^{1/2}
#                        2(2I+1)                 2        2I+1
#
#   where m=mI±mj to match sign above and x=(gj-gI)uB*B/dEhfs
uB = physical_constants["Bohr magneton in Hz/T"][0]
Ahfs, Bhfs = cs.getHFSCoefficients(6,0,0.5)
qhfs = Ahfs * (cs.I + jg)
gj = cs.getLandegjExact(lg, jg)
#Rescale to x
x=(gj - cs.gI) * uB * B / qhfs;

Ebr = np.zeros(np.shape(en))

ctr=-1
for mI in np.arange(-cs.I, 1+cs.I):
    for mj in np.arange(-jg, 1+jg):
        m=mI+mj
        ctr+=1
        if(abs(m)==(cs.I + jg)):
            Ebr[:,ctr] = qhfs * cs.I/(2 * cs.I + 1) + np.sign(mj)*1/2*(gj+2*cs.gI)*uB*B
        else:
            Ebr[:,ctr] = -qhfs/(2 * (2 * cs.I + 1)) + cs.gI * uB * m * B \
                         + np.sign(mj) * qhfs / 2 * np.sqrt(1 + 4 * m * x / (2 * cs.I + 1) + x**2)

ax[0].plot(B * 1e4, Ebr * 1e-9, 'r.');

ax[0].text(0,min(en[0,:])*1.5e-9,r"$F=3$",horizontalalignment='center',fontsize=12)
ax[0].text(0,max(en[0,:])*1.2e-9,r"$F=4$",horizontalalignment='center',fontsize=12)

#6P1/2

#Exact-Diagonalisation (Breit-Rabi)
B=np.linspace(0,2000,101)*1e-4
en,f,mf=cs.breitRabi(6,1,0.5,B)
ax[1].plot(B*1e4,en*1e-9,'k')
ax[1].set_xlabel(r'$B_z$ (G)')
ax[1].set_ylabel(r'$E$ (GHz)')
ax[1].set_title(r'Breit-Rabi Diagram $6P_{1/2}$')
ax[1].text(0, min(en[0, :]) * 1.5e-9,
           r"$F=3$",
           horizontalalignment='center', fontsize=12)
ax[1].text(0, max(en[0, :]) * 1.2e-9,
           r"$F=4$",
           horizontalalignment='center',
           fontsize=12)

#Exact-Diagonalisation (Breit-Rabi)
B = np.linspace(0, 500, 101) * 1e-4
en, f, mf = cs.breitRabi(6, 1, 1.5, B)
ax[2].plot(B * 1e4, en * 1e-9, 'k')
ax[2].set_xlabel(r'$B_z$ (G)')
ax[2].set_ylabel(r'$E$ (GHz)')
ax[2].set_title(r'Breit-Rabi Diagram $6P_{3/2}$')
ax[2].text(0, min(en[0,:]) * 2e-9,
           r"$F=2$",
           horizontalalignment='center', fontsize=12)
ax[2].text(0, max(en[0,:]) * 1.7e-9,
           r"$F=5$",
           horizontalalignment='center', fontsize=12)

plt.show()
```

```python

```