```python
import tensorflow as tf
import numpy as np
import scipy.io
import timeit
import sys
import matplotlib.pyplot as plt
from scipy.interpolate import griddata
import random
import seaborn as sns
import pandas as pd
import mpmath
```

The analytical solution (Srivastava and Yeh, 1991) is based on the following constituitive relationships:
\begin{equation}
K_* = K_s e^{\alpha \psi},
\end{equation}

\begin{equation}
\theta = \theta_r + (\theta_s - \theta_r) e^{\alpha \psi},
\end{equation}
where $K_s$ is the saturated hydraulic conductivity, $\theta_r$ is the residual water content, $\theta_s$ is the saturated water content, $\alpha$ is a pore-size distribution parameter. Using the constituitive relationships, the RRE is linearlized.

Dimensionless parameters are introduced:
\begin{equation}
z = \alpha_1 z_* \text{   for   } -L_{*1} \leq z_* \leq 0,
\end{equation}

\begin{equation}
z = \alpha_2 z_* \text{   for   } 0 \leq z_* \leq L_{*2},
\end{equation}

\begin{equation}
K_1 = K_{*1}/K_{s1},  q_{A1} = q_{A*}/K_{s1},  q_{B} = q_{B*}/K_{s1},
\end{equation}

\begin{equation}
K_2 = K_{*2}/K_{s2},  q_{A2} = q_{A*}/K_{s2},  q_{B} = q_{B*}/K_{s2},
\end{equation}

\begin{equation}
t = \frac{\alpha_1 K_{s1}t_*}{\theta_{s1} - \theta_{r1}}.
\end{equation}

The governing equations are:
\begin{equation}
\frac{\partial^2 K_1}{\partial z^2} + \frac{\partial K_1}{\partial z} = \frac{\partial K_1}{\partial t},
\end{equation}

\begin{equation}
\frac{\partial^2 K_2}{\partial z^2} + \frac{\partial K_2}{\partial z} = \beta\frac{\partial K_2}{\partial t},
\end{equation}
where
\begin{equation}
\beta = \frac{\alpha_1 K_{s1}(\theta_{s2} - \theta_{r2})}{\alpha_2 K_{s2}(\theta_{s1} - \theta_{r1})}.
\end{equation}

```python
# define computational domain
dom_coords_1 = np.array([[0.0, 0.0],   # t = 0 hour to t = 10 hours
                       [10.0, -10.0]])  # domain 1 is the lower layer

dom_coords_2 = np.array([[0.0, 10.0],   # t = 0 hour to t = 10 hours
                       [10.0, 0.0]])   # domain 2 is the upper layer
```

```python
t_star = np.linspace(dom_coords_1[0, 0], dom_coords_1[1, 0], 101)[:, None]
z1_star = np.linspace(dom_coords_1[0, 1], dom_coords_1[1, 1], 101)[:, None]
z2_star = np.linspace(dom_coords_2[0, 1], dom_coords_2[1, 1], 101)[:, None]
```

```python
# parameter values
alpha1 = 1.0 # /cm
alpha2 = 1.0 # /cm
K_s1 = 1.0 # cm/hour
K_s2 = 10.0 # cm/hour
theta_s1 = 0.40
theta_s2 = 0.40
theta_r1 = 0.06
theta_r2 = 0.06
L1_star = 10.0 # cm
L2_star = 10.0 # cm
```

```python
# transformation
L1 = L1_star * alpha1
L2 = L2_star * alpha2
z1 = z1_star * alpha1
z2 = z2_star * alpha2

t = alpha1 * K_s1 * t_star/(theta_s1 - theta_r1)
```

```python
# suface water flux condition (positive downward) [cm/hour]
q_A1_star = 0.1
q_A2_star = 0.1
q_B1_star = 0.9
q_B2_star = 0.9

# pressure at the lower boundary (water table)
psi_0 = 0.0

q_A1 = q_A1_star / K_s1
q_A2 = q_A2_star / K_s2
q_B1 = q_B1_star / K_s1
q_B2 = q_B2_star / K_s2

# initial condition for the lower layer
K_1_ini = q_A1 - (q_A1 - np.exp(alpha1 * psi_0)) * np.exp(-(L1 + z1))
psi_1_ini = np.log(K_1_ini) /alpha1
```

```python
t1_mesh, z1_mesh = np.meshgrid(t, z1)
t2_mesh, z2_mesh = np.meshgrid(t, z2)
```

```python
# initial condition for the upper layer
K_2_ini = q_A2 - (q_A2 - (q_A1 - (q_A1 - np.exp(alpha1 * psi_0)) \
       *np.exp(-L1))**(alpha2/alpha1))* np.exp(-z2)
psi_2_ini = np.log(K_2_ini) /alpha2
```

```python
# initial condition
plt.plot(psi_1_ini, z1_star, label = "layer 1")
plt.plot(psi_2_ini, z2_star, label = "layer 2")
plt.legend()
plt.xlabel("$\psi$ cm")
plt.ylabel("$z$ cm")
```

```python
# beta = 0.1 for this example, which means Case A and Case C are used for the analytical solution
beta = (alpha1 * K_s1 * (theta_s2 - theta_r2))/ (alpha2 * K_s2 * (theta_s1 - theta_r1))
print("beta is", beta)
```

    beta is 0.09999999999999999

Computing $RA_1$.
The values of $\lambda$ are obtained as the positive roots of the characteristic equation (21a).
we obtain $\lambda$ < 100 for higher accuracies for $t = 0.1$.

```python
lam_test = np.arange(0.0, 100, 0.001)
```

```python
def residue_RA(lam):
    mu = (beta * lam**2 + (beta - 1.0)/4.0)**(1/2)
    A = np.sin(lam*L1) + 2*lam*np.cos(lam*L1)
    B = np.sin(mu*L2) + 2*mu*np.cos(mu*L2)
    C = (1 + 4*mu**2)*(K_s2/K_s1)*np.sin(lam*L1)*np.sin(mu*L2)
    return A * B - C
```

```python
# two sign changes: one is due to discontinuity, and the other one is the root
pd.set_option('display.max_rows', None)
data = {"lam": lam_test.flatten()[0:-1],
        "sign": np.sign(residue_RA(lam_test)).flatten()[0:-1],
        "change": np.sign(residue_RA(lam_test)).flatten()[0:-1] - np.sign(residue_RA(lam_test)).flatten()[1:]}
df = pd.DataFrame(data)
```

    C:\Users\toshi\anaconda3\envs\tensorflow1\lib\site-packages\ipykernel_launcher.py:2: RuntimeWarning: invalid value encountered in sqrt

    C:\Users\toshi\anaconda3\envs\tensorflow1\lib\site-packages\ipykernel_launcher.py:4: RuntimeWarning: invalid value encountered in sign
      after removing the cwd from sys.path.
    C:\Users\toshi\anaconda3\envs\tensorflow1\lib\site-packages\ipykernel_launcher.py:5: RuntimeWarning: invalid value encountered in sign
      """

```python
# the roots can be found near where the sign changes
a_set = df.loc[np.abs(df["change"]) == 2]["lam"].values
```

```python
# program to find roots using Newton method (solver = 'newton')

def residue_RA(lam):
    mu = (beta * lam**2 + (beta - 1.0)/4.0)**(1/2)
    A = mpmath.sin(lam*L1) + 2*lam*mpmath.cos(lam*L1)
    B = mpmath.sin(mu*L2) + 2*mu*mpmath.cos(mu*L2)
    C = (1 + 4*mu**2)*(K_s2/K_s1)*mpmath.sin(lam*L1)*mpmath.sin(mu*L2)
    return A * B - C

roots = []

for a in a_set[0:]:  # change the argument to find roots if you encounter errors
    a = int(a/0.001)
    for i in range(a, a + 5):
        print(i)
        initial_value = 0.001 * (i + 1)
        root = mpmath.findroot(residue_RA, initial_value, solver = 'newton')
        root_value = np.array(root, np.float).tolist()
        roots.append(root_value)
```

    1599
    1600
    1601
    1602
    1603
    1788
    1789
    1790
    1791
    1792
    1932
    1933
    1934
    1935
    1936
    2210
    2211
    2212
    2213
    2214
    2441
    2442
    2443
    2444
    2445
    2588
    2589
    2590
    2591
    2592
    2849
    2850
    2851
    2852
    2853
    3128
    3129
    3130
    3131
    3132
    3326
    3327
    3328
    3329
    3330
    3508
    3509
    3510
    3511
    3512
    3783
    3784
    3785
    3786
    3787
    4057
    4058
    4059
    4060
    4061
    4250
    4251
    4252
    4253
    4254
    4445
    4446
    4447
    4448
    4449
    4722
    4723
    4724
    4725
    4726
    4997
    4998
    4999
    5000
    5001
    5190
    5191
    5192
    5193
    5194
    5387
    5388
    5389
    5390
    5391
    5663
    5664
    5665
    5666
    5667
    5940
    5941
    5942
    5943
    5944
    6140
    6141
    6142
    6143
    6144
    6332
    6333
    6334
    6335
    6336
    6607
    6608
    6609
    6610
    6611
    6886
    6887
    6888
    6889
    6890
    7095
    7096
    7097
    7098
    7099
    7280
    7281
    7282
    7283
    7284
    7551
    7552
    7553
    7554
    7555
    7832
    7833
    7834
    7835
    7836
    8051
    8052
    8053
    8054
    8055
    8230
    8231
    8232
    8233
    8234
    8496
    8497
    8498
    8499
    8500
    8778
    8779
    8780
    8781
    8782
    9011
    9012
    9013
    9014
    9015
    9183
    9184
    9185
    9186
    9187
    9441
    9442
    9443
    9444
    9445
    9725
    9726
    9727
    9728
    9729
    9969
    9970
    9971
    9972
    9973
    10138
    10139
    10140
    10141
    10142
    10388
    10389
    10390
    10391
    10392
    10671
    10672
    10673
    10674
    10675
    10926
    10927
    10928
    10929
    10930
    11096
    11097
    11098
    11099
    11100
    11335
    11336
    11337
    11338
    11339
    11618
    11619
    11620
    11621
    11622
    11881
    11882
    11883
    11884
    11885
    12058
    12059
    12060
    12061
    12062
    12283
    12284
    12285
    12286
    12287
    12564
    12565
    12566
    12567
    12568
    12835
    12836
    12837
    12838
    12839
    13021
    13022
    13023
    13024
    13025
    13232
    13233
    13234
    13235
    13236
    13510
    13511
    13512
    13513
    13514
    13786
    13787
    13788
    13789
    13790
    13986
    13987
    13988
    13989
    13990
    14182
    14183
    14184
    14185
    14186
    14457
    14458
    14459
    14460
    14461
    14737
    14738
    14739
    14740
    14741
    14951
    14952
    14953
    14954
    14955
    15135
    15136
    15137
    15138
    15139
    15403
    15404
    15405
    15406
    15407
    15686
    15687
    15688
    15689
    15690
    15915
    15916
    15917
    15918
    15919
    16090
    16091
    16092
    16093
    16094
    16350
    16351
    16352
    16353
    16354
    16634
    16635
    16636
    16637
    16638
    16877
    16878
    16879
    16880
    16881
    17047
    17048
    17049
    17050
    17051
    17298
    17299
    17300
    17301
    17302
    17581
    17582
    17583
    17584
    17585
    17837
    17838
    17839
    17840
    17841
    18009
    18010
    18011
    18012
    18013
    18245
    18246
    18247
    18248
    18249
    18528
    18529
    18530
    18531
    18532
    18794
    18795
    18796
    18797
    18798
    18972
    18973
    18974
    18975
    18976
    19194
    19195
    19196
    19197
    19198
    19475
    19476
    19477
    19478
    19479
    19748
    19749
    19750
    19751
    19752
    19938
    19939
    19940
    19941
    19942
    20144
    20145
    20146
    20147
    20148
    20422
    20423
    20424
    20425
    20426
    20700
    20701
    20702
    20703
    20704
    20905
    20906
    20907
    20908
    20909
    21096
    21097
    21098
    21099
    21100
    21369
    21370
    21371
    21372
    21373
    21650
    21651
    21652
    21653
    21654
    21871
    21872
    21873
    21874
    21875
    22050
    22051
    22052
    22053
    22054
    22316
    22317
    22318
    22319
    22320
    22599
    22600
    22601
    22602
    22603
    22835
    22836
    22837
    22838
    22839
    23007
    23008
    23009
    23010
    23011
    23263
    23264
    23265
    23266
    23267
    23547
    23548
    23549
    23550
    23551
    23797
    23798
    23799
    23800
    23801
    23967
    23968
    23969
    23970
    23971
    24211
    24212
    24213
    24214
    24215
    24495
    24496
    24497
    24498
    24499
    24756
    24757
    24758
    24759
    24760
    24931
    24932
    24933
    24934
    24935
    25160
    25161
    25162
    25163
    25164
    25442
    25443
    25444
    25445
    25446
    25711
    25712
    25713
    25714
    25715
    25896
    25897
    25898
    25899
    25900
    26109
    26110
    26111
    26112
    26113
    26389
    26390
    26391
    26392
    26393
    26664
    26665
    26666
    26667
    26668
    26863
    26864
    26865
    26866
    26867
    27061
    27062
    27063
    27064
    27065
    27336
    27337
    27338
    27339
    27340
    27615
    27616
    27617
    27618
    27619
    27830
    27831
    27832
    27833
    27834
    28014
    28015
    28016
    28017
    28018
    28283
    28284
    28285
    28286
    28287
    28565
    28566
    28567
    28568
    28569
    28796
    28797
    28798
    28799
    28800
    28970
    28971
    28972
    28973
    28974
    29230
    29231
    29232
    29233
    29234
    29514
    29515
    29516
    29517
    29518
    29759
    29760
    29761
    29762
    29763
    29929
    29930
    29931
    29932
    29933
    30178
    30179
    30180
    30181
    30182
    30462
    30463
    30464
    30465
    30466
    30719
    30720
    30721
    30722
    30723
    30892
    30893
    30894
    30895
    30896
    31126
    31127
    31128
    31129
    31130
    31409
    31410
    31411
    31412
    31413
    31676
    31677
    31678
    31679
    31680
    31857
    31858
    31859
    31860
    31861
    32076
    32077
    32078
    32079
    32080
    32356
    32357
    32358
    32359
    32360
    32630
    32631
    32632
    32633
    32634
    32824
    32825
    32826
    32827
    32828
    33026
    33027
    33028
    33029
    33030
    33303
    33304
    33305
    33306
    33307
    33582
    33583
    33584
    33585
    33586
    33792
    33793
    33794
    33795
    33796
    33979
    33980
    33981
    33982
    33983
    34250
    34251
    34252
    34253
    34254
    34532
    34533
    34534
    34535
    34536
    34758
    34759
    34760
    34761
    34762
    34935
    34936
    34937
    34938
    34939
    35197
    35198
    35199
    35200
    35201
    35481
    35482
    35483
    35484
    35485
    35722
    35723
    35724
    35725
    35726
    35893
    35894
    35895
    35896
    35897
    36145
    36146
    36147
    36148
    36149
    36429
    36430
    36431
    36432
    36433
    36683
    36684
    36685
    36686
    36687
    36855
    36856
    36857
    36858
    36859
    37093
    37094
    37095
    37096
    37097
    37376
    37377
    37378
    37379
    37380
    37641
    37642
    37643
    37644
    37645
    37820
    37821
    37822
    37823
    37824
    38042
    38043
    38044
    38045
    38046
    38324
    38325
    38326
    38327
    38328
    38596
    38597
    38598
    38599
    38600
    38786
    38787
    38788
    38789
    38790
    38993
    38994
    38995
    38996
    38997
    39271
    39272
    39273
    39274
    39275
    39548
    39549
    39550
    39551
    39552
    39754
    39755
    39756
    39757
    39758
    39945
    39946
    39947
    39948
    39949
    40218
    40219
    40220
    40221
    40222
    40499
    40500
    40501
    40502
    40503
    40721
    40722
    40723
    40724
    40725
    40900
    40901
    40902
    40903
    40904
    41165
    41166
    41167
    41168
    41169
    41448
    41449
    41450
    41451
    41452
    41686
    41687
    41688
    41689
    41690
    41858
    41859
    41860
    41861
    41862
    42112
    42113
    42114
    42115
    42116
    42396
    42397
    42398
    42399
    42400
    42648
    42649
    42650
    42651
    42652
    42819
    42820
    42821
    42822
    42823
    43061
    43062
    43063
    43064
    43065
    43344
    43345
    43346
    43347
    43348
    43606
    43607
    43608
    43609
    43610
    43783
    43784
    43785
    43786
    43787
    44009
    44010
    44011
    44012
    44013
    44291
    44292
    44293
    44294
    44295
    44562
    44563
    44564
    44565
    44566
    44749
    44750
    44751
    44752
    44753
    44960
    44961
    44962
    44963
    44964
    45238
    45239
    45240
    45241
    45242
    45515
    45516
    45517
    45518
    45519
    45717
    45718
    45719
    45720
    45721
    45911
    45912
    45913
    45914
    45915
    46185
    46186
    46187
    46188
    46189
    46466
    46467
    46468
    46469
    46470
    46684
    46685
    46686
    46687
    46688
    46866
    46867
    46868
    46869
    46870
    47133
    47134
    47135
    47136
    47137
    47415
    47416
    47417
    47418
    47419
    47650
    47651
    47652
    47653
    47654
    47823
    47824
    47825
    47826
    47827
    48080
    48081
    48082
    48083
    48084
    48364
    48365
    48366
    48367
    48368
    48613
    48614
    48615
    48616
    48617
    48783
    48784
    48785
    48786
    48787
    49028
    49029
    49030
    49031
    49032
    49312
    49313
    49314
    49315
    49316
    49572
    49573
    49574
    49575
    49576
    49747
    49748
    49749
    49750
    49751
    49977
    49978
    49979
    49980
    49981
    50259
    50260
    50261
    50262
    50263
    50528
    50529
    50530
    50531
    50532
    50713
    50714
    50715
    50716
    50717
    50927
    50928
    50929
    50930
    50931
    51206
    51207
    51208
    51209
    51210
    51482
    51483
    51484
    51485
    51486
    51680
    51681
    51682
    51683
    51684
    51878
    51879
    51880
    51881
    51882
    52153
    52154
    52155
    52156
    52157
    52433
    52434
    52435
    52436
    52437
    52648
    52649
    52650
    52651
    52652
    52832
    52833
    52834
    52835
    52836
    53100
    53101
    53102
    53103
    53104
    53383
    53384
    53385
    53386
    53387
    53614
    53615
    53616
    53617
    53618
    53788
    53789
    53790
    53791
    53792
    54048
    54049
    54050
    54051
    54052
    54332
    54333
    54334
    54335
    54336
    54577
    54578
    54579
    54580
    54581
    54748
    54749
    54750
    54751
    54752
    54996
    54997
    54998
    54999
    55000
    55280
    55281
    55282
    55283
    55284
    55538
    55539
    55540
    55541
    55542
    55711
    55712
    55713
    55714
    55715
    55944
    55945
    55946
    55947
    55948
    56227
    56228
    56229
    56230
    56231
    56495
    56496
    56497
    56498
    56499
    56677
    56678
    56679
    56680
    56681
    56894
    56895
    56896
    56897
    56898
    57174
    57175
    57176
    57177
    57178
    57449
    57450
    57451
    57452
    57453
    57644
    57645
    57646
    57647
    57648
    57845
    57846
    57847
    57848
    57849
    58121
    58122
    58123
    58124
    58125
    58400
    58401
    58402
    58403
    58404
    58612
    58613
    58614
    58615
    58616
    58798
    58799
    58800
    58801
    58802
    59068
    59069
    59070
    59071
    59072
    59350
    59351
    59352
    59353
    59354
    59578
    59579
    59580
    59581
    59582
    59754
    59755
    59756
    59757
    59758
    60016
    60017
    60018
    60019
    60020
    60299
    60300
    60301
    60302
    60303
    60542
    60543
    60544
    60545
    60546
    60713
    60714
    60715
    60716
    60717
    60964
    60965
    60966
    60967
    60968
    61247
    61248
    61249
    61250
    61251
    61503
    61504
    61505
    61506
    61507
    61676
    61677
    61678
    61679
    61680
    61912
    61913
    61914
    61915
    61916
    62195
    62196
    62197
    62198
    62199
    62461
    62462
    62463
    62464
    62465
    62641
    62642
    62643
    62644
    62645
    62861
    62862
    62863
    62864
    62865
    63142
    63143
    63144
    63145
    63146
    63415
    63416
    63417
    63418
    63419
    63608
    63609
    63610
    63611
    63612
    63812
    63813
    63814
    63815
    63816
    64089
    64090
    64091
    64092
    64093
    64367
    64368
    64369
    64370
    64371
    64576
    64577
    64578
    64579
    64580
    64765
    64766
    64767
    64768
    64769
    65036
    65037
    65038
    65039
    65040
    65318
    65319
    65320
    65321
    65322
    65543
    65544
    65545
    65546
    65547
    65720
    65721
    65722
    65723
    65724
    65984
    65985
    65986
    65987
    65988
    66267
    66268
    66269
    66270
    66271
    66507
    66508
    66509
    66510
    66511
    66679
    66680
    66681
    66682
    66683
    66931
    66932
    66933
    66934
    66935
    67215
    67216
    67217
    67218
    67219
    67469
    67470
    67471
    67472
    67473
    67640
    67641
    67642
    67643
    67644
    67880
    67881
    67882
    67883
    67884
    68163
    68164
    68165
    68166
    68167
    68427
    68428
    68429
    68430
    68431
    68605
    68606
    68607
    68608
    68609
    68829
    68830
    68831
    68832
    68833
    69110
    69111
    69112
    69113
    69114
    69382
    69383
    69384
    69385
    69386
    69572
    69573
    69574
    69575
    69576
    69779
    69780
    69781
    69782
    69783
    70057
    70058
    70059
    70060
    70061
    70335
    70336
    70337
    70338
    70339
    70540
    70541
    70542
    70543
    70544
    70732
    70733
    70734
    70735
    70736
    71005
    71006
    71007
    71008
    71009
    71286
    71287
    71288
    71289
    71290
    71507
    71508
    71509
    71510
    71511
    71686
    71687
    71688
    71689
    71690
    71952
    71953
    71954
    71955
    71956
    72235
    72236
    72237
    72238
    72239
    72472
    72473
    72474
    72475
    72476
    72644
    72645
    72646
    72647
    72648
    72899
    72900
    72901
    72902
    72903
    73183
    73184
    73185
    73186
    73187
    73435
    73436
    73437
    73438
    73439
    73605
    73606
    73607
    73608
    73609
    73848
    73849
    73850
    73851
    73852
    74131
    74132
    74133
    74134
    74135
    74393
    74394
    74395
    74396
    74397
    74570
    74571
    74572
    74573
    74574
    74797
    74798
    74799
    74800
    74801
    75078
    75079
    75080
    75081
    75082
    75349
    75350
    75351
    75352
    75353
    75536
    75537
    75538
    75539
    75540
    75747
    75748
    75749
    75750
    75751
    76026
    76027
    76028
    76029
    76030
    76302
    76303
    76304
    76305
    76306
    76504
    76505
    76506
    76507
    76508
    76699
    76700
    76701
    76702
    76703
    76973
    76974
    76975
    76976
    76977
    77253
    77254
    77255
    77256
    77257
    77472
    77473
    77474
    77475
    77476
    77653
    77654
    77655
    77656
    77657
    77920
    77921
    77922
    77923
    77924
    78203
    78204
    78205
    78206
    78207
    78437
    78438
    78439
    78440
    78441
    78610
    78611
    78612
    78613
    78614
    78867
    78868
    78869
    78870
    78871
    79151
    79152
    79153
    79154
    79155
    79400
    79401
    79402
    79403
    79404
    79571
    79572
    79573
    79574
    79575
    79815
    79816
    79817
    79818
    79819
    80099
    80100
    80101
    80102
    80103
    80360
    80361
    80362
    80363
    80364
    80535
    80536
    80537
    80538
    80539
    80764
    80765
    80766
    80767
    80768
    81046
    81047
    81048
    81049
    81050
    81316
    81317
    81318
    81319
    81320
    81501
    81502
    81503
    81504
    81505
    81714
    81715
    81716
    81717
    81718
    81994
    81995
    81996
    81997
    81998
    82269
    82270
    82271
    82272
    82273
    82469
    82470
    82471
    82472
    82473
    82666
    82667
    82668
    82669
    82670
    82941
    82942
    82943
    82944
    82945
    83221
    83222
    83223
    83224
    83225
    83436
    83437
    83438
    83439
    83440
    83620
    83621
    83622
    83623
    83624
    83888
    83889
    83890
    83891
    83892
    84170
    84171
    84172
    84173
    84174
    84402
    84403
    84404
    84405
    84406
    84576
    84577
    84578
    84579
    84580
    84835
    84836
    84837
    84838
    84839
    85119
    85120
    85121
    85122
    85123
    85366
    85367
    85368
    85369
    85370
    85536
    85537
    85538
    85539
    85540
    85783
    85784
    85785
    85786
    85787
    86067
    86068
    86069
    86070
    86071
    86326
    86327
    86328
    86329
    86330
    86499
    86500
    86501
    86502
    86503
    86732
    86733
    86734
    86735
    86736
    87015
    87016
    87017
    87018
    87019
    87283
    87284
    87285
    87286
    87287
    87465
    87466
    87467
    87468
    87469
    87682
    87683
    87684
    87685
    87686
    87962
    87963
    87964
    87965
    87966
    88237
    88238
    88239
    88240
    88241
    88433
    88434
    88435
    88436
    88437
    88633
    88634
    88635
    88636
    88637
    88909
    88910
    88911
    88912
    88913
    89188
    89189
    89190
    89191
    89192
    89401
    89402
    89403
    89404
    89405
    89586
    89587
    89588
    89589
    89590
    89856
    89857
    89858
    89859
    89860
    90138
    90139
    90140
    90141
    90142
    90367
    90368
    90369
    90370
    90371
    90542
    90543
    90544
    90545
    90546
    90803
    90804
    90805
    90806
    90807
    91087
    91088
    91089
    91090
    91091
    91331
    91332
    91333
    91334
    91335
    91502
    91503
    91504
    91505
    91506
    91751
    91752
    91753
    91754
    91755
    92035
    92036
    92037
    92038
    92039
    92292
    92293
    92294
    92295
    92296
    92464
    92465
    92466
    92467
    92468
    92700
    92701
    92702
    92703
    92704
    92983
    92984
    92985
    92986
    92987
    93249
    93250
    93251
    93252
    93253
    93430
    93431
    93432
    93433
    93434
    93649
    93650
    93651
    93652
    93653
    93930
    93931
    93932
    93933
    93934
    94204
    94205
    94206
    94207
    94208
    94397
    94398
    94399
    94400
    94401
    94600
    94601
    94602
    94603
    94604
    94877
    94878
    94879
    94880
    94881
    95156
    95157
    95158
    95159
    95160
    95365
    95366
    95367
    95368
    95369
    95553
    95554
    95555
    95556
    95557
    95824
    95825
    95826
    95827
    95828
    96106
    96107
    96108
    96109
    96110
    96332
    96333
    96334
    96335
    96336
    96509
    96510
    96511
    96512
    96513
    96772
    96773
    96774
    96775
    96776
    97055
    97056
    97057
    97058
    97059
    97297
    97298
    97299
    97300
    97301
    97468
    97469
    97470
    97471
    97472
    97719
    97720
    97721
    97722
    97723
    98003
    98004
    98005
    98006
    98007
    98258
    98259
    98260
    98261
    98262
    98430
    98431
    98432
    98433
    98434
    98668
    98669
    98670
    98671
    98672
    98951
    98952
    98953
    98954
    98955
    99216
    99217
    99218
    99219
    99220
    99395
    99396
    99397
    99398
    99399
    99617
    99618
    99619
    99620
    99621
    99898
    99899
    99900
    99901
    99902

```python
set(roots)
```

    {1.5990891427798726,
     1.788732691327088,
     1.932914344549613,
     2.2109110956513467,
     2.4419555167035885,
     2.5885848998466656,
     2.849989333932118,
     3.1281660491170373,
     3.3263461862747836,
     3.5089830748945015,
     3.783396450018647,
     4.0587499327999375,
     4.2500096554362,
     4.445555964175939,
     4.722261597342127,
     4.997520330348105,
     5.190937232329985,
     5.3875852765189185,
     5.663899469930367,
     5.940709869938875,
     6.140596614835184,
     6.33275567156643,
     6.607129174713171,
     6.886060588985002,
     7.095209499312184,
     7.280401466688491,
     7.551376030390238,
     7.8323370399311925,
     8.052565867522132,
     8.23044450603,
     8.49636374147896,
     8.778910896574448,
     9.011039166201218,
     9.1830697482723,
     9.441983167272015,
     9.725494441337466,
     9.969322368562036,
     10.138530274048614,
     10.388233654920867,
     10.671982483467206,
     10.926434114796178,
     11.096952757076904,
     11.335200046561987,
     11.618363938254081,
     11.881795898316074,
     12.05816134199271,
     12.283050046645862,
     12.564675538325883,
     12.83524123901948,
     13.021599981161451,
     13.232043976029324,
     13.510979803646782,
     13.78691103995045,
     13.986383685492326,
     14.182548338021387,
     14.457356685210335,
     14.737099288791287,
     14.951422765861851,
     15.135035478146149,
     15.403903238899153,
     15.686131192623376,
     15.915568038428134,
     16.09003453791654,
     16.35073878569254,
     16.63430083057258,
     16.87778792031103,
     17.04799247320583,
     17.298014892620287,
     17.5818530423035,
     17.837380305311136,
     18.009054889969647,
     18.245930540189043,
     18.528987245914934,
     18.794120668749507,
     18.97288575406163,
     19.194752835079072,
     19.475868961286285,
     19.748223167244905,
     19.938664451417566,
     20.144841261296854,
     20.422642802130465,
     20.700146618147762,
     20.905243784142034,
     21.096665490613677,
     21.36944521522357,
     21.65039765714026,
     21.871345260135367,
     22.050789185647496,
     22.316417308981585,
     22.59942216949087,
     22.83574787378633,
     23.007771571246685,
     23.26371901522057,
     23.547575723690084,
     23.79750807013827,
     23.967956594247525,
     24.211546230341284,
     24.495132983115276,
     24.75617125425975,
     24.93122756072884,
     25.160152468257767,
     25.442308614647423,
     25.711820575470526,
     25.896908812003076,
     26.109874943073603,
     26.38927822657481,
     26.66490500841016,
     26.863889556692698,
     27.061159119294334,
     27.33619656983615,
     27.615993811847783,
     27.830846095936195,
     28.014560849029866,
     28.283213658512377,
     28.565612463776073,
     28.796451204805997,
     28.970681191280395,
     29.23049057513832,
     29.514184601227317,
     29.759593065890563,
     29.929984806172033,
     30.17821705395824,
     30.462034619984344,
     30.719614250741852,
     30.892534456195584,
     31.126632864304202,
     31.40941075106034,
     31.676435610144512,
     31.857815694621625,
     32.07605391991435,
     32.35650997438538,
     32.63043881838727,
     32.82479975070046,
     33.026899688846825,
     33.303499500645344,
     33.5822062485808,
     33.792166497675645,
     33.97970628209389,
     34.250534985713415,
     34.532310974256845,
     34.75853318936432,
     34.9350861979419,
     35.197777341476126,
     35.481226092576726,
     35.72267335845731,
     35.89357884497937,
     36.145410432352826,
     36.42931503958064,
     36.68376431788118,
     36.85538911756186,
     37.093661935009315,
     37.376853342894364,
     37.641566154759616,
     37.82015879203966,
     38.0428288886042,
     38.32405564700252,
     38.59636827681748,
     38.78696222224,
     38.993306217048676,
     39.2710995220295,
     39.54873872479654,
     39.75451084722061,
     39.945606643276214,
     40.218145296597775,
     40.499279680878274,
     40.72139638545084,
     40.90033907234792,
     41.16535366141217,
     41.448503469395185,
     41.68631256019343,
     41.85808912541147,
     42.1129034130371,
     42.39680718255687,
     42.648302000506234,
     42.819174289219866,
     43.06101176135671,
     43.344490768519194,
     43.6069728126568,
     43.7833820332885,
     44.00995911201846,
     44.291784989158856,
     44.56250882104887,
     44.74990063619678,
     44.960117802384424,
     45.23887702049473,
     45.51544495342708,
     45.71749157319881,
     45.91197633558676,
     46.185931679348194,
     46.46639939369378,
     46.68474077765222,
     46.866131814599285,
     47.13310975153132,
     47.415916915687504,
     47.65028593541643,
     47.82319709880706,
     48.080585800164926,
     48.364425182999994,
     48.61305905456402,
     48.783578851174276,
     49.0285679637308,
     49.31224720988233,
     49.57252672412086,
     49.74719827622644,
     49.977321910123365,
     50.259629311327394,
     50.52875921312127,
     50.71336333188308,
     50.92719928775875,
     51.20676812274482,
     51.48224356648535,
     51.680899329614626,
     51.87866465525323,
     52.15383310023448,
     52.4336031467763,
     52.648401016127465,
     52.83229855812179,
     53.10098556555707,
     53.38340988653587,
     53.61446737441787,
     53.78872776035599,
     54.04839662812041,
     54.33212001649836,
     54.57793918434023,
     54.748429047040915,
     54.99626654235868,
     55.28007903144342,
     55.53815235637299,
     55.711445050649495,
     55.94484783996595,
     56.22754875066923,
     56.4950587870766,
     56.67720610669772,
     56.8944732101477,
     57.17473544232577,
     57.449085073822914,
     57.644612558834176,
     57.84558402295245,
     58.121813522455895,
     58.4008498688213,
     58.612279882020026,
     58.79874145803496,
     59.06894533675702,
     59.350947530622044,
     59.578781823129134,
     59.7545764538522,
     60.01629925881578,
     60.29986134642249,
     60.542884027117346,
     60.713620046805545,
     60.96406868532684,
     61.24795911639664,
     61.50380262689339,
     61.676023944667456,
     61.91249438662608,
     62.19551841306711,
     62.46136864171307,
     62.64134130386411,
     62.8618916535012,
     63.14275550063443,
     63.41593714946038,
     63.60855696756366,
     63.81267964411834,
     64.08985017439029,
     64.36811250891209,
     64.57631762116566,
     64.76539855955836,
     65.03696631214747,
     65.31850683423295,
     65.5431828535302,
     65.72067625259793,
     65.9842702171477,
     66.26762912605196,
     66.50785686580033,
     66.67908458680958,
     66.93194933306637,
     67.21586952957189,
     67.46944687957124,
     67.64087207368169,
     67.88023387086244,
     68.16352180241984,
     68.42766247629135,
     68.6057132399554,
     68.82942313625892,
     69.11081269241872,
     69.38277739295187,
     69.57268550444225,
     69.7799151533709,
     70.05792785859981,
     70.33537213436817,
     70.54047645948188,
     70.73222831541355,
     71.0050332299864,
     71.28607167243781,
     71.50764112724117,
     71.68698192099966,
     71.95229366389803,
     72.23540934023758,
     72.47283414486189,
     72.64477721778965,
     72.89989146346213,
     73.18379776109897,
     73.43506458325183,
     73.60594743653633,
     73.84804734874952,
     74.13154742502378,
     74.3939217650586,
     74.57028533634853,
     74.79704592914061,
     75.07889613888426,
     75.34958948807117,
     75.53696768075024,
     75.74726521343875,
     76.02603595823757,
     76.30261476934344,
     76.5047322008677,
     76.69920149521533,
     76.97313539084057,
     77.25363013473333,
     77.47213803750353,
     77.65346140629822,
     77.92035844148852,
     78.20319168541128,
     78.4378007424199,
     78.61066580783138,
     78.86788301438096,
     79.15173466413779,
     79.40064180921354,
     79.57122086961452,
     79.81592130214112,
     80.09958691145333,
     80.36013305822898,
     80.53503292480471,
     80.76474440122267,
     81.04699793457125,
     81.31636104116674,
     81.50138351083318,
     81.7147114319083,
     81.99416675772241,
     82.26982959889115,
     82.46906939375938,
     82.66629667635232,
     82.94126500360997,
     83.22117299241245,
     83.43666200148992,
     83.62009104391728,
     83.88845639514504,
     84.17096823030235,
     84.4027471837285,
     84.57672669165763,
     84.83591512414371,
     85.11967325672309,
     85.3661691539245,
     85.53667137760691,
     85.7838457349754,
     86.06763391882056,
     86.32628639822366,
     86.49993900343861,
     86.73250690473772,
     87.01511211670491,
     87.2830823081631,
     87.46591989204883,
     87.68223998088733,
     87.9623144498927,
     88.23700790009333,
     88.43347841180586,
     88.6334975954157,
     88.90941620232971,
     89.1886927772784,
     89.40120619350648,
     89.5868527194551,
     89.85658136967506,
     90.13873260495185,
     90.36766791836288,
     90.54294177220034,
     90.803981074716,
     91.0876079921273,
     91.33164045612156,
     91.50228330110585,
     91.75181302405834,
     92.03568345897374,
     92.29237435358041,
     92.4649916181877,
     92.7003244914274,
     92.98323403147752,
     93.24974541871431,
     93.43056834472881,
     93.649840144344,
     93.93047452271378,
     94.20414238242675,
     94.39795361566829,
     94.60079152284162,
     94.87758443608601,
     95.15618320580225,
     95.36576710207841,
     95.55373211896544,
     95.82472858419692,
     96.10647948952229,
     96.33256020946943,
     96.50929672332327,
     96.77207563016299,
     97.055534296444,
     97.29705197058009,
     97.4680445358772,
     97.71981719652074,
     98.00373147096346,
     98.25839141022159,
     98.43018218489864,
     98.66819010184345,
     98.95135992922364,
     99.21634389178443,
     99.39532354569985,
     99.61750339646373,
     99.89864336603767}

```python
# these roots are obtained through Newton's method
RA_roots = [1.5990891427798726,
 1.788732691327088,
 1.932914344549613,
 2.2109110956513467,
 2.4419555167035885,
 2.5885848998466656,
 2.849989333932118,
 3.1281660491170373,
 3.3263461862747836,
 3.5089830748945015,
 3.783396450018647,
 4.0587499327999375,
 4.2500096554362,
 4.445555964175939,
 4.722261597342127,
 4.997520330348105,
 5.190937232329985,
 5.3875852765189185,
 5.663899469930367,
 5.940709869938875,
 6.140596614835184,
 6.33275567156643,
 6.607129174713171,
 6.886060588985002,
 7.095209499312184,
 7.280401466688491,
 7.551376030390238,
 7.8323370399311925,
 8.052565867522132,
 8.23044450603,
 8.49636374147896,
 8.778910896574448,
 9.011039166201218,
 9.1830697482723,
 9.441983167272015,
 9.725494441337466,
 9.969322368562036,
 10.138530274048614,
 10.388233654920867,
 10.671982483467206,
 10.926434114796178,
 11.096952757076904,
 11.335200046561987,
 11.618363938254081,
 11.881795898316074,
 12.05816134199271,
 12.283050046645862,
 12.564675538325883,
 12.83524123901948,
 13.021599981161451,
 13.232043976029324,
 13.510979803646782,
 13.78691103995045,
 13.986383685492326,
 14.182548338021387,
 14.457356685210335,
 14.737099288791287,
 14.951422765861851,
 15.135035478146149,
 15.403903238899153,
 15.686131192623376,
 15.915568038428134,
 16.09003453791654,
 16.35073878569254,
 16.63430083057258,
 16.87778792031103,
 17.04799247320583,
 17.298014892620287,
 17.5818530423035,
 17.837380305311136,
 18.009054889969647,
 18.245930540189043,
 18.528987245914934,
 18.794120668749507,
 18.97288575406163,
 19.194752835079072,
 19.475868961286285,
 19.748223167244905,
 19.938664451417566,
 20.144841261296854,
 20.422642802130465,
 20.700146618147762,
 20.905243784142034,
 21.096665490613677,
 21.36944521522357,
 21.65039765714026,
 21.871345260135367,
 22.050789185647496,
 22.316417308981585,
 22.59942216949087,
 22.83574787378633,
 23.007771571246685,
 23.26371901522057,
 23.547575723690084,
 23.79750807013827,
 23.967956594247525,
 24.211546230341284,
 24.495132983115276,
 24.75617125425975,
 24.93122756072884,
 25.160152468257767,
 25.442308614647423,
 25.711820575470526,
 25.896908812003076,
 26.109874943073603,
 26.38927822657481,
 26.66490500841016,
 26.863889556692698,
 27.061159119294334,
 27.33619656983615,
 27.615993811847783,
 27.830846095936195,
 28.014560849029866,
 28.283213658512377,
 28.565612463776073,
 28.796451204805997,
 28.970681191280395,
 29.23049057513832,
 29.514184601227317,
 29.759593065890563,
 29.929984806172033,
 30.17821705395824,
 30.462034619984344,
 30.719614250741852,
 30.892534456195584,
 31.126632864304202,
 31.40941075106034,
 31.676435610144512,
 31.857815694621625,
 32.07605391991435,
 32.35650997438538,
 32.63043881838727,
 32.82479975070046,
 33.026899688846825,
 33.303499500645344,
 33.5822062485808,
 33.792166497675645,
 33.97970628209389,
 34.250534985713415,
 34.532310974256845,
 34.75853318936432,
 34.9350861979419,
 35.197777341476126,
 35.481226092576726,
 35.72267335845731,
 35.89357884497937,
 36.145410432352826,
 36.42931503958064,
 36.68376431788118,
 36.85538911756186,
 37.093661935009315,
 37.376853342894364,
 37.641566154759616,
 37.82015879203966,
 38.0428288886042,
 38.32405564700252,
 38.59636827681748,
 38.78696222224,
 38.993306217048676,
 39.2710995220295,
 39.54873872479654,
 39.75451084722061,
 39.945606643276214,
 40.218145296597775,
 40.499279680878274,
 40.72139638545084,
 40.90033907234792,
 41.16535366141217,
 41.448503469395185,
 41.68631256019343,
 41.85808912541147,
 42.1129034130371,
 42.39680718255687,
 42.648302000506234,
 42.819174289219866,
 43.06101176135671,
 43.344490768519194,
 43.6069728126568,
 43.7833820332885,
 44.00995911201846,
 44.291784989158856,
 44.56250882104887,
 44.74990063619678,
 44.960117802384424,
 45.23887702049473,
 45.51544495342708,
 45.71749157319881,
 45.91197633558676,
 46.185931679348194,
 46.46639939369378,
 46.68474077765222,
 46.866131814599285,
 47.13310975153132,
 47.415916915687504,
 47.65028593541643,
 47.82319709880706,
 48.080585800164926,
 48.364425182999994,
 48.61305905456402,
 48.783578851174276,
 49.0285679637308,
 49.31224720988233,
 49.57252672412086,
 49.74719827622644,
 49.977321910123365,
 50.259629311327394,
 50.52875921312127,
 50.71336333188308,
 50.92719928775875,
 51.20676812274482,
 51.48224356648535,
 51.680899329614626,
 51.87866465525323,
 52.15383310023448,
 52.4336031467763,
 52.648401016127465,
 52.83229855812179,
 53.10098556555707,
 53.38340988653587,
 53.61446737441787,
 53.78872776035599,
 54.04839662812041,
 54.33212001649836,
 54.57793918434023,
 54.748429047040915,
 54.99626654235868,
 55.28007903144342,
 55.53815235637299,
 55.711445050649495,
 55.94484783996595,
 56.22754875066923,
 56.4950587870766,
 56.67720610669772,
 56.8944732101477,
 57.17473544232577,
 57.449085073822914,
 57.644612558834176,
 57.84558402295245,
 58.121813522455895,
 58.4008498688213,
 58.612279882020026,
 58.79874145803496,
 59.06894533675702,
 59.350947530622044,
 59.578781823129134,
 59.7545764538522,
 60.01629925881578,
 60.29986134642249,
 60.542884027117346,
 60.713620046805545,
 60.96406868532684,
 61.24795911639664,
 61.50380262689339,
 61.676023944667456,
 61.91249438662608,
 62.19551841306711,
 62.46136864171307,
 62.64134130386411,
 62.8618916535012,
 63.14275550063443,
 63.41593714946038,
 63.60855696756366,
 63.81267964411834,
 64.08985017439029,
 64.36811250891209,
 64.57631762116566,
 64.76539855955836,
 65.03696631214747,
 65.31850683423295,
 65.5431828535302,
 65.72067625259793,
 65.9842702171477,
 66.26762912605196,
 66.50785686580033,
 66.67908458680958,
 66.93194933306637,
 67.21586952957189,
 67.46944687957124,
 67.64087207368169,
 67.88023387086244,
 68.16352180241984,
 68.42766247629135,
 68.6057132399554,
 68.82942313625892,
 69.11081269241872,
 69.38277739295187,
 69.57268550444225,
 69.7799151533709,
 70.05792785859981,
 70.33537213436817,
 70.54047645948188,
 70.73222831541355,
 71.0050332299864,
 71.28607167243781,
 71.50764112724117,
 71.68698192099966,
 71.95229366389803,
 72.23540934023758,
 72.47283414486189,
 72.64477721778965,
 72.89989146346213,
 73.18379776109897,
 73.43506458325183,
 73.60594743653633,
 73.84804734874952,
 74.13154742502378,
 74.3939217650586,
 74.57028533634853,
 74.79704592914061,
 75.07889613888426,
 75.34958948807117,
 75.53696768075024,
 75.74726521343875,
 76.02603595823757,
 76.30261476934344,
 76.5047322008677,
 76.69920149521533,
 76.97313539084057,
 77.25363013473333,
 77.47213803750353,
 77.65346140629822,
 77.92035844148852,
 78.20319168541128,
 78.4378007424199,
 78.61066580783138,
 78.86788301438096,
 79.15173466413779,
 79.40064180921354,
 79.57122086961452,
 79.81592130214112,
 80.09958691145333,
 80.36013305822898,
 80.53503292480471,
 80.76474440122267,
 81.04699793457125,
 81.31636104116674,
 81.50138351083318,
 81.7147114319083,
 81.99416675772241,
 82.26982959889115,
 82.46906939375938,
 82.66629667635232,
 82.94126500360997,
 83.22117299241245,
 83.43666200148992,
 83.62009104391728,
 83.88845639514504,
 84.17096823030235,
 84.4027471837285,
 84.57672669165763,
 84.83591512414371,
 85.11967325672309,
 85.3661691539245,
 85.53667137760691,
 85.7838457349754,
 86.06763391882056,
 86.32628639822366,
 86.49993900343861,
 86.73250690473772,
 87.01511211670491,
 87.2830823081631,
 87.46591989204883,
 87.68223998088733,
 87.9623144498927,
 88.23700790009333,
 88.43347841180586,
 88.6334975954157,
 88.90941620232971,
 89.1886927772784,
 89.40120619350648,
 89.5868527194551,
 89.85658136967506,
 90.13873260495185,
 90.36766791836288,
 90.54294177220034,
 90.803981074716,
 91.0876079921273,
 91.33164045612156,
 91.50228330110585,
 91.75181302405834,
 92.03568345897374,
 92.29237435358041,
 92.4649916181877,
 92.7003244914274,
 92.98323403147752,
 93.24974541871431,
 93.43056834472881,
 93.649840144344,
 93.93047452271378,
 94.20414238242675,
 94.39795361566829,
 94.60079152284162,
 94.87758443608601,
 95.15618320580225,
 95.36576710207841,
 95.55373211896544,
 95.82472858419692,
 96.10647948952229,
 96.33256020946943,
 96.50929672332327,
 96.77207563016299,
 97.055534296444,
 97.29705197058009,
 97.4680445358772,
 97.71981719652074,
 98.00373147096346,
 98.25839141022159,
 98.43018218489864,
 98.66819010184345,
 98.95135992922364,
 99.21634389178443,
 99.39532354569985,
 99.61750339646373,
 99.89864336603767]
```

```python
# redifine residue_RA for plotting
def residue_RA(lam):
    mu = (beta * lam**2 + (beta - 1.0)/4.0)**(1/2)
    A = np.sin(lam*L1) + 2*lam*np.cos(lam*L1)
    B = np.sin(mu*L2) + 2*mu*np.cos(mu*L2)
    C = (1 + 4*mu**2)*(K_s2/K_s1)*np.sin(lam*L1)*np.sin(mu*L2)
    return A * B - C
```

```python
# check if there are no missing roots
plt.plot(lam_test, residue_RA(lam_test))
plt.scatter(RA_roots, residue_RA(np.array(RA_roots)), c = 'r')
plt.xlim(0, 11)
plt.ylim(-.25, .25)
```

```python
plt.plot(lam_test, residue_RA(lam_test))
plt.scatter(RA_roots, residue_RA(np.array(RA_roots)), c = 'r')
plt.xlim(10, 21)
plt.ylim(-.25, .25)
```

```python
plt.plot(lam_test, residue_RA(lam_test))
plt.scatter(RA_roots, residue_RA(np.array(RA_roots)), c = 'r')
plt.xlim(20, 31)
plt.ylim(-.25, .25)
```

```python
plt.plot(lam_test, residue_RA(lam_test))
plt.scatter(RA_roots, residue_RA(np.array(RA_roots)), c = 'r')
plt.xlim(30, 41)
plt.ylim(-.25, .25)
```

```python
plt.plot(lam_test, residue_RA(lam_test))
plt.scatter(RA_roots, residue_RA(np.array(RA_roots)), c = 'r')
plt.xlim(40, 51)
plt.ylim(-.25, .25)
```

```python
plt.plot(lam_test, residue_RA(lam_test))
plt.scatter(RA_roots, residue_RA(np.array(RA_roots)), c = 'r')
plt.xlim(50, 61)
plt.ylim(-.25, .25)
```

```python
plt.plot(lam_test, residue_RA(lam_test))
plt.scatter(RA_roots, residue_RA(np.array(RA_roots)), c = 'r')
plt.xlim(60, 71)
plt.ylim(-.25, .25)
```

```python
plt.plot(lam_test, residue_RA(lam_test))
plt.scatter(RA_roots, residue_RA(np.array(RA_roots)), c = 'r')
plt.xlim(70, 81)
plt.ylim(-.25, .25)
```

```python
plt.plot(lam_test, residue_RA(lam_test))
plt.scatter(RA_roots, residue_RA(np.array(RA_roots)), c = 'r')
plt.xlim(80, 91)
plt.ylim(-.25, .25)
```

```python
plt.plot(lam_test, residue_RA(lam_test))
plt.scatter(RA_roots, residue_RA(np.array(RA_roots)), c = 'r')
plt.xlim(90, 101)
plt.ylim(-.25, .25)
```

```python
residue_RA(np.array(RA_roots))
```

    array([-9.32587341e-15,  1.48769885e-14, -5.32907052e-15, -8.88178420e-15,
            8.88178420e-15, -6.21724894e-15, -4.88498131e-14,  8.88178420e-15,
            8.12683254e-14,  1.10134124e-13, -1.87405647e-13, -5.32907052e-15,
            4.27435864e-14, -9.59232693e-14, -7.65609798e-13,  1.38555833e-13,
            3.52384788e-13, -1.77635684e-13,  5.25801624e-13,  5.36459765e-13,
            9.01501096e-14, -3.19744231e-13,  1.09423581e-12, -3.76587650e-13,
            4.45865567e-13, -3.33955086e-13,  7.06990022e-13, -8.52651283e-14,
            1.42819090e-12, -1.12265752e-12,  4.61852778e-13,  7.81597009e-13,
           -1.66977543e-12, -7.10542736e-13, -1.38555833e-12, -2.32347475e-12,
           -1.00897068e-12, -2.41584530e-12,  2.81374923e-12, -1.22213351e-12,
            3.25428573e-12, -1.64845915e-12, -1.63424829e-12,  9.69890834e-13,
           -3.90798505e-12,  2.57216470e-12, -3.95061761e-12, -4.99156272e-13,
           -1.13686838e-13,  5.69144731e-12,  5.31485966e-12, -7.97051314e-12,
            5.05906428e-12,  2.16449081e-12,  1.10844667e-12,  7.31148475e-12,
            2.64321898e-12, -3.97903932e-12, -7.44648787e-12, -9.03810360e-12,
           -1.58024704e-11, -3.29691829e-12,  5.79802872e-12, -2.33342234e-11,
            1.75361947e-11,  1.80477855e-11,  2.98427949e-12,  2.58637556e-12,
            3.80850906e-12, -5.85487214e-12,  1.00612851e-11, -2.69437805e-11,
            1.97246663e-11, -5.28643795e-12, -1.44098067e-11, -5.22959454e-12,
           -1.05302433e-11,  1.40403245e-11,  1.80051529e-11,  4.11546353e-11,
           -3.26849658e-12,  4.28030944e-11, -9.52837809e-12,  2.48405740e-11,
            5.40865130e-11,  1.98383532e-11, -1.67688086e-11,  2.13162821e-11,
           -1.82183157e-11, -4.37694325e-11, -2.01794137e-11,  7.04858394e-12,
            5.15001375e-11, -1.63140612e-11,  1.13686838e-11, -7.73070497e-12,
            1.65414349e-11,  4.96811481e-11,  2.44426701e-11,  8.35598257e-12,
           -5.42286216e-11,  7.60280727e-12,  4.66116035e-12,  6.56541488e-11,
           -2.69437805e-11, -7.59676766e-11,  1.13686838e-13,  3.00257597e-11,
            3.50155460e-11, -1.03540287e-10,  2.27373675e-11, -5.00222086e-12,
            6.38920028e-11,  9.72022463e-11,  7.98081601e-11, -2.53521648e-11,
            5.46833689e-11, -5.99129635e-11, -2.02931005e-11,  2.83080226e-11,
           -1.04591891e-11, -2.59205990e-11, -1.07036158e-10,  2.72848411e-11,
            8.29913915e-12,  9.36779543e-11,  5.89466254e-11,  4.34283720e-11,
            8.60040927e-11, -8.09450285e-11,  2.00515160e-11, -2.04636308e-11,
            8.62883098e-11,  2.56932253e-11, -8.60751470e-11, -1.04591891e-10,
            8.03765943e-11, -1.08457243e-10, -4.76916284e-11, -2.75122147e-11,
           -3.97903932e-12, -4.38831194e-11, -1.51771928e-10, -3.93356459e-11,
            1.08457243e-10,  4.45652404e-11, -6.33235686e-11, -3.19460014e-11,
            2.75122147e-11,  9.77706804e-11,  1.25055521e-10, -1.06183506e-10,
           -1.36424205e-11, -9.43600753e-12, -1.11867848e-10,  8.23661139e-11,
            9.07220965e-11,  4.97664132e-11,  5.25233190e-11,  3.81461973e-10,
           -9.14042175e-11,  4.77768936e-11,  3.43334250e-11,  4.82032192e-11,
            7.04858394e-12,  1.21644916e-10, -1.17324817e-10,  2.30215846e-10,
           -1.01408659e-10, -1.45291779e-10,  9.43600753e-11,  2.67050382e-10,
           -2.28396857e-10, -3.48336471e-10, -9.45874490e-11, -3.74711817e-10,
            2.73985279e-10, -9.27684596e-11, -9.74296199e-11,  2.09183781e-11,
           -1.09253051e-10, -2.21689334e-10, -9.99307304e-11,  3.55157681e-10,
           -1.51144874e-10,  1.74622983e-10, -2.54601673e-10, -1.60298441e-10,
           -8.88462637e-11, -4.03360900e-10,  2.42948772e-10,  4.68389771e-11,
           -9.87938620e-11, -6.68478606e-11, -1.20508048e-11,  4.09727363e-10,
            2.73075784e-10,  5.04996933e-10,  2.20552465e-11,  5.02495823e-11,
            4.70663508e-11, -2.30443220e-10, -3.54702934e-10,  2.05773176e-10,
           -3.83352017e-10,  1.40062184e-10, -2.50111043e-11,  4.59749572e-10,
            9.23137122e-11, -6.17035312e-10,  5.13864507e-11, -3.47011309e-10,
            4.16548573e-10,  9.29389898e-11, -3.69709596e-10, -2.48974175e-10,
           -1.12777343e-10,  7.32370609e-10,  6.34372554e-10, -1.14823706e-10,
            3.70164344e-10,  5.41831469e-10, -3.16049409e-11,  6.82121026e-12,
            1.54159352e-10, -2.73075784e-10, -2.07592166e-10,  2.56477506e-10,
            7.91260391e-11, -9.36779543e-11,  1.07490905e-09,  3.10592441e-10,
           -3.09455572e-10, -1.49611878e-10, -3.57886165e-10, -2.03272066e-10,
           -1.80392590e-10,  4.32009983e-11,  3.04055447e-10,  3.08318704e-10,
            1.68256520e-10, -1.81898940e-12,  2.75122147e-10, -2.35104380e-10,
           -2.13276508e-10,  2.61024979e-10, -8.36962499e-10, -9.69521352e-10,
            8.23092705e-11, -4.73391992e-10,  5.06133802e-10, -6.92807589e-10,
           -2.14640750e-10,  5.93900040e-10, -2.36468622e-10,  9.02673492e-11,
           -4.92036634e-10, -2.09183781e-11, -9.90439730e-10,  5.92308425e-11,
            7.83984433e-10, -7.06563696e-11, -4.62023309e-10, -1.48503432e-10,
            3.28782335e-10, -1.02534159e-09, -1.93267624e-10,  1.43450052e-09,
            5.11590770e-10,  2.76486389e-10, -2.29192665e-10, -7.41692929e-10,
           -1.89584171e-09, -2.27373675e-12,  2.91947799e-10,  1.00726538e-09,
            3.81987775e-10,  2.75122147e-10, -5.42513590e-10, -6.54836185e-11,
           -1.23782229e-09,  3.51974450e-10,  6.54381438e-10,  6.82121026e-12,
           -5.60930857e-10,  6.28460839e-10,  7.56017471e-11,  9.95896698e-10,
           -2.76557444e-10, -8.55834514e-10,  1.55068847e-10, -3.78349796e-10,
           -7.13839654e-10, -1.21872290e-10, -3.65616870e-10,  7.41238182e-10,
            8.01719580e-10, -4.61568561e-10, -7.32143235e-11,  8.73114914e-10,
            2.06955519e-09,  6.78028300e-10,  7.09405867e-11, -2.03272066e-10,
           -1.00271791e-09,  1.80671123e-09,  1.16324372e-09, -3.52883944e-10,
            1.27238309e-09, -4.91354513e-10,  4.79303708e-10,  1.22804522e-09,
            5.66615199e-10, -2.96097369e-10,  1.54614099e-09,  5.43991518e-11,
           -1.31785782e-09,  3.36171979e-10,  5.29325916e-10, -1.51339918e-09,
            1.33513822e-09, -2.24645191e-09,  8.27640179e-10,  8.58562998e-10,
            1.06410880e-09,  2.87900548e-09,  3.53793439e-10, -4.40195436e-10,
           -8.35825631e-10, -2.24190444e-09, -6.68478606e-10, -2.55568011e-10,
            6.01630745e-10, -1.27238309e-09,  1.55910129e-09,  9.96806193e-10,
            1.18393473e-09,  5.73891157e-10, -3.56760665e-09,  1.29330147e-09,
           -1.02687636e-10,  1.14869181e-09, -2.72802936e-09, -7.53061613e-10,
            1.52112989e-09,  1.28147803e-09,  1.46746970e-09, -1.31058187e-09,
            1.94631866e-10, -9.90439730e-10, -1.57024260e-09,  1.83717930e-10,
           -1.09139364e-10,  6.60293153e-10,  9.01309249e-10, -3.14230419e-10,
            5.70253178e-10,  3.69254849e-10,  6.18456397e-11,  1.15824150e-09,
            7.42147677e-10,  1.01636033e-09, -6.60293153e-10, -6.60065780e-10,
            1.35514711e-10, -3.66355835e-10, -2.19461072e-09,  3.11263193e-09,
           -1.72531145e-09, -7.63520802e-10, -3.81078280e-10,  1.42154022e-09,
            9.97715688e-10, -2.68300937e-10,  2.74667400e-10, -1.51794666e-09,
            9.38598532e-10,  5.23868948e-10,  1.88902050e-09, -7.72161002e-10,
            1.38015821e-09,  4.91127139e-10,  4.06544132e-10, -9.36779543e-10,
            2.70347300e-09, -1.74804882e-09, -2.88309820e-10,  1.27056410e-09,
           -2.80260792e-09,  1.51885615e-09, -1.47349510e-09,  1.11504050e-09,
           -1.76896719e-09,  1.13413989e-09,  3.23780114e-10, -1.42063072e-09,
           -2.55522536e-09,  3.64616426e-09, -1.74168235e-09, -3.00133252e-10,
           -1.54932422e-09,  2.63571565e-09,  1.65709935e-09, -1.46519596e-09,
            3.39059625e-09,  4.29099600e-09, -2.96131475e-09, -1.83717930e-10,
            2.81943358e-10, -1.49702828e-09,  8.27640179e-10, -8.80390871e-10,
           -1.30057742e-09,  3.47881723e-11])

```python
# summation in Equation 22a and 22b in their paper (Srivastva and Yeh, 1991)
RA1 = np.zeros(t1_mesh.shape)
RA2 = np.zeros(t2_mesh.shape)

for root in RA_roots:
    mu = (beta*root**2 + (beta - 1.0)/4.0)**(1/2)
    A_ss = 4 * beta * root *K_s2/K_s1 + root*(L1 + beta* L2)
    A_sc = K_s2/K_s1*(beta * root * L2)/(2*mu)*(1 + 4*mu**2) + 2*root*mu*L1- (beta * root)/(2*mu)*(2 + L2)
    A_cs = K_s2/K_s1*L1/2*(1 + 4*mu**2) + 2*beta*root**2*L2 - 1.0 - L1/2.0
    A_cc = -mu*(2.0 + L1) - beta*root**2/mu*(2.0 + L2)

    D_n = (1/4 + root**2)/(mu*root)* \
         (A_ss*np.sin(root*L1)*np.sin(mu*L2) + \
          A_sc*np.sin(root*L1)*np.cos(mu*L2) + \
          A_cs*np.cos(root*L1)*np.sin(mu*L2) + \
          A_cc*np.cos(root*L1)*np.cos(mu*L2))

    RA1 += np.sin(root*(L1 + z1_mesh))*np.exp(-(1/4 + root**2)*t1_mesh)/D_n
    RA2 += (np.sin(root*L1) * (np.sin(mu*(L2 - z2_mesh)) + 2*mu*np.cos(mu*(L2 - z2_mesh)))\
            * np.exp(-(1/4 + root**2)*t2_mesh))/(D_n * (np.sin(mu*L2) + 2*mu*np.cos(mu*L2)))
```

In the example, $\beta = 0.1$ and thus we only use Case C. Also, Because $(\frac{1 - \beta}{4})^{1/2} = 0.4743...$, we only take into accout of $\lambda < 0.474$.

```python
((1 - beta)/(4))**(1/2)
```

    0.4743416490252569

```python
lam_test = np.arange(0, 0.474, 0.001)
```

```python
# two sign changes: one is due to discontinuity, and the other one is the root
pd.set_option('display.max_rows', None)
data = {"lam": lam_test.flatten()[0:-1],
        "sign": np.sign(residue_RC(lam_test)).flatten()[0:-1],
        "change": np.sign(residue_RC(lam_test)).flatten()[0:-1] - np.sign(residue_RC(lam_test)).flatten()[1:]}
df = pd.DataFrame(data)

a_set = df.loc[np.abs(df["change"]) == 2]["lam"].values
```

```python
# program to find roots using Newton method

def residue_RC(mu):
    lam = ((1 - beta)/(4*beta) - mu**2/beta)**(1/2)

    A = mpmath.sin(lam*L1) + 2*lam*mpmath.cos(lam*L1)
    B = mpmath.sinh(mu*L2) + 2*mu*mpmath.cosh(mu*L2)
    C = (1 - 4*mu**2)*(K_s2/K_s1)*mpmath.sin(lam*L1)*mpmath.sinh(mu*L2)
    return A * B - C

roots = []

for a in a_set:
    a = int(a/0.001)
    for i in range(a, a + 10):
        print(i)
        initial_value = 0.001 * (i + 1)
        root = mpmath.findroot(residue_RC, initial_value, solver = 'newton')
        root_value = np.array(root, np.float).tolist()
        roots.append(root_value)
```

```python
set(roots)
```

    {0.23059017983575425,
     0.34753161404921556,
     0.41281772029876773,
     0.4502656328080065,
     0.46889797480709333}

```python
# these roots are obtained through Newton's method
RC_roots = [
0.23059017983575425,
 0.34753161404921556,
 0.41281772029876773,
 0.4502656328080065,
 0.46889797480709333
]
```

```python
# redifine for the plotting
def residue_RC(mu):
    lam = ((1 - beta)/(4*beta) - mu**2/beta)**(1/2)

    A = np.sin(lam*L1) + 2*lam*np.cos(lam*L1)
    B = np.sinh(mu*L2) + 2*mu*np.cosh(mu*L2)
    C = (1 - 4*mu**2)*(K_s2/K_s1)*np.sin(lam*L1)*np.sinh(mu*L2)
    return A * B - C
```

```python
plt.plot(lam_test, residue_RC(lam_test))
plt.scatter(RC_roots, residue_RC(np.array(RC_roots)), c = 'r')
plt.xlim(0, 0.5)
plt.ylim(-.25, .25)
```

```python
residue_RC(np.array(RC_roots))
```

    array([ 5.32907052e-14,  1.91846539e-13, -7.10542736e-14,  9.94759830e-14,
            1.84741111e-13])

```python
# summation in Equation 28a and 28b in their paper (Srivastva and Yeh, 1991)
RC1 = np.zeros(t1_mesh.shape)
RC2 = np.zeros(t2_mesh.shape)

for mu in RC_roots:
    root = ((1 - beta)/(4*beta) - mu**2/beta)**(1/2)
    C_ss = 4 * beta * root *K_s2/K_s1 + root*(L1 + beta* L2)
    C_sc = -K_s2/K_s1*(beta * root * L2)/(2*mu)*(1 - 4*mu**2) + 2*root*mu*L1 + (beta * root)/(2*mu)*(2 + L2)
    C_cs = K_s2/K_s1*L1/2*(1 - 4*mu**2) + 2*beta*root**2*L2 - 1.0 - L1/2.0
    C_cc = -mu*(2.0 + L1) + beta*root**2/mu*(2 + L2)

    D_n = (1/4 + root**2)/(mu*root)* \
         (C_ss*np.sin(root*L1)*np.sinh(mu*L2) + \
          C_sc*np.sin(root*L1)*np.cosh(mu*L2) + \
          C_cs*np.cos(root*L1)*np.sinh(mu*L2) + \
          C_cc*np.cos(root*L1)*np.cosh(mu*L2))

    RC1 += np.sin(root*(L1 + z1_mesh))*np.exp(-(1/4 + root**2)*t1_mesh)/D_n
    RC2 += (np.sin(root*L1) * (np.sinh(mu*(L2 - z2_mesh)) + 2*mu*np.cosh(mu*(L2 - z2_mesh)))\
            * np.exp(-(1/4 + root**2)*t2_mesh))/(D_n * (np.sinh(mu*L2) + 2*mu*np.cosh(mu*L2)))
```

```python
# Equation 30a and 30b in their paper (Srivastva and Yeh, 1991)
# note the analytical solution assumes alpha1 = alpha2

K_1 = q_B1 - (q_B1 - np.exp(alpha1 * psi_0)) * np.exp(-(L1 + z1_mesh)) \
      - 4 * (q_B1 - q_A1) * np.exp((L2 - z1_mesh)/2) * (RA1 + RC1)

K_2 = q_B2 - (q_B2 - q_B1 + (q_B1 - np.exp(alpha2 * psi_0)) * np.exp(-L1)) * np.exp(-z2_mesh) \
      - 4 * (q_B1 - q_A1) * np.exp((L2 - z2_mesh)/2) * (RA2 + RC2)

```

```python
# transformation
K_1_star = K_1 * K_s1
K_2_star = K_2 * K_s2
psi_1 = np.log(K_1_star/K_s1)/alpha1
psi_2 = np.log(K_2_star/K_s2)/alpha2
```

```python
psi_ini = np.concatenate((psi_2_ini, psi_1_ini), axis = 0)
z_star = np.concatenate((z2_star, z1_star), axis = 0)
psi = np.concatenate((psi_2, psi_1), axis = 0)
```

```python
# the initial condition is used for the solution at t = 0 because the numerical solution is unstable at t = 0.
psi[:, 0] = psi_ini.flatten()
```

```python
plt.plot(psi[:, 0], z_star, label = "t = 0 hour")
plt.plot(psi[:, 1], z_star, label = "t = 0.1 hour")
plt.plot(psi[:, 5], z_star, label = "t = 0.5 hour")
plt.plot(psi[:, 10], z_star, label = "t = 1 hour")
plt.plot(psi[:, 20], z_star, label = "t = 2 hour")
plt.plot(psi[:, 50], z_star, label = "t = 5 hour")
plt.plot(psi[:, 100], z_star, label = "t = 10 hour")

plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left',)
plt.xlabel("$\\psi$ cm")
plt.ylabel("$z$ cm")
plt.xlim(-5, 0)
plt.ylim(-10, 10)
```

```python
# save solutions
# psi: each column for each time (0, 0.1, 0.2, ..., 10 hours)
# each row for each depth (10, 9.9, ..., -10.0 cm)
```

```python
# numpy format
np.save("Srivastava_psi_heterogeneous", psi)
# matlab format
matlab_data = {'data':psi}
scipy.io.savemat("Srivastava_psi_heterogeneous", matlab_data)
```


```python
%load_ext autoreload
%autoreload 2

%matplotlib inline
```

```python
#export
from exp.nb_06 import *
```

## ConvNet

Let's get the data and training interface from where we left in the last notebook.

[Jump_to lesson 10 video](https://course.fast.ai/videos/?lesson=10&t=5899)

```python
x_train,y_train,x_valid,y_valid = get_data()

x_train,x_valid = normalize_to(x_train,x_valid)
train_ds,valid_ds = Dataset(x_train, y_train),Dataset(x_valid, y_valid)

nh,bs = 50,512
c = y_train.max().item()+1
loss_func = F.cross_entropy

data = DataBunch(*get_dls(train_ds, valid_ds, bs), c)
```

```python
mnist_view = view_tfm(1,28,28)
cbfs = [Recorder,
        partial(AvgStatsCallback,accuracy),
        CudaCallback,
        partial(BatchTransformXCallback, mnist_view)]
```

```python
nfs = [8,16,32,64,64]
```

```python
learn,run = get_learn_run(nfs, data, 0.4, conv_layer, cbs=cbfs)
```

    Metrics:
    <function accuracy at 0x7feebb395e18>
    Metrics:
    <function accuracy at 0x7feebb395e18>

```python
%time run.fit(2, learn)
```

    train: [1.6451540625, tensor(0.4413, device='cuda:0')]
    valid: [0.405144921875, tensor(0.8681, device='cuda:0')]
    train: [0.269253515625, tensor(0.9165, device='cuda:0')]
    valid: [0.1485333984375, tensor(0.9509, device='cuda:0')]
    CPU times: user 12 s, sys: 3.02 s, total: 15 s
    Wall time: 27 s

## Batchnorm

### Custom

Let's start by building our own `BatchNorm` layer from scratch.

[Jump_to lesson 10 video](https://course.fast.ai/videos/?lesson=10&t=6018)

```python
class BatchNorm(nn.Module): # nf number of filters
    def __init__(self, nf, mom=0.1, eps=1e-5):
        super().__init__()
        # NB: pytorch bn mom is opposite of what you'd expect
        self.mom,self.eps = mom,eps
        self.mults = nn.Parameter(torch.ones (nf,1,1))
        self.adds  = nn.Parameter(torch.zeros(nf,1,1))
        self.register_buffer('vars',  torch.ones(1,nf,1,1))
        self.register_buffer('means', torch.zeros(1,nf,1,1))

    def update_stats(self, x):
        m = x.mean((0,2,3), keepdim=True)
        v = x.var ((0,2,3), keepdim=True)
        self.means.lerp_(m, self.mom)
        self.vars.lerp_ (v, self.mom)
        return m,v

    def forward(self, x):
        if self.training:
            with torch.no_grad(): m,v = self.update_stats(x)
            print(m)
        else: m,v = self.means,self.vars
        x = (x-m) / (v+self.eps).sqrt()
        return x*self.mults + self.adds
```

```python
def conv_layer(ni, nf, ks=3, stride=2, bn=True, **kwargs):
    # No bias needed if using bn
    layers = [nn.Conv2d(ni, nf, ks, padding=ks//2, stride=stride, bias=not bn),
              GeneralRelu(**kwargs)]
    if bn: layers.append(BatchNorm(nf))
    return nn.Sequential(*layers)
```

```python
#export
def init_cnn_(m, f):
    if isinstance(m, nn.Conv2d):
        f(m.weight, a=0.1)
        if getattr(m, 'bias', None) is not None: m.bias.data.zero_()
    for l in m.children(): init_cnn_(l, f)

def init_cnn(m, uniform=False):
    f = init.kaiming_uniform_ if uniform else init.kaiming_normal_
    init_cnn_(m, f)

def get_learn_run(nfs, data, lr, layer, cbs=None, opt_func=None, uniform=False, **kwargs):
    model = get_cnn_model(data, nfs, layer, **kwargs)
    init_cnn(model, uniform=uniform)
    return get_runner(model, data, lr=lr, cbs=cbs, opt_func=opt_func)
```

We can then use it in training and see how it helps keep the activations means to 0 and the std to 1.

```python
learn,run = get_learn_run(nfs, data, 0.9, conv_layer, cbs=cbfs)
```

    Metrics:
    <function accuracy at 0x7feebb395e18>
    Metrics:
    <function accuracy at 0x7feebb395e18>

```python
learn.model[0][2].mults, learn.model[0][2].adds, learn.model[0][0].weight, learn.model
```

    (Parameter containing:
     tensor([[[1.4311]],

             [[1.6707]],

             [[1.4983]],

             [[0.7902]],

             [[1.1550]],

             [[0.9599]],

             [[1.3351]],

             [[1.0985]]], device='cuda:0', requires_grad=True),
     Parameter containing:
     tensor([[[ 0.1510]],

             [[-0.1416]],

             [[-0.0154]],

             [[ 0.2880]],

             [[ 0.0074]],

             [[-0.2211]],

             [[-0.0891]],

             [[-0.0758]]], device='cuda:0', requires_grad=True),
     Parameter containing:
     tensor([[[[ 3.2673e-01, -4.0895e-01, -4.6189e-01, -2.3045e-01, -1.4938e-01],
               [ 1.2137e-01, -2.0384e-01, -9.2740e-01, -4.6862e-01, -2.8914e-01],
               [-7.7978e-01, -1.4199e+00, -6.2296e-01, -1.7435e-01,  2.2356e-01],
               [-8.2712e-01, -6.1046e-01, -4.0323e-02,  2.8165e-01,  1.0199e+00],
               [-7.7368e-01, -4.8329e-01,  4.2812e-01,  7.2656e-01,  7.5583e-01]]],

             [[[-5.2581e-01,  2.5834e-02,  5.2102e-03,  3.2452e-01,  5.2763e-02],
               [-7.8672e-01, -7.5356e-01, -5.6800e-01, -4.3620e-01, -7.5392e-01],
               [ 1.6299e-01, -5.4327e-01, -8.5035e-01, -6.8228e-01, -3.4117e-01],
               [ 7.8579e-01,  8.7069e-01,  1.0799e-01,  1.0407e-01,  2.9231e-01],
               [ 6.4458e-01,  3.1410e-01,  7.9506e-01,  6.1506e-01,  6.0299e-01]]],

             [[[-1.4695e-01, -1.8160e-01, -1.4631e-01,  6.1111e-02, -2.4257e-01],
               [-3.6590e-01, -2.0526e-02,  1.6840e-01,  3.6619e-01, -1.0112e-01],
               [-2.0774e-01, -2.5164e-01,  1.7113e-01,  4.3212e-01,  4.8603e-01],
               [-1.6399e-01,  2.0958e-01,  4.7967e-01,  8.8115e-01,  5.5735e-01],
               [ 4.8136e-01,  3.7458e-01,  2.7670e-01,  6.3475e-01,  6.4163e-01]]],

             [[[-5.3348e-01,  1.1957e-01, -1.6873e-02,  3.8616e-01,  4.3201e-01],
               [ 3.0384e-02, -5.3512e-01,  9.0575e-02,  5.8774e-01,  5.4507e-02],
               [-4.5384e-01, -2.4965e-01,  2.1142e-01,  5.8174e-01,  2.7110e-01],
               [-5.4699e-02,  2.9106e-01,  7.1928e-01, -1.6924e-01, -9.2375e-02],
               [-2.5668e-01, -5.9731e-02, -7.6865e-03,  3.9907e-01, -6.3391e-01]]],

             [[[ 2.4487e-02,  2.7824e-01,  1.9594e-01, -4.0299e-01, -7.3391e-01],
               [ 2.2083e-01,  1.8199e-01,  2.7246e-01, -1.5728e-01, -1.3238e+00],
               [ 2.2811e-01,  2.2175e-01,  1.7981e-01, -5.6557e-01, -7.0445e-01],
               [ 4.3397e-01,  2.2295e-01, -4.2498e-01, -7.8814e-01, -3.0519e-01],
               [ 3.6147e-01, -4.3284e-01, -6.0052e-01, -3.2484e-01,  4.1524e-02]]],

             [[[-3.2247e-01,  1.7705e-01, -3.3328e-01, -1.8279e-01, -3.4857e-01],
               [ 7.2165e-02, -3.8473e-01,  1.4174e-03,  2.9784e-01, -6.1307e-01],
               [ 4.3566e-01,  8.0468e-01,  3.7751e-01, -1.1627e-01,  2.6047e-01],
               [ 6.3873e-01,  2.2535e-01, -3.4538e-01,  1.1229e-02, -2.3223e-01],
               [ 4.9852e-01,  6.8637e-01,  4.1385e-01, -5.0141e-01,  5.0434e-01]]],

             [[[-4.4840e-01, -4.3504e-02,  3.7060e-01,  4.2998e-01,  8.3362e-02],
               [ 2.1569e-01,  4.8446e-01,  7.3055e-01,  4.5360e-01,  1.0359e+00],
               [-7.9489e-02,  6.4323e-01,  5.4897e-01,  1.2266e-01,  1.0892e-01],
               [-4.1418e-01,  1.0374e-01,  1.9722e-01,  9.0322e-02,  1.9639e-01],
               [-3.4150e-01, -6.1588e-01, -3.8859e-01, -3.1969e-01, -2.8520e-01]]],

             [[[-2.6920e-01,  5.2437e-01, -2.4185e-01, -4.1845e-01,  5.4522e-02],
               [ 7.4007e-02,  1.7739e-01, -3.6910e-01,  2.6857e-01, -2.6047e-01],
               [ 6.1573e-02,  2.9422e-01,  9.8547e-02,  2.1217e-01, -5.3816e-01],
               [ 3.8280e-01,  6.9459e-01,  1.9069e-01,  4.7545e-01, -5.8502e-01],
               [ 4.4372e-01,  2.0806e-01,  4.6171e-01,  1.3153e-01, -6.4725e-01]]]],
            device='cuda:0', requires_grad=True),
     Sequential(
       (0): Sequential(
         (0): Conv2d(1, 8, kernel_size=(5, 5), stride=(2, 2), padding=(2, 2), bias=False)
         (1): GeneralRelu()
         (2): RunningBatchNorm()
       )
       (1): Sequential(
         (0): Conv2d(8, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
         (1): GeneralRelu()
         (2): RunningBatchNorm()
       )
       (2): Sequential(
         (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
         (1): GeneralRelu()
         (2): RunningBatchNorm()
       )
       (3): Sequential(
         (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
         (1): GeneralRelu()
         (2): RunningBatchNorm()
       )
       (4): Sequential(
         (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
         (1): GeneralRelu()
         (2): RunningBatchNorm()
       )
       (5): AdaptiveAvgPool2d(output_size=1)
       (6): Lambda()
       (7): Linear(in_features=64, out_features=10, bias=True)
     ))

```python
def append_stats(hook, mod, inp, outp):
    if not hasattr(hook,'stats'): hook.stats = ([],[])
    means,stds = hook.stats
    means.append(outp.data.mean())
    stds.append(outp.data.std())
```

```python
with Hooks(learn.model, append_stats) as hooks:
    run.fit(1, learn)
    fig,(ax0,ax1) = plt.subplots(1,2, figsize=(10,4))
    for h in hooks[:-1]:
        ms,ss = h.stats
        ax0.plot(ms[:10])
        ax1.plot(ss[:10])
        h.remove()
    plt.legend(range(6));

    fig,(ax0,ax1) = plt.subplots(1,2, figsize=(10,4))
    for h in hooks[:-1]:
        ms,ss = h.stats
        ax0.plot(ms)
        ax1.plot(ss)
```

```python
learn,run = get_learn_run(nfs, data, 1.0, conv_layer, cbs=cbfs)
```

    Metrics:
    <function accuracy at 0x7f27aa6f57b8>
    Metrics:
    <function accuracy at 0x7f27aa6f57b8>

```python
%time run.fit(3, learn)
```

    train: [0.25310638671875, tensor(0.9195, device='cuda:0')]
    valid: [0.1629533935546875, tensor(0.9513, device='cuda:0')]
    train: [0.077422421875, tensor(0.9765, device='cuda:0')]
    valid: [0.11033271484375, tensor(0.9655, device='cuda:0')]
    train: [0.054823671875, tensor(0.9835, device='cuda:0')]
    valid: [0.07338250732421875, tensor(0.9768, device='cuda:0')]
    CPU times: user 12.5 s, sys: 2.98 s, total: 15.5 s
    Wall time: 15.5 s

### Builtin batchnorm

[Jump_to lesson 10 video](https://course.fast.ai/videos/?lesson=10&t=6679)

```python
#export
def conv_layer(ni, nf, ks=3, stride=2, bn=True, **kwargs):
    layers = [nn.Conv2d(ni, nf, ks, padding=ks//2, stride=stride, bias=not bn),
              GeneralRelu(**kwargs)]
    if bn: layers.append(nn.BatchNorm2d(nf, eps=1e-5, momentum=0.1))
    return nn.Sequential(*layers)
```

```python
learn,run = get_learn_run(nfs, data, 1., conv_layer, cbs=cbfs)
```

    Metrics:
    <function accuracy at 0x7f27aa6f57b8>
    Metrics:
    <function accuracy at 0x7f27aa6f57b8>

```python
%time run.fit(3, learn)
```

    train: [0.261401640625, tensor(0.9238, device='cuda:0')]
    valid: [0.1479271240234375, tensor(0.9517, device='cuda:0')]
    train: [0.072921455078125, tensor(0.9772, device='cuda:0')]
    valid: [0.08903743896484374, tensor(0.9733, device='cuda:0')]
    train: [0.052027099609375, tensor(0.9838, device='cuda:0')]
    valid: [0.06108431396484375, tensor(0.9804, device='cuda:0')]
    CPU times: user 11.3 s, sys: 2.37 s, total: 13.7 s
    Wall time: 13.8 s

### With scheduler

Now let's add the usual warm-up/annealing.

```python
sched = combine_scheds([0.3, 0.7], [sched_lin(0.6, 2.), sched_lin(2., 0.1)])
```

```python
learn,run = get_learn_run(nfs, data, 0.9, conv_layer, cbs=cbfs
                          +[partial(ParamScheduler,'lr', sched)])
```

    Metrics:
    <function accuracy at 0x7f27aa6f57b8>
    Metrics:
    <function accuracy at 0x7f27aa6f57b8>

```python
run.fit(8, learn)
```

    train: [0.31647228515625, tensor(0.9062, device='cuda:0')]
    valid: [0.141310986328125, tensor(0.9565, device='cuda:0')]
    train: [0.09898029296875, tensor(0.9701, device='cuda:0')]
    valid: [0.11971199951171875, tensor(0.9652, device='cuda:0')]
    train: [0.06806630859375, tensor(0.9786, device='cuda:0')]
    valid: [0.07510550537109376, tensor(0.9766, device='cuda:0')]
    train: [0.0429486865234375, tensor(0.9863, device='cuda:0')]
    valid: [0.0596386962890625, tensor(0.9814, device='cuda:0')]
    train: [0.0292974560546875, tensor(0.9906, device='cuda:0')]
    valid: [0.05085914001464844, tensor(0.9857, device='cuda:0')]
    train: [0.01904771240234375, tensor(0.9945, device='cuda:0')]
    valid: [0.044835992431640624, tensor(0.9863, device='cuda:0')]
    train: [0.01245054443359375, tensor(0.9970, device='cuda:0')]
    valid: [0.04183270263671875, tensor(0.9878, device='cuda:0')]
    train: [0.008412266845703125, tensor(0.9986, device='cuda:0')]
    valid: [0.042111663818359374, tensor(0.9877, device='cuda:0')]

## More norms

### Layer norm

From [the paper](https://arxiv.org/abs/1607.06450): "*batch normalization cannot be applied to online learning tasks or to extremely large distributed models where the minibatches have to be small*".

General equation for a norm layer with learnable affine:

$$y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta$$

The difference with BatchNorm is
1. we don't keep a moving average
2. we don't average over the batches dimension but over the hidden dimension, so it's independent of the batch size

[Jump_to lesson 10 video](https://course.fast.ai/videos/?lesson=10&t=6717)

```python
class LayerNorm(nn.Module):
    __constants__ = ['eps']
    def __init__(self, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.mult = nn.Parameter(tensor(1.))
        self.add  = nn.Parameter(tensor(0.))

    def forward(self, x):
        m = x.mean((1,2,3), keepdim=True)
        v = x.var ((1,2,3), keepdim=True)
        x = (x-m) / ((v+self.eps).sqrt())
        return x*self.mult + self.add
```

```python
def conv_ln(ni, nf, ks=3, stride=2, bn=True, **kwargs):
    layers = [nn.Conv2d(ni, nf, ks, padding=ks//2, stride=stride, bias=True),
              GeneralRelu(**kwargs)]
    if bn: layers.append(LayerNorm())
    return nn.Sequential(*layers)
```

```python
learn,run = get_learn_run(nfs, data, 0.8, conv_ln, cbs=cbfs)
```

    Metrics:
    <function accuracy at 0x7f27aa6f57b8>
    Metrics:
    <function accuracy at 0x7f27aa6f57b8>

```python
%time run.fit(3, learn)
```

    train: [nan, tensor(0.1210, device='cuda:0')]
    valid: [nan, tensor(0.0991, device='cuda:0')]
    train: [nan, tensor(0.0986, device='cuda:0')]
    valid: [nan, tensor(0.0991, device='cuda:0')]
    train: [nan, tensor(0.0986, device='cuda:0')]
    valid: [nan, tensor(0.0991, device='cuda:0')]
    CPU times: user 15.4 s, sys: 3.99 s, total: 19.4 s
    Wall time: 19.4 s

*Thought experiment*: can this distinguish foggy days from sunny days (assuming you're using it before the first conv)?

### Instance norm

From [the paper](https://arxiv.org/abs/1607.08022):

The key difference between **contrast** and batch normalization is that the latter applies the normalization to a  whole batch of images instead for single ones:

\begin{equation}\label{eq:bnorm}
    y_{tijk} =  \frac{x_{tijk} - \mu_{i}}{\sqrt{\sigma_i^2 + \epsilon}},
    \quad
    \mu_i = \frac{1}{HWT}\sum_{t=1}^T\sum_{l=1}^W \sum_{m=1}^H x_{tilm},
    \quad
    \sigma_i^2 = \frac{1}{HWT}\sum_{t=1}^T\sum_{l=1}^W \sum_{m=1}^H (x_{tilm} - mu_i)^2.
\end{equation}

In order to combine the effects of instance-specific normalization and batch normalization, we propose to replace the latter by the *instance normalization* (also known as *contrast normalization*) layer:

\begin{equation}\label{eq:inorm}
    y_{tijk} =  \frac{x_{tijk} - \mu_{ti}}{\sqrt{\sigma_{ti}^2 + \epsilon}},
    \quad
    \mu_{ti} = \frac{1}{HW}\sum_{l=1}^W \sum_{m=1}^H x_{tilm},
    \quad
    \sigma_{ti}^2 = \frac{1}{HW}\sum_{l=1}^W \sum_{m=1}^H (x_{tilm} - mu_{ti})^2.
\end{equation}

[Jump_to lesson 10 video](https://course.fast.ai/videos/?lesson=10&t=7114)

```python
class InstanceNorm(nn.Module):
    __constants__ = ['eps']
    def __init__(self, nf, eps=1e-0):
        super().__init__()
        self.eps = eps
        self.mults = nn.Parameter(torch.ones (nf,1,1))
        self.adds  = nn.Parameter(torch.zeros(nf,1,1))

    def forward(self, x):
        m = x.mean((2,3), keepdim=True)
        v = x.var ((2,3), keepdim=True)
        res = (x-m) / ((v+self.eps).sqrt())
        return res*self.mults + self.adds
```

```python
def conv_in(ni, nf, ks=3, stride=2, bn=True, **kwargs):
    layers = [nn.Conv2d(ni, nf, ks, padding=ks//2, stride=stride, bias=True),
              GeneralRelu(**kwargs)]
    if bn: layers.append(InstanceNorm(nf))
    return nn.Sequential(*layers)
```

```python
learn,run = get_learn_run(nfs, data, 0.1, conv_in, cbs=cbfs)
```

    Metrics:
    <function accuracy at 0x7f27aa6f57b8>
    Metrics:
    <function accuracy at 0x7f27aa6f57b8>

```python
%time run.fit(3, learn)
```

    train: [nan, tensor(0.0986, device='cuda:0')]
    valid: [nan, tensor(0.0991, device='cuda:0')]
    train: [nan, tensor(0.0986, device='cuda:0')]
    valid: [nan, tensor(0.0991, device='cuda:0')]
    train: [nan, tensor(0.0986, device='cuda:0')]
    valid: [nan, tensor(0.0991, device='cuda:0')]
    CPU times: user 16.5 s, sys: 4.32 s, total: 20.8 s
    Wall time: 20.8 s

*Question*: why can't this classify anything?

Lost in all those norms? The authors from the [group norm paper](https://arxiv.org/pdf/1803.08494.pdf) have you covered:

### Group norm

[Jump_to lesson 10 video](https://course.fast.ai/videos/?lesson=10&t=7213)

*From the PyTorch docs:*

`GroupNorm(num_groups, num_channels, eps=1e-5, affine=True)`

The input channels are separated into `num_groups` groups, each containing
``num_channels / num_groups`` channels. The mean and standard-deviation are calculated
separately over the each group. $\gamma$ and $\beta$ are learnable
per-channel affine transform parameter vectorss of size `num_channels` if
`affine` is ``True``.

This layer uses statistics computed from input data in both training and
evaluation modes.

Args:
-    num_groups (int): number of groups to separate the channels into
-    num_channels (int): number of channels expected in input
-    eps: a value added to the denominator for numerical stability. Default: 1e-5
-    affine: a boolean value that when set to ``True``, this module
        has learnable per-channel affine parameters initialized to ones (for weights)
        and zeros (for biases). Default: ``True``.

Shape:
- Input: `(N, num_channels, *)`
- Output: `(N, num_channels, *)` (same shape as input)

Examples::

    >>> input = torch.randn(20, 6, 10, 10)
    >>> # Separate 6 channels into 3 groups
    >>> m = nn.GroupNorm(3, 6)
    >>> # Separate 6 channels into 6 groups (equivalent with InstanceNorm)
    >>> m = nn.GroupNorm(6, 6)
    >>> # Put all 6 channels into a single group (equivalent with LayerNorm)
    >>> m = nn.GroupNorm(1, 6)
    >>> # Activating the module
    >>> output = m(input)

## Fix small batch sizes

### What's the problem?

When we compute the statistics (mean and std) for a BatchNorm Layer on a small batch, it is possible that we get a standard deviation very close to 0. because there aren't many samples (the variance of one thing is 0. since it's equal to its mean).

[Jump_to lesson 10 video](https://course.fast.ai/videos/?lesson=10&t=7304)

```python
data = DataBunch(*get_dls(train_ds, valid_ds, 2), c)
```

```python
def conv_layer(ni, nf, ks=3, stride=2, bn=True, **kwargs):
    layers = [nn.Conv2d(ni, nf, ks, padding=ks//2, stride=stride, bias=not bn),
              GeneralRelu(**kwargs)]
    if bn: layers.append(nn.BatchNorm2d(nf, eps=1e-5, momentum=0.1))
    return nn.Sequential(*layers)
```

```python
learn,run = get_learn_run(nfs, data, 0.4, conv_layer, cbs=cbfs)
```

    Metrics:
    <function accuracy at 0x7f27aa6f57b8>
    Metrics:
    <function accuracy at 0x7f27aa6f57b8>

```python
%time run.fit(1, learn)
```

    train: [2.32158203125, tensor(0.1925, device='cuda:0')]
    valid: [1.951426171875, tensor(0.2837, device='cuda:0')]
    CPU times: user 1min 30s, sys: 1.81 s, total: 1min 32s
    Wall time: 1min 32s

### Running Batch Norm

To solve this problem we introduce a Running BatchNorm that uses smoother running mean and variance for the mean and std.

[Jump_to lesson 10 video](https://course.fast.ai/videos/?lesson=10&t=7516)

```python
class RunningBatchNorm(nn.Module):
    def __init__(self, nf, mom=0.1, eps=1e-5):
        super().__init__()
        self.mom,self.eps = mom,eps
        self.mults = nn.Parameter(torch.ones (nf,1,1))
        self.adds = nn.Parameter(torch.zeros(nf,1,1))
        self.register_buffer('sums', torch.zeros(1,nf,1,1))
        self.register_buffer('sqrs', torch.zeros(1,nf,1,1))
        self.register_buffer('batch', tensor(0.))
        self.register_buffer('count', tensor(0.))
        self.register_buffer('step', tensor(0.))
        self.register_buffer('dbias', tensor(0.))

    def update_stats(self, x):
        bs,nc,*_ = x.shape
        self.sums.detach_()
        self.sqrs.detach_()
        dims = (0,2,3)
        s = x.sum(dims, keepdim=True)
        ss = (x*x).sum(dims, keepdim=True)
        c = self.count.new_tensor(x.numel()/nc)
        mom1 = 1 - (1-self.mom)/math.sqrt(bs-1)
        self.mom1 = self.dbias.new_tensor(mom1)
        self.sums.lerp_(s, self.mom1)
        self.sqrs.lerp_(ss, self.mom1)
        self.count.lerp_(c, self.mom1)
        self.dbias = self.dbias*(1-self.mom1) + self.mom1
        self.batch += bs
        self.step += 1

    def forward(self, x):
        if self.training: self.update_stats(x)
        sums = self.sums
        sqrs = self.sqrs
        c = self.count
        if self.step<100:
            sums = sums / self.dbias
            sqrs = sqrs / self.dbias
            c    = c    / self.dbias
        means = sums/c
        vars = (sqrs/c).sub_(means*means)
        if bool(self.batch < 20): vars.clamp_min_(0.01)
        x = (x-means).div_((vars.add_(self.eps)).sqrt())
        return x.mul_(self.mults).add_(self.adds)
```

```python
def conv_rbn(ni, nf, ks=3, stride=2, bn=True, **kwargs):
    layers = [nn.Conv2d(ni, nf, ks, padding=ks//2, stride=stride, bias=not bn),
              GeneralRelu(**kwargs)]
    if bn: layers.append(RunningBatchNorm(nf))
    return nn.Sequential(*layers)
```

```python
learn,run = get_learn_run(nfs, data, 0.4, conv_rbn, cbs=cbfs)
```

    Metrics:
    <function accuracy at 0x7f27aa6f57b8>
    Metrics:
    <function accuracy at 0x7f27aa6f57b8>

```python
%time run.fit(1, learn)
```

    train: [0.4153578125, tensor(0.8900, device='cuda:0')]
    valid: [0.2260486572265625, tensor(0.9460, device='cuda:0')]
    CPU times: user 5min 15s, sys: 26.9 s, total: 5min 42s
    Wall time: 5min 44s

This solves the small batch size issue!

### What can we do in a single epoch?

Now let's see with a decent batch size what result we can get.

[Jump_to lesson 10 video](https://course.fast.ai/videos/?lesson=10&t=8068)

```python
data = DataBunch(*get_dls(train_ds, valid_ds, 32), c)
```

```python
learn,run = get_learn_run(nfs, data, 0.9, conv_rbn, cbs=cbfs
                          +[partial(ParamScheduler,'lr', sched_lin(1., 0.2))])
```

    Metrics:
    <function accuracy at 0x7f27aa6f57b8>
    Metrics:
    <function accuracy at 0x7f27aa6f57b8>

```python
%time run.fit(1, learn)
```

    train: [0.162093759765625, tensor(0.9495, device='cuda:0')]
    valid: [0.16422022705078124, tensor(0.9780, device='cuda:0')]
    CPU times: user 15.7 s, sys: 612 ms, total: 16.3 s
    Wall time: 16.3 s

## Export

```python
nb_auto_export()
```

    <IPython.core.display.Javascript object>

```python

```


```python
# 그래프, 수학 기능 추가
# Add graph and math features
import pylab as py
import numpy as np
import numpy.linalg as nl
# 기호 연산 기능 추가
# Add symbolic operation capability
import sympy as sy

```

```python
sy.init_printing()

```

# 룽게-쿠타법 (RK4) : 고차 상미방
Runge-Kutta Method (RK4) : Higher Order ODE

## 단진자
Simple Pendulum

다음 미분 방정식은 단진자의 운동을 묘사한다.

Following differential equation describes the motion of a simple pendulum.

Ref : Wikipedia contributors, 'Pendulum (mathematics)', Wikipedia, The Free Encyclopedia, 2 June 2018, 13:28 UTC, <https://en.wikipedia.org/w/index.php?title=Pendulum_(mathematics)&oldid=844080803> [accessed 5 August 2018]

$$
\frac{d^2\theta}{dt^2} + \frac{g}{l}sin\theta = 0
$$

상태변수는 다음과 같다고 가정하자.

Let's assume that the state variables are as follows.

$$
\mathbf{x}
=
\begin{pmatrix}
x_0\\
x_1
\end{pmatrix}
=
\begin{pmatrix}
\theta\\
\frac{d}{dt}\theta
\end{pmatrix}
$$

상태변수의 미분은 다음과 같다.
Differentiation of the state variables are as follows.

$$
\frac{d}{dt}
\begin{pmatrix}
    x_0\\
    x_1
\end{pmatrix}
=
\begin{pmatrix}
    x_1\\
    -\frac{g}{l}sinx_0
\end{pmatrix}
$$

python 함수로 구현해 보면 다음과 같을 것이다.

One possible python implementation would be as follows.

```python
g_mpsps = 9.8
l_m = 0.3

legends = ('$\\theta(deg)$', '$\\frac{d}{dt}\\theta(deg/s)$')
ylabel = ''

# Initial state
x_0 = np.array([np.deg2rad(90), 0])

def pendulum_NL(t, x):
    """
    Parameters
    ==========
    x: array of theta and d(theta)/dt
    t: time value

    Return Value
    ============
    One dimensional array of dx/dt
    """

    return np.array([x[1], (-g_mpsps/l_m)*np.sin(x[0])])

```

룽게-쿠타법을 오일러법, 훈법과 비교해보자.
Let's compare the Runge-Kutta method with Euler method, and Heun's method.

```python
import ode_solver
import scipy.integrate as si

```

```python
py.figure(figsize=(12, 12))

# Time array
delta_t = 0.001

t_sec_array = np.arange(0, 6 + delta_t*0.5, delta_t)

# *** scipy ***
sol_pendulum = si.solve_ivp(pendulum_NL, (t_sec_array[0], t_sec_array[-1]), x_0, t_eval=t_sec_array)
py.plot(sol_pendulum.t, sol_pendulum.y.T, 'o', label='solve_ivp')

# *** Euler ***
t_euler_out, x_euler_out = ode_solver.euler(pendulum_NL, t_sec_array, x_0)
py.plot(t_euler_out, x_euler_out, '-', label='Euler')

# *** Heun ***
t_heun__out, x_heun__out = ode_solver.heun(pendulum_NL, t_sec_array, x_0)
py.plot(t_heun__out, x_heun__out, '-', label='Heun')

# *** RK4 ***
t_rk4___out, x_rk4___out = ode_solver.rk4(pendulum_NL, t_sec_array, x_0)
py.plot(t_rk4___out, x_rk4___out, '.', label='RK4')

py.xlabel('t(sec)')
py.ylabel('x(m)')

py.legend(loc=0)
py.grid(True)

```

## 4차 선형 상미방
Fourth Order Linear ODE

아래 4차 선형 상미분 방정식에 대해서도 여러 해법을 적용해 보자.

Let's apply various solvers to the following fourth order linear ODE.

$$
         \frac{d^4x}{dt^4}
         + 12 \frac{d^3x}{dt^3}
         + 54 \frac{d^2x}{dt^2}
         + 108 \frac{dx}{dt}
         + 80 x = 0
$$

$$
\mathbf{q} = \begin{pmatrix}q_0 & q_1 & q_2 & q_3 \end{pmatrix}^T = \begin{pmatrix}x & \frac{dx}{dt} & \frac{d^2x}{dt^2} & \frac{d^3x}{dt^3} \end{pmatrix}^T
$$

$$
\frac{d\mathbf{q}}{dt}
=
\begin{bmatrix}
0 & 1 & 0 & 0 \\
0 & 0 & 1 & 0 \\
0 & 0 & 0 & 1 \\
- 80 & - 108 & - 54 & -12
\end{bmatrix}
\begin{pmatrix}
q_0 \\ q_1 \\ q_2 \\ q_3
\end{pmatrix}
=
\mathbf{Aq}
$$

아래는 해당 4차 선형 상미방을 `dataclass`로 구현한 예이다. (Python 3.7 이후)

Following is a `dataclass` implementation of the fourth order LODE. (Python 3.7 or later)

```python
import dataclasses

@dataclasses.dataclass
class LODE(object):
    A : np.ndarray

    def slope(self, t, q):
        """
        Parameters
        ==========
        q: array of q_0, q_1, q_2, ..., q_n
        t: time value

        Return Value
        ============
        One dimensional array of dq/dt
        """

        q_column = np.array(q).T
        qdot_column = self.A @ q_column

        qdot_array = np.array(qdot_column.T).flatten()

        return qdot_array

lode4 = LODE(np.array([
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1],
        [-80, -108, -54, -12],
]))

ylabel = '$\mathbf{q}$'

# Initial state
x_0 = np.array([1, 0, 0, 0])

```

시간 배열을 정한다.

Define an array for time.

```python
delta_t = 0.08
t_sec_array = np.arange(0, 6 + delta_t*0.5, delta_t)

```

각 해법을 적용한다.

Apply each solver.

```python
# *** scipy ***
sol = si.solve_ivp(lode4.slope, (t_sec_array[0], t_sec_array[-1]), x_0, t_eval=t_sec_array)

# *** Euler ***
t_euler_out, x_euler_out_ode4 = ode_solver.euler(lode4.slope, t_sec_array, x_0)

# *** Heun ***
t_heun__out, x_heun__out_ode4 = ode_solver.heun(lode4.slope, t_sec_array, x_0)

# *** RK4 ***
t_rk4___out, x_rk4___out_ode4 = ode_solver.rk4(lode4.slope, t_sec_array, x_0)

```

$q_i$ 별로 그리기 위해 넘파이 배열로 변환한다.

To plot each $q_i$, convert to NumPy arrays.

```python
x_euler_out_ode4_array = np.array(x_euler_out_ode4)
x_heun__out_ode4_array = np.array(x_heun__out_ode4)
x_rk4___out_ode4_array = np.array(x_rk4___out_ode4)

```

```python
py.figure(figsize=(10, 10))

for i_state in range(x_euler_out_ode4_array.shape[1]):

    py.subplot(x_euler_out_ode4_array.shape[1], 1, i_state+1)
    py.plot(sol.t, sol.y[i_state, :], 'o', label='solve_ivp')
    py.plot(t_euler_out, x_euler_out_ode4_array[:, i_state], '-', label='Euler')
    py.plot(t_heun__out, x_heun__out_ode4_array[:, i_state], '-', label='Heun')
    py.plot(t_rk4___out, x_rk4___out_ode4_array[:, i_state], '.', label='RK4')
    py.xlabel('t(sec)')
    py.ylabel(f'$q_{i_state}$')
    py.legend(loc=1)
    py.grid(True)

```

## 연습 문제
Exercises

도전 과제 1: 다음 2계 선형 상미분 방정식의 수치해를 RK4법으로 구하시오:

Try This 1: Find the numerical solutions of the following second order linear ordinary equation using RK4 Method:

$$
\begin{align}
\frac{d^2}{dt^2}x(t) + 2\frac{d}{dt}x(t) + x(t) &= 0 \\
x(0) &= 0 \\
\frac{d}{dt}x(0) &= 1
\end{align}
$$

```python

```

도전 과제 2: 다음 2계 선형 상미분 방정식의 수치해를 RK4법으로 구하시오:

Try This 2: Find the numerical solutions of the following second order linear ordinary equation using RK4 Method:

$$
\begin{align}
\frac{d^2}{dt^2}x(t) + \frac{d}{dt}x(t) + 4x(t) &= sin(t[rad]) \\
x(0) &= 0 \\
\frac{d}{dt}x(0) &= 0
\end{align}
$$

```python

```

도전 과제 3: 다음 연립 상미분 방정식의 수치해를 RK4법으로 구하시오:

Try This 3: Find the numerical solutions of the following system of ordinary equations using RK4 Method:

$$
\begin{cases}
    \begin{align}
        \frac{d}{dt}x(t) &= 2 x(t) - y(t) \\
        \frac{d}{dt}y(t) &= 2.1 y(t) - x(t)
    \end{align}
\end{cases}\\
x(0) = y(0) = 1.0 \\
0 \le t \le 2
$$

ref : Zill & Cullen

```python

```

도전 과제 4: 다음 연립 상미분 방정식의 수치해를 RK4법으로 구하시오:

Try This 4: Find the numerical solutions of the following system of ordinary equations using RK4 Method:

$$
\begin{cases}
    \begin{align}
        \frac{d}{dt}x(t) + \frac{d}{dt}y(t) + 2y(t)&= 0 \\
        \frac{d}{dt}x(t) - 3x(t) - 2 y(t) &= 0
    \end{align}
\end{cases}\\
x(0) = y(0) = 1.0 \\
0 \le t \le 2
$$

ref : Zill & Cullen

```python

```

## 부록: 단진자의 방향장
Appendix: The direction field of a simple pendulum

각 상태에서 상태변수의 변화의 방향 $\left(\frac{d}{dt}\theta, \frac{d^2}{dt^2}\theta \right)$ 을 표시해 보자.
At each state, let's present the direction of state variable change $\left(\frac{d}{dt}\theta, \frac{d^2}{dt^2}\theta \right)$.

```python
import ode_plot

```

```python
time_list = []

# list of theta
theta_deg_array = np.arange(-540, 540+1, 30)
theta_rad_list = np.deg2rad(theta_deg_array)

# list of theta_dot
theta_dot_deg_array = np.arange(-540, 540+1, 45)
theta_dot_rad_list = np.deg2rad(theta_dot_deg_array)

# plot the direction filed
ode_plot.ode_slopes_2states_cartesian(pendulum_NL, theta_rad_list, theta_dot_rad_list, time_list)

# Convert pendulum solution curves to NumPy arrays
x_euler_out_array, x_heun__out_array, x_rk4___out_array = (
    np.array(x_euler_out), np.array(x_heun__out), np.array(x_rk4___out))

# Plot the solution curves
py.plot(py.rad2deg(x_euler_out_array[:, 0]), py.rad2deg(x_euler_out_array[:, 1]), '-', label='Euler', alpha=0.7)
py.plot(py.rad2deg(x_heun__out_array[:, 0]), py.rad2deg(x_heun__out_array[:, 1]), '.-', label='Heun', alpha=0.7)
py.plot(py.rad2deg(x_rk4___out_array[:, 0]), py.rad2deg(x_rk4___out_array[:, 1]), '-', label='RK4')

ode_plot.set_axis(
    py.gca(),
    theta_deg_array[0], theta_deg_array[-1],
    theta_dot_deg_array[0], theta_dot_deg_array[-1],
)

py.legend(loc=0)
ode_plot.title_axis_labels()

py.savefig('pendulum_direction_field.svg')

```

## Final Bell
마지막 종

```python
# stackoverfow.com/a/24634221
import os
os.system("printf '\a'");

```

```python

```


# Corrugated Shells geometry

## Init symbols for *sympy*

```python
from sympy import *
from geom_util import *
from sympy.vector import CoordSys3D
import matplotlib.pyplot as plt
import sys
sys.path.append("../")

%matplotlib inline

%reload_ext autoreload
%autoreload 2
%aimport geom_util
```

```python
# Any tweaks that normally go in .matplotlibrc, etc., should explicitly go here
%config InlineBackend.figure_format='retina'
plt.rcParams['figure.figsize'] = (12, 12)

plt.rc('text', usetex=True)

plt.rc('font', family='serif')
# SMALL_SIZE = 42
# MEDIUM_SIZE = 42
# BIGGER_SIZE = 42

# plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
# plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title
# plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
# plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
# plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
# plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize
# plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title

init_printing()
```

```python
N = CoordSys3D('N')
alpha1, alpha2, alpha3 = symbols("alpha_1 alpha_2 alpha_3", real = True, positive=True)
```

## Cylindrical coordinates

```python
R, L, ga, gv = symbols("R L g_a g_v", real = True, positive=True)
```

```python
a1 = pi / 2 + (L / 2 - alpha1)/R

a2 = 2 * pi * alpha1 / L

x1 = (R + ga * cos(gv * a2)) * cos(a1)
x2 = alpha2
x3 = (R + ga * cos(gv * a2)) * sin(a1)

r = x1*N.i + x2*N.j + x3*N.k

r1=r.diff(alpha1)
r1
```

```python
A = r1.dot(r1)
trigsimp(A)
```

```python
z = 2*ga*gv*pi/L*sin(gv*a2)
w = 1 + ga/R*cos(gv*a2)

dr1x=w*sin(a1) - z*cos(a1)
dr1z=-w*cos(a1) - z*sin(a1)

# r1 = dr1x*N.i + dr1z*N.k
# r2 =N.j

mag=sqrt((w)**2+(z)**2)

nx = -dr1z/mag
nz = dr1x/mag

n = nx*N.i+nz*N.k

dnx=nx.diff(alpha1)
dnz=nz.diff(alpha1)

dn= dnx*N.i+dnz*N.k

n
```

```python
Ralpha = r+alpha3*n

R1=r1+alpha3*dn
R2=Ralpha.diff(alpha2)
R3=n

Ralpha
```

```python
R1
```

```python
R2
```

```python
R3
```

### Draw

```python
import plot

%aimport plot

# x1 = (R + alpha3 + ga * cos(gv * a2)) * cos(a1)
# x2 = alpha2
# x3 = (R + alpha3 + ga * cos(gv * a2)) * sin(a1)

x1 = Ralpha.dot(N.i)
x3 = Ralpha.dot(N.k)

alpha1_x = lambdify([R, L, ga, gv, alpha1, alpha3], x1, "numpy")
alpha3_z = lambdify([R, L, ga, gv, alpha1, alpha3], x3, "numpy")

R_num = 1/0.8
L_num = 2
h_num = 0.1
ga_num = 0.1
gv_num = 20

x1_start = 0
x1_end = L_num
x3_start = -h_num/2
x3_end = h_num/2

def alpha_to_x(a1, a2, a3):
    x=alpha1_x(R_num, L_num, ga_num, gv_num, a1, a3)
    z=alpha3_z(R_num, L_num, ga_num, gv_num, a1, a3)
    return x, 0, z

plot.plot_init_geometry_2(x1_start, x1_end, x3_start, x3_end, alpha_to_x)
```

```python
%aimport plot

R3_1=R3.dot(N.i)
R3_3=R3.dot(N.k)

R3_1_x = lambdify([R, L, ga, gv, alpha1, alpha3], R3_1, "numpy")
R3_3_z = lambdify([R, L, ga, gv, alpha1, alpha3], R3_3, "numpy")

def R3_to_x(a1, a2, a3):
    x=R3_1_x(R_num, L_num, ga_num, gv_num, a1, a3)
    z=R3_3_z(R_num, L_num, ga_num, gv_num, a1, a3)
    return x, 0, z

plot.plot_vectors(x1_start, x1_end, 0, alpha_to_x, R3_to_x)
```

```python
%aimport plot

R1_1=r1.dot(N.i)
R1_3=r1.dot(N.k)

R1_1_x = lambdify([R, L, ga, gv, alpha1, alpha3], R1_1, "numpy")
R1_3_z = lambdify([R, L, ga, gv, alpha1, alpha3], R1_3, "numpy")

def R1_to_x(a1, a2, a3):
    x=R1_1_x(R_num, L_num, ga_num, gv_num, a1, a3)
    z=R1_3_z(R_num, L_num, ga_num, gv_num, a1, a3)
    return x, 0, z

plot.plot_vectors(x1_start, x1_end, 0, alpha_to_x, R1_to_x)
```

### Lame params

```python
A=

H1 = sqrt((alpha3*((-(1 + ga*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R)*sin((L/2 - alpha1)/R) - ga*gv*sin(gv*(pi/2 + (L/2 - alpha1)/R))*cos((L/2 - alpha1)/R)/R)*(-ga*gv*(1 + ga*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R)*sin(gv*(pi/2 + (L/2 - alpha1)/R))/R**2 + ga**2*gv**3*sin(gv*(pi/2 + (L/2 - alpha1)/R))*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R**3)/((1 + ga*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R)**2 + ga**2*gv**2*sin(gv*(pi/2 + (L/2 - alpha1)/R))**2/R**2)**(3/2) + ((1 + ga*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R)*cos((L/2 - alpha1)/R)/R + ga*gv**2*cos((L/2 - alpha1)/R)*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R**2 - 2*ga*gv*sin((L/2 - alpha1)/R)*sin(gv*(pi/2 + (L/2 - alpha1)/R))/R**2)/sqrt((1 + ga*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R)**2 + ga**2*gv**2*sin(gv*(pi/2 + (L/2 - alpha1)/R))**2/R**2)) + (1 + ga*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R)*cos((L/2 - alpha1)/R) - ga*gv*sin((L/2 - alpha1)/R)*sin(gv*(pi/2 + (L/2 - alpha1)/R))/R)**2 + (alpha3*(((1 + ga*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R)*cos((L/2 - alpha1)/R) - ga*gv*sin((L/2 - alpha1)/R)*sin(gv*(pi/2 + (L/2 - alpha1)/R))/R)*(-ga*gv*(1 + ga*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R)*sin(gv*(pi/2 + (L/2 - alpha1)/R))/R**2 + ga**2*gv**3*sin(gv*(pi/2 + (L/2 - alpha1)/R))*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R**3)/((1 + ga*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R)**2 + ga**2*gv**2*sin(gv*(pi/2 + (L/2 - alpha1)/R))**2/R**2)**(3/2) + ((1 + ga*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R)*sin((L/2 - alpha1)/R)/R + ga*gv**2*sin((L/2 - alpha1)/R)*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R**2 + 2*ga*gv*sin(gv*(pi/2 + (L/2 - alpha1)/R))*cos((L/2 - alpha1)/R)/R**2)/sqrt((1 + ga*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R)**2 + ga**2*gv**2*sin(gv*(pi/2 + (L/2 - alpha1)/R))**2/R**2)) + (1 + ga*cos(gv*(pi/2 + (L/2 - alpha1)/R))/R)*sin((L/2 - alpha1)/R) + ga*gv*sin(gv*(pi/2 + (L/2 - alpha1)/R))*cos((L/2 - alpha1)/R)/R)**2)
H2=S(1)
H3=S(1)

H=[H1, H2, H3]
DIM=3
dH = zeros(DIM,DIM)
for i in range(DIM):
    dH[i,0]=H[i].diff(alpha1)
    dH[i,1]=H[i].diff(alpha2)
    dH[i,2]=H[i].diff(alpha3)

trigsimp(H1)

```