# Large deformations and images

```python
import numpy as np
%matplotlib notebook
import matplotlib as mpl
import matplotlib.pyplot as plt
import nibabel as nib
import os
import summerschool as ss
import scipy.interpolate as spi
```

```python
# load a pair of images
fname = os.path.join('mouse_images','PMD2052_orig_target_STS_clean.img')
imgI = nib.load(fname)
nx = imgI.header['dim'][1:4]
dx = imgI.header['pixdim'][1:4]
fname = os.path.join('mouse_images','PMD3097_orig_target_STS_clean.img')
imgJ = nib.load(fname)
```

```python
# linearly transform the first to match the second
# load affine from linear tutorial and apply it as in last tutorial
A = np.zeros((4,4))
with open('affine.txt','rt') as f:
    for i in range(4):
        line = f.readline()
        A[i,:] = [float(a) for a in line.split()]

x0 = np.arange(nx[0])*dx[0]
x1 = np.arange(nx[1])*dx[1]
x2 = np.arange(nx[2])*dx[2]
x0 = x0 - np.mean(x0)
x1 = x1 - np.mean(x1)
x2 = x2 - np.mean(x2)
X0,X1,X2 = np.meshgrid(x0,x1,x2,indexing='ij')
X0s,X1s,X2s = ss.sample_points_from_affine(X0,X1,X2,A)
interp_args = {
    'method':'linear',
    'bounds_error':False,
    'fill_value':0
    }
AI = spi.interpn([x0,x1,x2],imgI.get_data()[:,:,:,0],np.stack([X0s,X1s,X2s],axis=-1),**interp_args)
f,ax = plt.subplots(2,3)
ss.imshow_slices(x0,x1,x2,imgI.get_data()[:,:,:,0],ax[0])
ss.imshow_slices(x0,x1,x2,AI,ax[1])
```

    <IPython.core.display.Javascript object>

    [<matplotlib.image.AxesImage at 0x7f952b1ca5c0>,
     <matplotlib.image.AxesImage at 0x7f952b1ca898>,
     <matplotlib.image.AxesImage at 0x7f952b1cab70>]

```python
# start with a 2D example
# set up a domain
x0 = np.arange(nx[0])*dx[0] - np.mean(np.arange(nx[0])*dx[0])
x1 = np.arange(nx[1])*dx[1] - np.mean(np.arange(nx[1])*dx[1])
X0,X1 = np.meshgrid(x0,x1,indexing='ij')
I = AI[:,:,imgI.shape[2]//2]
J = imgJ.get_data()[:,:,imgJ.shape[2]//2,0] # note last axes is time, we'd like to remove it

f,ax = plt.subplots(1,2)
ax[0].imshow(I,cmap='gray')
ax[0].set_aspect('equal')
ax[1].imshow(J,cmap='gray')
ax[1].set_aspect('equal')
```

    <IPython.core.display.Javascript object>

## The failure of smooth displacement vector fields
Smooth displacement fields generate invertible transformations as long as they are small enough.

But as they grow bigger, they can become lose invertability no matter how smooth they are

```python
# example deformation and action on image
gaussian = np.exp(-(X0**2 + X1**2)/(2.0*(10.0*dx[0])**2))
magnitudes = np.array([10,20,30,40])*dx[0]

f,ax = plt.subplots(len(magnitudes),1,sharex=True,sharey=True)
plt.subplots_adjust(hspace=0.5)

for magnitude,a in zip(magnitudes,ax):
    v0 = gaussian*magnitude
    v1 = np.zeros_like(v0)
    phi0 = X0+v0
    phi1 = X1+v1
    phiinv0 = X0-v0
    phiinv1 = X1-v1
    interp_args = {'method':'linear','bounds_error':False,'fill_value':0}
    phiI = spi.interpn([x0,x1],I,np.stack([phiinv0,phiinv1],axis=-1),**interp_args)

    stride = 10
    imshow_args = dict(cmap='gray',extent=[x1[0],x1[-1],x0[0],x0[-1]],origin='lower')
    a.imshow(phiI,**imshow_args)
    ss.plot_grid(phi1,phi0,ax=a,rstride=stride,cstride=stride,color='r',linewidth=0.5,alpha=0.75)
    a.set_aspect('equal')
    a.set_title('mag: {:5.2f}'.format(magnitude))

# note that in this visualization the grid shows the deformation "phi",
# but the image is deformed by composing with "phi inverse"
# In this case, "phi inverse" doesn't exist, so what we've visualizing is an approximation
```

    <IPython.core.display.Javascript object>

## Combining small deformations under composition works
If any two transformations are invertable, than so is their composition.  We can build up large transformations by composing small ones.

```python
f,ax = plt.subplots(len(magnitudes),1,sharex=True,sharey=True)

mag = magnitudes[0]
interp_args = dict(bounds_error=False,fill_value=0,method='linear')
stride = 10
imshow_args = dict(cmap='gray',extent=[x1[0],x1[-1],x0[0],x0[-1]],origin='lower')

phi0 = np.array(X0)
phi1 = np.array(X1)
phiinv0 = np.array(X0)
phiinv1 = np.array(X1)
for a in ax:
    v0 = gaussian*mag
    v1 = np.zeros_like(v0)
    # update using composition
    # recall that on a computer, composition is just interpolation
    phi0_ = phi0 + spi.interpn([x0,x1],v0,np.stack([phi0,phi1],axis=-1),**interp_args)
    phi1 = phi1 + spi.interpn([x0,x1],v1,np.stack([phi0,phi1],axis=-1),**interp_args)
    phi0 = phi0_

    # update to the inverse
    phiinv0 = spi.interpn([x0,x1],phiinv0,np.stack([X0-v0,X1-v1],axis=-1),**interp_args)
    phiinv1 = spi.interpn([x0,x1],phiinv1,np.stack([X0-v0,X1-v1],axis=-1),**interp_args)

    phiI = spi.interpn([x0,x1],I,np.stack([phiinv0,phiinv1],axis=-1),**interp_args)

    a.imshow(phiI,**imshow_args)
    ss.plot_grid(phi1,phi0,ax=a,rstride=stride,cstride=stride,color='r',linewidth=0.5,alpha=0.5)
    a.set_aspect('equal')

```

## Flows arise in the limit of composition

Let $\varphi_{t_0} = id$, then $\varphi_{t_1} = (id+v_{t_0}\Delta t)\circ \varphi_{t_0}$, and $\varphi_{t_{n+1}} = (id+v_{t_n}\Delta t)\circ \varphi_{t_n}$.

This can be written as $\varphi_{t_{n+1}} = \varphi_{t_n} + \Delta T v_{t_n}(\varphi_{t_n})$.  Rearranging gives
\begin{align*}
\frac{\varphi_{t_{n+1}} - \varphi_{t_n}}{\Delta t} = v_{t_n}(\varphi_{t_n})
\end{align*}

Passing to the limit gives us Euler's equation
\begin{align*}
\frac{d}{dt} \varphi_t &= v_t(\varphi_t)
\end{align*}
which describes an the flow of the ODE $\dot x = v_t(x)$.

## The inverse deforms under optical flow

Since images deform by composing with $\varphi^{-1}$, we need to understand how $\varphi^{-1}$ flows over time under $v$.

Differentiate the identity transformation to get
\begin{align*}
0 &=\frac{d}{dt}id\\
&= \frac{d}{dt}(\varphi_t(\varphi_t^{-1}))\\
&=\frac{d}{dt}\varphi_t(\varphi_t^{-1}) + D\varphi_t(\varphi_t^{-1})\frac{d}{dt}\varphi_t^{-1}\\
&=v_t + D\varphi_t(\varphi_t^{-1})\frac{d}{dt}\varphi_t^{-1}
\end{align*}

We plugged in Euler's equation for the first term.  Note that we can understand the second term by applying the chain rule
\begin{align}
id &= D(x)\\
&= D(\varphi(\varphi^{-1}))\\
&= D\varphi(\varphi^{-1})D\varphi^{-1}
\end{align}
which implies that $[D\varphi(\varphi^{-1})]^{-1} = D\varphi^{-1}$.  We can use this to solve the above equation
\begin{align*}
\frac{d}{dt}\varphi_t^{-1} &= -D\varphi_t^{-1}v_t
\end{align*}
This is just the optical flow equation https://en.wikipedia.org/wiki/Optical_flow on each component of $\varphi^{-1}$.

This equation is generally solved numerically using the method of characteristics https://en.wikipedia.org/wiki/Method_of_characteristics which we will employ here.

## The image deforms under optical flow
Using the above formula for $\frac{d}{dt}\varphi^{-1}_t$, we can find a formula for $\frac{d}{dt}I_t$.
\begin{align*}
\frac{d}{dt} I_t &= \frac{d}{dt}I_0(\varphi^{-1}_t)\\
&= D I_0 (\varphi^{-1}_t)\frac{d}{dt} \varphi^{-1}_t\\
&\stackrel{(a)}{=} -D I_0 (\varphi^{-1}_t) D\varphi^{-1}_t v_t\\
&\stackrel{(b)}{=} -D[I_0 (\varphi^{-1}_t)] v_t\\
&= -D I_t v_t
\end{align*}
where in (a) we applied the optical flow equation derived above, and in (b) we simply group terms using the chain rule.

```python
# example of flows
nt = 5
dt = 1.0/nt

# encode our energy operator in the Fourier domain using discrete derivatives
f0 = np.arange(nx[0])/(dx[0]*nx[0])
f1 = np.arange(nx[1])/(dx[1]*nx[1])
F0,F1 = np.meshgrid(f0,f1,indexing='ij')
a = 5.0*dx[0] # a scale in mm
p = 2
Ahat = (1.0 - 2.0*a**2*((np.cos(2.0*np.pi*dx[0]*F0) - 1.0)/dx[0]**2 + (np.cos(2.0*np.pi*dx[1]*F1) - 1.0)/dx[1]**2))**(2.0*p)
Khat = 1.0/Ahat
K = np.fft.ifft2(Khat,nx[:2]).real

# make a random flow
mag = 500.0*dx[0]

vt0 = np.random.randn(nx[0],nx[1],nt)*mag
vt1 = np.random.randn(nx[0],nx[1],nt)*mag
# smooth it
vt0 = np.fft.ifftn(np.fft.fftn(vt0,axes=[0,1])*Khat[:,:,None],axes=[0,1]).real
vt1 = np.fft.ifftn(np.fft.fftn(vt1,axes=[0,1])*Khat[:,:,None],axes=[0,1]).real

# now lets flow it forward
phi0 = np.array(X0)
phi1 = np.array(X1)
phiinv0 = np.array(X0)
phiinv1 = np.array(X1)

f,ax = plt.subplots(1,2)
for t in range(nt):
    v0 = vt0[:,:,t]
    v1 = vt1[:,:,t]
    # update phi using Euler's equation
    phi0_ = phi0 + spi.interpn([x0,x1],v0,np.stack([phi0,phi1],axis=-1),**interp_args)*dt
    phi1 = phi1 + spi.interpn([x0,x1],v1,np.stack([phi0,phi1],axis=-1),**interp_args)*dt
    phi0 = phi0_

    # update phiinv using method of characteristics
    X0s = X0-v0*dt
    X1s = X1-v1*dt
    interp_args_phiinv = dict(interp_args)
    interp_args_phiinv['fill_value'] = None
    #phiinv0 = spi.interpn([x0,x1],phiinv0,np.stack([X0s,X1s],axis=-1),**interp_args_phiinv)
    #phiinv1 = spi.interpn([x0,x1],phiinv1,np.stack([X0s,X1s],axis=-1),**interp_args_phiinv)
    # here we need to deal with boundary conditions better
    phiinv0 = spi.interpn([x0,x1],phiinv0-X0,np.stack([X0s,X1s],axis=-1),**interp_args_phiinv)+X0s
    phiinv1 = spi.interpn([x0,x1],phiinv1-X1,np.stack([X0s,X1s],axis=-1),**interp_args_phiinv)+X1s

    # act on the image
    phiI = spi.interpn([x0,x1],I,np.stack([phiinv0,phiinv1],axis=-1),**interp_args)

    [a.cla() for a in ax]
    ax[0].imshow(phiI,**imshow_args)
    ss.plot_grid(phi1,phi0,ax=ax[0],rstride=stride,cstride=stride,color='r',linewidth=0.5,alpha=0.5)
    ax[0].set_aspect('equal')
    ax[0].set_title('t = {:3.1f}'.format((t+1)*dt))

    ax[1].quiver(X1[::stride,::stride],X0[::stride,::stride],
             v1[::stride,::stride],v0[::stride,::stride])
    ax[1].set_aspect('equal')
    f.canvas.draw()
```

## Finding the optimal flow
We write an optimization problem like
\begin{align*}
E &= \frac{1}{2}\int_0^1 \int_X v^T(x)Av(x) dx dt + \frac{1}{2\sigma^2}\int |I(\varphi_1^{-1}(x)) - J(x)|^2dx
\end{align*}
where the the diffeomorphism is linked to the velocity by Euler's equation $\frac{d}{dt}\varphi_t = v_t(\varphi_t)$, and $A$ is our energy operator from before.

One of the nicer approaches to solving this problem comes from optimal control.  We'll define the state of our system as the image $I$, which is governed by dynamics $\dot I = -DI v$ (optical flow) with $I_0 = I_{\text{template}}$, where $v$ is the control.

These dynamics are enforced by a Lagrange multiplier, or costate, $\lambda$.  We can write an augmented cost function $L$ as
\begin{align*}
L = E + \int_0^1 \int_X \lambda_t(x)(\frac{d}{dt}I_t(x) + DI_t(x) v_t(x))dxdt
\end{align*}
We can find minimizers by taking perturbations with respect to $I$ and $v$.

First take a perturbation with respect to $v$
\begin{align}
&\frac{d}{d\epsilon}  \frac{1}{2}\int_0^1 \int_X (v_t + \epsilon \delta v_t)^T(x)A(v_t + \epsilon \delta v_t)(x) dx dt + \frac{1}{2\sigma^2}\int_X |I_1(x) - J(x)|^2dx + \int_0^1 \int_X \lambda_t(x)(\frac{d}{dt}I_t(x) + DI_t(x) (v_t + \epsilon \delta v_t)(x))dxdt \bigg|_{\epsilon = 0}\\
&=  \int_0^1 \int_X Av_t^t(x) \delta v(x) + \int_0^1\int_X \lambda_t(x) DI_t(x) \delta v_xdx dt
\end{align}

If this is true for all perturbations, we must have
\begin{align}
d_vE &= Av_t + \lambda_t DI_t
\end{align}
for numerical reasons we usually work with the gradient vector by acting with $K$, the Green's kernel of $A$
\begin{align}
\nabla_v E &= v_t + K*\lambda_t DI_t
\end{align}

Now we'll extremize over $I$
\begin{align*}
&\frac{d}{d \epsilon}\frac{1}{2}\int_0^1 \int_X v_t^T(x)Av_t(x) dx dt + \frac{1}{2\sigma^2}\int_X |I_1(x) + \epsilon \delta I_1(x) - J(x)|^2dx + \int_0^1 \int_X \lambda_t(x)(\frac{d}{dt}(I_t + \epsilon \delta I_t)(x)  + D(I_t + \epsilon \delta I_t)(x) v_t(x))dxdt \bigg|_{\epsilon = 0}\\
&= \frac{1}{\sigma^2}\int_X (I_1(x)  - J(x))\delta I_1(x) dx + \int_0^1\int_X \lambda_t(x) \frac{d}{dt}\delta I_t(x) dxdt + \int_0^1 \int_X \lambda_t(x) D \delta I_t(x) v_t(x)dxdt
\end{align*}
We'll use integration by parts on the middle term, remembering that $\delta I_0 = 0$, and we'll use the divergence theorem on the final term.
\begin{align*}
\frac{1}{\sigma^2}\int_X (I_1(x)  - J(x))\delta I_1(x) dx + -\int_0^1\int_X \frac{d}{dt}\lambda_t(x) \delta I_t(x) dxdt + \int_X \lambda_1(x) \delta I_1(x)dx - \int_0^1 \int_X \text{div}[\lambda_t(x) v_t(x)] \delta I_t(x)dxdt
\end{align*}

For this to vanish for all perturbations the we must have have the boundary conditions $\lambda_1 = -\frac{1}{\sigma^2}(I_1 - J)$, and the dynamical equation
\begin{align*}
\frac{d}{dt}\lambda = \text{div}[\lambda v]
\end{align*}
This equation has an explicit solution in terms of the diffeomorphism.  If $\varphi_{1t}$ is the diffeomorphism generated from integrating $v$ from time 1 backwards to time $t$, then the solution is
\begin{align}
\lambda_t &= \lambda_1(\varphi_{1t}^{-1})|D\varphi_{1t}^{-1}|
\end{align}

Putting this all together gives a nice equation for the gradient that we can use for an optimization algorithm
\begin{align*}
\nabla_v E &= v -K*(I(\varphi_t^{-1}) - J(\varphi_{1t}^{-1})) \nabla I_t |D\varphi_{1t}^{-1}|
\end{align*}

```python
# now run matching
# for display, create plots with shared axis
f,ax = plt.subplots(2,2,sharex=True,sharey=True)
# for plotting energy, no shared axis
f.delaxes(ax[1][0])
ax[1][0] = f.add_subplot(2,2,3)

# parameters
sigma = 2e0
niter = 50
epsilon = 2.0e-2
# mouse
# for mouse
sigma = 2.0e1
epsilon = 5.0e-3

# initialize
vt0 = np.zeros((nx[0],nx[1],nt))
vt1 = np.zeros((nx[0],nx[1],nt))
It = np.zeros((nx[0],nx[1],nt+1))
It[:,:,0] = I

EMAll = []
ERAll = []
EAll = []
for it in range(niter):
    # compute the deformed image by flowing over time
    phi0 = X0
    phi1 = X1
    phiinv0 = X0
    phiinv1 = X1
    ER = 0.0
    for t in range(nt):
        v0 = vt0[:,:,t]
        v1 = vt1[:,:,t]
        # update phi using Euler's equation
        # this is not necessary for the algorithm, only visualization
        phi0_ = phi0 + spi.interpn([x0,x1],v0,np.stack([phi0,phi1],axis=-1),**interp_args)*dt
        phi1 = phi1 + spi.interpn([x0,x1],v1,np.stack([phi0,phi1],axis=-1),**interp_args)*dt
        phi0 = phi0_

        # update phiinv using method of characteristics
        X0s = X0-v0*dt
        X1s = X1-v1*dt
        #phiinv0 = spi.interpn([x0,x1],phiinv0,np.stack([X0s,X1s],axis=-1),**interp_args_phiinv)
        #phiinv1 = spi.interpn([x0,x1],phiinv1,np.stack([X0s,X1s],axis=-1),**interp_args_phiinv)
        # here we need to deal with boundary conditions better
        phiinv0 = spi.interpn([x0,x1],phiinv0-X0,np.stack([X0s,X1s],axis=-1),**interp_args_phiinv)+X0s
        phiinv1 = spi.interpn([x0,x1],phiinv1-X1,np.stack([X0s,X1s],axis=-1),**interp_args_phiinv)+X1s

        # act on the image
        It[:,:,t+1] = spi.interpn([x0,x1],I,np.stack([phiinv0,phiinv1],axis=-1),**interp_args)

        # calculate the energy of the flow
        Av0 = np.fft.ifftn(np.fft.fftn(v0)*Ahat).real
        Av1 = np.fft.ifftn(np.fft.fftn(v1)*Ahat).real
        ER += np.sum(v0*Av0 + v1*Av1)*0.5*dx[0]*dx[1]*dt

    ax[0,0].cla()
    ax[0,0].imshow(It[:,:,-1],**imshow_args)
    ss.plot_grid(phi1,phi0,ax=ax[0,0],color='r',alpha=0.5,linewidth=0.5,rstride=stride,cstride=stride)

    # error
    err = It[:,:,-1] - J
    lambda1 = -(err)/sigma**2
    ax[0,1].cla()
    ax[0,1].imshow(err,**imshow_args)

    EM = np.sum(err**2/(2.0*sigma**2))*dx[0]*dx[1]
    EMAll.append(EM)
    ERAll.append(ER)
    E = ER+EM
    EAll.append(E)
    ax[1,0].cla()
    ax[1][0].plot(EAll)
    ax[1][0].plot(EMAll)
    ax[1][0].plot(ERAll)

    # now we start flowing the error back in time
    phiinv0 = X0
    phiinv1 = X1
    for t in range(nt-1,-1,-1):
        v0 = vt0[:,:,t]
        v1 = vt1[:,:,t]

        # update phiinv using method of characteristics
        X0s = X0+v0*dt
        X1s = X1+v1*dt
        #phiinv0 = spi.interpn([x0,x1],phiinv0,np.stack([X0s,X1s],axis=-1),**interp_args_phiinv)
        #phiinv1 = spi.interpn([x0,x1],phiinv1,np.stack([X0s,X1s],axis=-1),**interp_args_phiinv)
        # here we need to deal with boundary conditions better
        phiinv0 = spi.interpn([x0,x1],phiinv0-X0,np.stack([X0s,X1s],axis=-1),**interp_args_phiinv)+X0s
        phiinv1 = spi.interpn([x0,x1],phiinv1-X1,np.stack([X0s,X1s],axis=-1),**interp_args_phiinv)+X1s

        # find the determinant of Jacobian
        phiinv0_0,phiinv0_1 = np.gradient(phiinv0,dx[0],dx[1])
        phiinv1_0,phiinv1_1 = np.gradient(phiinv1,dx[0],dx[1])
        detjac = phiinv0_0*phiinv1_1 - phiinv0_1*phiinv1_0

        # find lambda_t
        lambdat = spi.interpn([x0,x1],lambda1,np.stack([phiinv0,phiinv1],axis=-1),**interp_args)*detjac

        # get the gradient of the image at this time
        I_0,I_1 = np.gradient(It[:,:,t],dx[0],dx[1])

        # initialize the gradient with the matching term
        grad0 = lambdat*I_0
        grad1 = lambdat*I_1

        # smooth it
        grad0 = np.fft.ifftn(np.fft.fftn(grad0)*Khat).real
        grad1 = np.fft.ifftn(np.fft.fftn(grad1)*Khat).real

        # add the regularization term
        grad0 += v0
        grad1 += v1

        # update
        vt0[:,:,t] -= epsilon*grad0
        vt1[:,:,t] -= epsilon*grad1

    f.canvas.draw()
```

## LDDMM Image Matching in 3D
We'll run the same code in 3D, using a downsampled image so it doesn't crash anyone's laptop.  The code is exactly the same, except for display and a slightly more complicated definition of the determinant.

This code should run with quite high performance, provided that numpy has been installed appropriately and can use multiple processors.  Please see this guide https://roman-kh.github.io/numpy-multicore/ .

You are welcome to use this code in your research.  Please give attribution to the original source on GitHub and author (Daniel Tward).  A python module is also provided for running LDDMM image matching outside of this jupyter notebook.

```python
# load images
I = AI
down = 3
I = ss.downsample_image(I,down)
nx = I.shape
dx = imgI.header['pixdim'][1:4]*down
J = imgJ.get_data()[:,:,:,0]
J = ss.downsample_image(J,down)
x0 = np.arange(nx[0])*dx[0]
x1 = np.arange(nx[1])*dx[1]
x2 = np.arange(nx[2])*dx[2]
X0,X1,X2 = np.meshgrid(x0,x1,x2,indexing='ij')
f,ax = plt.subplots(2,3)
ss.imshow_slices(x0,x1,x2,I,ax[0])
ss.imshow_slices(x0,x1,x2,J,ax[1])
```

```python
# encode our energy operator in the Fourier domain using discrete derivatives
f0 = np.arange(nx[0])/(dx[0]*nx[0])
f1 = np.arange(nx[1])/(dx[1]*nx[1])
f2 = np.arange(nx[2])/(dx[2]*nx[2])
F0,F1,F2 = np.meshgrid(f0,f1,f2,indexing='ij')
a = 2.0*dx[0] # a scale in mm
p = 2
Ahat = (1.0 - 2.0*a**2*((np.cos(2.0*np.pi*dx[0]*F0) - 1.0)/dx[0]**2
                        + (np.cos(2.0*np.pi*dx[1]*F1) - 1.0)/dx[1]**2
                        + (np.cos(2.0*np.pi*dx[2]*F2) - 1.0)/dx[2]**2))**(2.0*p)
Khat = 1.0/Ahat
K = np.fft.ifftn(Khat).real
f,ax = plt.subplots(1,2)
ax[0].imshow(Ahat[:,:,0])
ax[0].set_title('A hat')
ax[1].imshow(K[:,:,0])
ax[1].set_title('K')
```

```python
# for display
f,ax = plt.subplots(3,3)

# parameters
sigma = 2e0 # 1/2/sigma^2 is the parameter in front of image matching integral square error
niter = 50 # how many iterations of gradient descent to run
epsilon = 2.0e-2 # step size in gradient descent
# mouse
sigma = 5.0e0
epsilon = 2.0e-3
# initialize
vt0, vt1, vt2 = np.zeros((nx[0],nx[1],nx[2],nt)), np.zeros((nx[0],nx[1],nx[2],nt)), np.zeros((nx[0],nx[1],nx[2],nt))
It = np.tile(I[:,:,:,None],(1,1,1,nt+1))
EMAll = []
ERAll = []
EAll = []

# start an optimization loop
for it in range(niter):
    # compute the deformed image by flowing over time
    phiinv0,phiinv1,phiinv2 = X0,X1,X2
    ER = 0.0
    for t in range(nt):
        # update phiinv using method of characteristics
        v0, v1, v2 = vt0[:,:,:,t], vt1[:,:,:,t], vt2[:,:,:,t]
        X0s, X1s, X2s = X0-v0*dt, X1-v1*dt, X2-v2*dt

        # deal with boundary conditions: subtract identity, use flat boundary conditions, then add identity
        phiinv0 = spi.interpn([x0,x1,x2],phiinv0-X0,np.stack([X0s,X1s,X2s],axis=-1),**interp_args_phiinv)+X0s
        phiinv1 = spi.interpn([x0,x1,x2],phiinv1-X1,np.stack([X0s,X1s,X2s],axis=-1),**interp_args_phiinv)+X1s
        phiinv2 = spi.interpn([x0,x1,x2],phiinv2-X2,np.stack([X0s,X1s,X2s],axis=-1),**interp_args_phiinv)+X2s

        # deform the image
        It[:,:,:,t+1] = spi.interpn([x0,x1,x2],I,np.stack([phiinv0,phiinv1,phiinv2],axis=-1),**interp_args)

        # calculate the energy of the flow
        Av0 = np.fft.ifftn(np.fft.fftn(v0)*Ahat).real
        Av1 = np.fft.ifftn(np.fft.fftn(v1)*Ahat).real
        Av2 = np.fft.ifftn(np.fft.fftn(v2)*Ahat).real
        ER += np.sum(v0*Av0 + v1*Av1 + v2*Av2)*0.5*dx[0]*dx[1]*dx[2]*dt

    [a.cla() for a in ax[0]]
    ss.imshow_slices(x0,x1,x2,It[:,:,:,-1],ax[0])

    # error
    err = It[:,:,:,-1] - J
    lambda1 = -(err)/sigma**2
    [a.cla() for a in ax[1]]
    ss.imshow_slices(x0,x1,x2,err,ax[1])
    EM = np.sum(err**2/(2.0*sigma**2))*dx[0]*dx[1]
    E = ER+EM
    EMAll.append(EM)
    ERAll.append(ER)
    EAll.append(E)
    ax[2][0].cla()
    ax[2][0].plot(EAll)
    ax[2][0].plot(EMAll)
    ax[2][0].plot(ERAll)

    # now we start flowing the error back in time
    phiinv0, phiinv1, phiinv2 = X0,X1,X2
    for t in range(nt-1,-1,-1):
        # update phiinv using method of characteristics, note "+" because we are integrating backward
        v0, v1, v2 = vt0[:,:,:,t], vt1[:,:,:,t], vt2[:,:,:,t]
        X0s, X1s, X2s = X0+v0*dt, X1+v1*dt, X2+v2*dt
        phiinv0 = spi.interpn([x0,x1,x2],phiinv0-X0,np.stack([X0s,X1s,X2s],axis=-1),**interp_args_phiinv)+X0s
        phiinv1 = spi.interpn([x0,x1,x2],phiinv1-X1,np.stack([X0s,X1s,X2s],axis=-1),**interp_args_phiinv)+X1s
        phiinv2 = spi.interpn([x0,x1,x2],phiinv2-X2,np.stack([X0s,X1s,X2s],axis=-1),**interp_args_phiinv)+X2s

        # find the determinant of Jacobian
        phiinv0_0,phiinv0_1,phiinv0_2 = np.gradient(phiinv0,dx[0],dx[1],dx[2])
        phiinv1_0,phiinv1_1,phiinv1_2 = np.gradient(phiinv1,dx[0],dx[1],dx[2])
        phiinv2_0,phiinv2_1,phiinv2_2 = np.gradient(phiinv2,dx[0],dx[1],dx[2])
        detjac = phiinv0_0*(phiinv1_1*phiinv2_2 - phiinv1_2*phiinv2_1)\
            - phiinv0_1*(phiinv1_0*phiinv2_2 - phiinv1_2*phiinv2_0)\
            + phiinv0_2*(phiinv1_0*phiinv2_1 - phiinv1_1*phiinv2_0)

        # find lambda_t
        lambdat = spi.interpn([x0,x1,x2],lambda1,np.stack([phiinv0,phiinv1,phiinv2],axis=-1),**interp_args)*detjac

        # get the gradient of the image at this time
        I_0, I_1, I_2 = np.gradient(It[:,:,:,t],dx[0],dx[1],dx[2])

        # initialize the gradient with the matching term
        grad0, grad1, grad2 = lambdat*I_0, lambdat*I_1, lambdat*I_2

        # smooth it
        grad0 = np.fft.ifftn(np.fft.fftn(grad0)*Khat).real
        grad1 = np.fft.ifftn(np.fft.fftn(grad1)*Khat).real
        grad2 = np.fft.ifftn(np.fft.fftn(grad2)*Khat).real

        # add the regularization term
        grad0 += v0
        grad1 += v1
        grad2 += v2

        # update
        vt0[:,:,:,t] -= epsilon*grad0
        vt1[:,:,:,t] -= epsilon*grad1
        vt2[:,:,:,t] -= epsilon*grad2

    f.canvas.draw()
```


```python
from IPython.core.display import HTML
from IPython.display import Image
HTML("""
<style>
.output_png {
    display: table-cell;
    text-align: center;
    vertical-align: middle;
}
</style>
""")
```

# *Circuitos Elétricos I - Semana 4*

### Problema 1

Determine o circuito equivalente de Thévenin ($v_{th}$, $R_{th}$) do ponto de vista dos terminais $(a,b)$ do circuito abaixo.

a) Determine a $v_{th}$ pelo método das tensões de nó.\
b) Determine a corrente de curto-circuito $i_{cc}$ pelo método das tensões de nó.\
c) Determine a $v_{th}$ via superposição.\
d) Determine a $R_{th}$ via resistência equivalente.\
e) Determine $v_{th}$, $R_{th}$ via transformação de fontes.

```python
Image("./figures/J6C1.png", width=700)
```

```python
import sympy as sp
import numpy as np
```

```python
# define as N variáveis desconhecidas
v1, v2, v3, v4 = sp.symbols('v1, v2, v3, v4')

# define os sistema de N equações
eq1 = sp.Eq()
eq2 = sp.Eq()
eq3 = sp.Eq()

# resolve o sistema
soluc = sp.solve((eq1, eq2, eq3), dict=True)

v1 = np.array([sol[v1] for sol in soluc])
v2 = np.array([sol[v2] for sol in soluc])
v4 = np.array([sol[v4] for sol in soluc])
v3 = 2

print('Solução do sistema:\n\n v1 = %.2f V,\n v2 = %.2f V,\n v3 = %.2f V,\n v4 = %.2f V.' %(v1, v2, v3, v4))
```

```python

```

Jupyter notebooks/Circuitos Eletricos I - Semana 4.2.ipynb


```python
from sympy.physics.optics import RayTransferMatrix, ThinLens, BeamParameter, FreeSpace
from sympy import symbols, init_printing, pi, sqrt, pprint, solveset, S, plot, cos, Eq
from scipy.constants import c
import matplotlib.pyplot as plt
from matplotlib.ticker import (MultipleLocator, FormatStrFormatter,
                               AutoMinorLocator)
import numpy as np

init_printing()

from ipywidgets.widgets import (
    IntSlider,
    FloatLogSlider,
    FloatSlider,
    interactive_output,
    FloatRangeSlider,
    BoundedFloatText,
    Checkbox,
    HBox,
    VBox,
    IntRangeSlider,
    BoundedIntText,
    Text,
    Button,
    Label,
    Select,
    Layout
)
from IPython.display import display
import datetime, os
from os.path import isdir
from pathlib import Path
```

```python
save_location = "D:\SurfDrive\backup_for_surfacePro\Work\THz_setup\Simulations"

w, lam, w0, z = symbols("w, lam, w0, z", positive=True)

# defining the confocal length, R(z) and w(z) in sympy
z_r = pi * w0 ** 2 / lam
#R = z * (1 + (z_r / z) ** 2)
w = w0 * sqrt(1 + (z / z_r) ** 2)

# Defining functions
def ghz_cm(x):
    return c * 1e2 / x * 1e-9

print(\
    "Confocal length: $Z_c$, \
\nRadius of the curvature of beam wave front: \
\nR(z) and the Beam radius: w(z) are defined such as: \n"
)

z_r, w
```

# Simulating Gaussian beam propagation from inside FELion Trap

We will try to obtain the optimum conditions such the waist radius w0 and distance to the w0 from outside mirror, etc.

```python
class gaussian_simulation:
    def __init__(self):

        W = BoundedFloatText(value=0.34, min=0.1, max=1, step=0.01, description="w0:")
        L = IntRangeSlider(
            value=[325, 1100], min=100, max=1500, step=1, description="GHz"
        )

        width = BoundedIntText(value=7, min=5, max=25, step=1, description="width")
        height = BoundedIntText(value=5, min=5, max=25, step=1, description="height")
        dpi = BoundedIntText(value=200, min=70, max=250, step=1, description="Dpi")

        entry1 = BoundedFloatText(
            value=-9, min=-11, max=10, step=0.1, description="Entry1"
        )
        Mirror = BoundedIntText(value=11, min=10, max=20, step=1, description="Mirror Position")

        self.save_name = Text(value="")

        save_button = Button(description="Save", layout={"width": "300px"})

        save_button.on_click(self.save_plot)

        self.print_save = Label(value="Status: ")

        row1 = HBox([W, L])
        row2 = HBox([width, height, dpi])
        row3 = HBox([entry1, Mirror])

        out = interactive_output(
            self.simulate,
            {
                "w0_": W,
                "wn": L,
                "width": width,
                "height": height,
                "dpi": dpi,
                "entry1": entry1,
                "mirror":Mirror
            },
        )

        display(
            VBox(
                [row1, row2, row3, out, self.save_name, save_button, self.print_save]
            )
        )

    def save_plot(self, event):
        if not isdir('./Simulations'):
            os.mkdir('./Simulations')

        self.fig0.savefig(f"./Simulations/{self.save_name.value}.png")
        self.fig0.savefig(f"./Simulations/{self.save_name.value}.pdf")

        t = datetime.datetime.now()
        self.print_save.value = f'File "{self.save_name.value}.png" Saved; {t.ctime()}'

    def beam_plot(self, ax, w0_, wn):

        self.gaussian_beam(w0_, wn)

        markers = '--'
        line_width = 1
        alpha = 0.7

        for c, beam in enumerate(self.beams):
            color = f"C{c}"
            ax.plot(
                z_,
                beam,
                markers+color, lw=line_width, alpha=alpha,
                label=f"{self.wn[c]:.2f} GHz ({ghz_cm(self.wn[c])*10:.2f} mm)",
                zorder=1
            )
            ax.plot(-z_, -beam, markers+color, lw=line_width, alpha=alpha, zorder=1)
            ax.plot(-z_, beam, markers+color, lw=line_width, alpha=alpha, zorder=1)
            ax.plot(z_, -beam, markers+color, lw=line_width, alpha=alpha, zorder=1)

    def simulate(self, w0_, wn, width, height, dpi, entry1, mirror):

        entry1 = round(entry1, 1)
        entry2 = round(10.4 - abs(entry1), 1)

        self.fig0, ax0 = plt.subplots(figsize=(width, height), dpi=dpi)
        self.beam_plot(ax0, w0_, wn)

        ax0.hlines(0, xmin=-z_.max(), xmax=z_.max())

        ## Annotation begins

        fontsz = 6

        # Distance b/w two entry points
        ax0.hlines(-0.2, xmin=entry1, xmax=entry2)
        plt.annotate(
            s=f"d={(entry2-entry1)}",
            xy=(0, 0),
            xytext=(-3, -0.1),
            ha="center",
            va="center",
            fontsize=fontsz
        )

        # Annotate beam waist
        ax0.vlines(x=0, ymin=0, ymax=w0_)
        plt.annotate(
            s="$w_0^{in}$=%.2f cm"%w0_,
            xy=(0, w0_+0.2),
            xytext=(0, 0.8),
            ha="center",
            va="center",
            arrowprops=dict(arrowstyle="->", connectionstyle="arc3"),
            fontsize=fontsz
        )

        # Annotate Entry1
        ax0.vlines(x=entry1, ymin=-1, ymax=1)

        plt.annotate(
            s="Main entry window:\ndiameter=2cm",
            xy=(entry1, 1),
            xytext=(entry1 + 6, 1.1),
            ha="center",
            va="center",
            arrowprops=dict(arrowstyle="->", connectionstyle="arc3"),
            fontsize=fontsz
        )

        # Annotate Entry2
        ax0.vlines(x=entry2, ymin=-0.43, ymax=0.43)

        plt.annotate(
            s="Trap bender entry:\ndiameter=0.86cm",
            xy=(entry2, -0.43),
            xytext=(entry2 - 3, -0.8),
            ha="center",
            va="center",
            # rotation = 30,
            arrowprops=dict(arrowstyle="->", connectionstyle="arc3"),
            fontsize=fontsz
        )

        # Inner entry: diameter: 0.64cm
        ax0.vlines(x=entry2+0.7, ymin=-0.32, ymax=0.32)

        # Annotate z-position

        plt.annotate(s="z=0", xy=(0, 0), xytext=(0, -0.05), ha="center", va="center",fontsize=fontsz)

        plt.annotate(
            s=f"z={entry1}",
            xy=(entry1, 0),
            xytext=(entry1 - 2, -0.2),
            ha="center",
            va="center",
            arrowprops=dict(arrowstyle="->", connectionstyle="arc3"),
            fontsize=fontsz
        )

        # d_in window
        mirror_pos = entry1-mirror
        ax0.vlines(x=mirror_pos, ymin=-self.beams.max(), ymax=self.beams.max())

        plt.annotate(
            s="Mirror($d_{in}$) at %.1f cm"%mirror_pos,
            xy=(mirror_pos+5, -self.beams.max()),
            ha="center",
            va="center",
            arrowprops=dict(arrowstyle="->", connectionstyle="arc3"),
            fontsize=fontsz
        )

        ## Annotation ends

        ax0.set_xlabel("z direction : Beam Propagation (cm)")
        ax0.set_ylabel("w(z) : Gaussian Beam radius (cm)")
        ax0.set_title(f"Gaussian Beam simulation: {wn[0]}-{wn[-1]} GHz")
        ax0.legend(bbox_to_anchor=(0.16, 0.86), ncol=3, fontsize=5)
        ax0.xaxis.set_minor_locator(AutoMinorLocator())

        plt.show()
        # plt.close()

        self.save_name.value = f"{w0_}cm-{wn[0]}-{wn[-1]}GHz_simulation"

    def gaussian_beam(self, w0_, wn):

        global w

        # Defining the z-propagation range
        z_ = np.linspace(0, 25, 10)

        self.wn = np.linspace(*wn, 10)

        beam = []
        beams = []

        for l in self.wn:

            temp = []
            beam = []

            for i in z_:
                temp = float(w.subs({w0: w0_, lam: ghz_cm(l), z: i}).evalf())
                beam.append(temp)
            beams.append(beam)

        self.beams = np.array(beams)
#%matplotlib widget
%matplotlib inline
start = gaussian_simulation()
```

    VBox(children=(HBox(children=(BoundedFloatText(value=0.34, description='w0:', max=1.0, min=0.1, step=0.01), In…

```python

```

# Ellipsoidal mirror (Goldsmith, 1998)

An ellipsoidal mirror in the geometrical optics limit, transforms radiation from a point source at one focal point to a point source located at the second focal point; this is true for any part of the ellipsoidal surface of revolution,
formed by rotating the ellipse about the line joining its foci.

First consider a lens that has a focal length equal to the distance from the first focal point to the center of the part of the ellipsoidal surface being used (R1).
This lens will collimate the radiation from the first focal point into a parallel beam.

If we immediately follow this first lens by a second, having focal length equal to the distance from the center of the mirror to the second focal point (R2) , the second lens will bring the parallel beam to a focus at the desired position.

## Thin lens
ABCD matrix for thin lens of focal length $f_{lens}$

```python
f_lens = symbols('f_lens')
M_lens = ThinLens(f_lens)
M_lens
```

## Elliptical mirror effective focal length

Thus, the ABCD matrix for an ellipsoidal mirror is the product of two matrices, each representing a thin lens, and having zero separation:
where R1 and R2 are the distances from the center of the section of the ellipsoid used to the respective foci

```python
R1, R2 = symbols('R1, R2')
M_l1 = ThinLens(R1)
M_l2 = ThinLens(R2)
M_d = FreeSpace(0)
M_ellipsoid = M_l2*M_d*M_l1
M_ellipsoid
```

So that comparing with the expression for the thin lens, we see that the Effective focal length of an elliptical mirror: f_ellipsoid is given by

```python
f_ellipsoid = 1/(1/R1+1/R2)
print('\n')
f_ellipsoid.simplify()
```

# Mode Matching (Goldsmith 1998)

Quasioptical system designers are often required to couple the power in a beam (e.g., that produced by a feed hom) to a specified beam waist radius (e.g., that required by some component in a receiver system). Dealing with beam transformation from this perspective is often called mode matching

## Case 1:
In practice most often the case is: We are given the input and output waist radii, and thus the magnification, together with the input distance. Solving for the required focal length, we obtain

```python
M, E, d_in, d_out, w0_in, w0_out, Z_c = symbols('M, E, d_in, d_out, w0_in, w0_out, Z_c')

fp = Z_c*(d_in/(E*Z_c))*(1+sqrt(1-E*(1+(d_in/Z_c)**-2)))
fm = Z_c*(d_in/(E*Z_c))*(1-sqrt(1-E*(1+(d_in/Z_c)**-2)))
```

Positive and negative focus "f" for M>1

```python
fm, fp
```

```python
f1 = Z_c*(1+(d_in/Z_c)**2)/(2*d_in/Z_c)
print('For M=1\nFocus is given by:')
f1
```

Assuming for the case of our ellipse,
$$d_{in} = R1$$  $$f = f_{ellipsoid}$$

we can solve for R2 i.e., $d_{out}$

The eccentricity can also be defined in terms of R1, R2, and the angle α between input and output beams (Goldsmith, 1998)

```python
alpha = symbols('alpha')
e_ellipsoid = sqrt(R1**2 + R2**2 - 2*R1*R2*cos(alpha))/(R1 + R2)
e_ellipsoid
```

```python
a, b = symbols('a, b')
print('Eccentricity of an elliptical mirror is already defined as :\n')

e_standard = sqrt(1-b**2/a**2)
e_standard
```

where a and b are the semi-major and semi-minor axis respectivey.

```python
print('The semi-major axis a is also defined as: \n')
a = (R1 + R2)/2
a
```

```python
class f_soln:
    def __init__(self):
        mag = BoundedFloatText(
            value=0.41, min=0.1, max=10, step=0.01, description="M: "
        )
        W = BoundedFloatText(
            value=0.34, min=0.1, max=1, step=0.01, description="w0_in: "
        )
        R1 = BoundedFloatText(value=18, min=5, max=20, step=0.1, description="R1: ")
        L = IntRangeSlider(
            value=[325, 500], min=50, max=1500, step=1, description="GHz"
        )

        width = BoundedIntText(value=10, min=5, max=25, step=1, description="width")
        height = BoundedIntText(value=6, min=5, max=25, step=1, description="height")
        dpi = BoundedIntText(value=100, min=70, max=250, step=1, description="Dpi")

        self.save_name = Text(value="")

        save_button = Button(description="Save", layout={"width": "300px"})

        save_button.on_click(self.save_plot)

        self.print_save = Label(value="Status: ")

        row1 = HBox([width, height, dpi])
        row2 = HBox([mag, W, R1, L])

        out = interactive_output(
            self.plot,
            {
                "M": mag,
                "r1": R1,
                "w0_in": W,
                "lam": L,
                "width": width,
                "height": height,
                "dpi": dpi,
            },
        )

        display(VBox([row1, row2, out, self.save_name, save_button, self.print_save]))

    def plot(self, M, r1, w0_in, lam, width, height, dpi):

        wn = np.linspace(*lam, 50)

        x, y1, y2 = [], [], []
        no_soln, soln = [], []
        for l in wn:
            zc = np.pi * w0_in ** 2 / ghz_cm(l)

            temp0 = r1 / zc
            x.append(temp0)

            try:

                if not M == 1:
                    temp1 = fm / zc
                    temp1_1 = float(
                        temp1.subs({E: 1 - M ** -2, d_in: r1, Z_c: zc}).evalf()
                    )
                elif M == 1:
                    temp1 = f1 / Zc
                    temp1_1 = float(temp1.subs({d_in: r1, Z_c: zc}).evalf())

                y1.append(temp1_1)
                soln.append(round(l, 2))
            except:
                no_soln.append(round(l, 2))

        self.fig, axs = plt.subplots(figsize=(width, height), dpi=dpi)

        x_ = x[: len(y1)]
        axs.plot(x_, y1, color='C0', label=f"f/$Z_c$; $w0(trap)$={w0_in}cm;" + "$d_{in}$=%.2f cm" % r1,)

        axs.legend()
        axs.grid()
        axs.set_xlabel("$d_{in}$/$Z_c$")
        axs.set_ylabel("f/$Z_c$", color='C0')
        axs.set_title(
            f"For M={M}: Solution exist for {soln[0]}-{soln[-1]} GHz"
        )
        axs.tick_params(axis='y', labelcolor='C0')

        if len(no_soln) > 0 and len(soln) > 0:
            print(f"\nSolution: Exists for {soln[0]}-{soln[-1]} GHz\n")

            print(f"Does not exist for: {no_soln[0]}-{no_soln[-1]} GHz\n")
        if len(soln) < 1:
            print(f"x: {x}\ny: {y1}")

        # Finding R2 position:
        x_ = np.array(x_)
        y1 = np.array(y1)
        r2 = r1 / ((x_ / y1) - 1)
        r2_ind = np.logical_and(r2 <= 50, r2 > 0)

        axs1 = axs.twinx()
        axs1.plot(x_[r2_ind], r2[r2_ind], color='C1', label="$d_{out}$")

        axs1.set_ylabel('$d_{out}$ (cm)', color='C1')
        axs1.tick_params(axis='y', labelcolor='C1')
        axs1.legend()

        #self.fig.legend()

        try:
            self.save_name.value = f"M-{M}_{soln[0]}-{soln[-1]}GHz_din_{r1}cm"
        except:
            pass
        plt.show()

    def save_plot(self, event):

        self.fig.savefig(f"./{self.save_name.value}.png")

        t = datetime.datetime.now()
        self.print_save.value = f'File "{self.save_name.value}.png" Saved; {t.ctime()}'
```

## Simulation to determine optimal focal length solution for elliptical mirror.

```python
start = f_soln()
```

    VBox(children=(HBox(children=(BoundedIntText(value=10, description='width', max=25, min=5), BoundedIntText(val…

The above simulation shows that we might need different focal length R2 for different wavelength, so we cannot use this method. Let us examine another case as defined below.

## Case 2:

The input and output distance to waist is a fixed sum d, then the expression for input and output distance is given by:

```python
f, f0, d, M, d_in, d_out, R_z, Z_c, Z_r = symbols("f, f0, d, M, d_in, d_out, R_z, Z_c, Z_r")
Wz, Rz = symbols('Wz, Rz')

eq_din_m = Eq(d_in, f - sqrt(f ** 2 - f0 ** 2) / M)
eq_din_p = Eq(d_in, f + sqrt(f ** 2 - f0 ** 2) / M)

eq_dout_m = Eq(d_out, f - sqrt(f ** 2 - f0 ** 2) * M)
eq_dout_p = Eq(d_out, f + sqrt(f ** 2 - f0 ** 2) * M)

eq_dout_din = Eq(d_out, f * (1 + M ** 2 * (d_in / f - 1)))
eq_Rz = Eq(Rz, z * (1 + (Z_c / z) ** 2))
eq_Wz = Eq(Wz, w0*sqrt(1+(z/Z_c)**2))

eq_din_m, eq_din_p, eq_dout_m, eq_dout_p, eq_dout_din.simplify(), eq_Rz, eq_Wz
```

```python
solveset(eq_dout_din, d_in), solveset(eq_dout_din, d_out)
```

We will now try to find the focal length for an optimal d_out corresponds to M (which provides w_out since we already know w_in) using the known set d_in value (from trap) and
$$f0 = \pi*(w0_{in}*w0_{out})/\lambda$$

```python
eq_dout2 = Eq(d_out, f * (1 + (M ** 2 * (d_in / f - 1))))
eq_din2 = Eq(d_in, f * (1 + sqrt(M ** -2 - (Z_r / f) ** 2)))
Rz = z * (1 + (Z_c / z) ** 2) # Radius of curvature
eq_din2, eq_dout2, eq_dout2.subs({d_in:eq_din2.rhs})
```

```python
class f_case2:
    def __init__(self):

        # Beam properties
        Win = BoundedFloatText(
            value=0.34, min=0.1, max=1, step=0.01, description="$w0_{in}$: "
        )
        Wout = BoundedFloatText(
            value=0.15, min=0.01, max=1, step=0.01, description="$w0_{out}$: "
        )
        D1 = BoundedFloatText(
            value=20, min=5, max=30, step=0.1, description="$d_{in}$: "
        )
        L = IntRangeSlider(
            value=[325, 500], min=50, max=1500, step=1, description="GHz"
        )

        # subplot adjust
        plot0_top = BoundedFloatText(
            value=0.79, min=0, max=1, step=0.01, description="top-sub"
        )
        plot0_bottom = BoundedFloatText(
            value=0.1, min=0, max=1, step=0.01, description="bottom-sub"
        )
        plot0_left = BoundedFloatText(
            value=0.1, min=0, max=1, step=0.01, description="left-sub"
        )
        plot0_right = BoundedFloatText(
            value=0.9, min=0, max=1, step=0.01, description="right-sub"
        )
        plot0_wspace = BoundedFloatText(
            value=0.34, min=0, max=1, step=0.01, description="wspace-sub"
        )

        # Legend label adjust
        Lg_X = BoundedFloatText(value=0.2, min=0.1, max=2, step=0.1, description="lg_x")
        Lg_Y = BoundedFloatText(value=0.8, min=0, max=2, step=0.1, description="lg_y")
        Lg_X1 = BoundedFloatText(value=1, min=0.1, max=2, step=0.1, description="lg_x")
        Lg_Y1 = BoundedFloatText(value=1, min=0, max=2, step=0.1, description="lg_y")

        # Figure property adjust
        width = BoundedIntText(value=15, min=5, max=25, step=1, description="width")
        height = BoundedIntText(value=6, min=5, max=25, step=1, description="height")
        dpi = BoundedIntText(value=100, min=70, max=300, step=1, description="Dpi")

        # Figure Label adjust
        Xlabel = Text(value="Frequency (GHz)")
        Ylabel_0 = Text(value=" + ve f (cm)")
        Ylabel_1 = Text(value=" - ve f (cm)")
        Title = Text(
            value="Determining focal length (f) from given M and $d_{in}$ then calculating $d_{out}$"
        )

        # Widget to save plot1
        self.save_name = Text(value="")
        save_button = Button(description="Save", layout={"width": "300px"})
        save_button.on_click(self.save_plot)
        self.print_save = Label(value="Status: ")

        # Plot 2 widgets
        self.focal_list1 = Select(
            options=["1"],
            description="+ve F:",
            layout=Layout(width="30%", height="160px"),
        )
        self.focal_list2 = Select(
            options=["2"],
            description="-ve F:",
            layout=Layout(width="30%", height="160px"),
        )
        custom = Checkbox(value=True, description="f-Auto")
        f1 = BoundedFloatText(value=5, min=1, max=25, step=0.1, description="f (+ve) ")
        f2 = BoundedFloatText(
            value=-15, min=-25, max=-1, step=0.1, description="f(-ve): "
        )
        Alpha = BoundedFloatText(
            value=90, min=0, max=360, step=1, description=r"2\(\theta\)"
        )

        # Widget to save plot2

        self.save_name1 = Text(value="")
        save_button1 = Button(description="Save", layout={"width": "300px"})
        save_button1.on_click(self.save_plot1)
        self.print_save1 = Label(value="Status: ")

        # Making interactive outputs
        out = interactive_output(
            self.plot,
            {
                "d1": D1,
                "w0_in": Win,
                "w0_out": Wout,
                "lam": L,
                "width": width,
                "height": height,
                "dpi": dpi,
                "xlabel": Xlabel,
                "ylabel_0": Ylabel_0,
                "ylabel_1": Ylabel_1,
                "title": Title,
                "lg_x": Lg_X,
                "lg_y": Lg_Y,
                "top": plot0_top,
                "bottom": plot0_bottom,
                "left": plot0_left,
                "right": plot0_right,
                "wspace": plot0_wspace,
            },
        )
        out1 = interactive_output(
            self.plot2,
            {
                "focal_list1": self.focal_list1,
                "focal_list2": self.focal_list2,
                "d1": D1,
                "width": width,
                "height": height,
                "dpi": dpi,
                "lam": L,
                "w0_in": Win,
                "w0_out": Wout,
                "lg_x": Lg_X1,
                "lg_y": Lg_Y1,
                "f_auto": custom,
                "f_manual1": f1,
                "f_manual2": f2,
                "theta": Alpha,
            },
        )

        # Widget to save plot3
        self.save_name2 = Text(value="")
        save_button2 = Button(description="Save", layout={"width": "300px"})
        save_button2.on_click(self.save_plot2)
        self.print_save2 = Label(value="Status: ")

        out2 = interactive_output(
            self.new_beams_fn,
            {
                "f_auto": custom,
                "focal_list1": self.focal_list1,
                "focal_list2": self.focal_list2,
                "f_manual1": f1,
                "f_manual2": f2,
                "win_new": Win,
                "w_out_use": Wout,
                "d1_in": D1,
                "width": width,
                "height": height,
                "dpi": dpi,
            },
        )

        # Alignment of widgets
        row1 = HBox([width, height, dpi])
        row2 = HBox([Win, Wout, D1, L])
        row3 = HBox([Xlabel, VBox([Ylabel_0, Ylabel_1]), Title])
        row4 = HBox([out])
        row5 = HBox(
            [
                Lg_X,
                Lg_Y,
                VBox([plot0_top, plot0_bottom, plot0_left, plot0_right, plot0_wspace]),
            ]
        )
        row6 = VBox([self.save_name, save_button, self.print_save])
        row7 = HBox([self.focal_list1, self.focal_list2])
        row8 = HBox([out1])
        row9 = HBox([Lg_X1, Lg_Y1])
        row10 = HBox(
            [
                VBox([self.save_name1, save_button1, self.print_save1]),
                VBox([custom, f1, f2, Alpha]),
            ]
        )
        row11 = VBox([out2])
        row12 = VBox([self.save_name2, save_button2, self.print_save2])

        display(
            VBox(
                [
                    row1,
                    row2,
                    row3,
                    row4,
                    row5,
                    row6,
                    row7,
                    row8,
                    row9,
                    row10,
                    row11,
                    row12,
                ]
            )
        )

    def plot(
        self,
        d1,
        w0_in,
        w0_out,
        lam,
        width,
        height,
        dpi,
        xlabel,
        ylabel_0,
        ylabel_1,
        title,
        lg_x,
        lg_y,
        top,
        bottom,
        left,
        right,
        wspace,
    ):

        M_ = w0_out / w0_in
        lam = np.linspace(*lam, 10)

        f0_ = []
        f_m, f_p = [], []
        d_out_m, d_out_p = [], []
        no_soln, soln = [], []

        freq = []

        for l in lam:
            tempx = (np.pi * w0_in * w0_out) / ghz_cm(l)

            try:
                f0_.append(tempx)
                tempy = solveset(
                    eq_din_p.subs({M: M_, d_in: d1, f0: tempx}), f
                )  # find focal length, f,  from d_in (+ve) equation

                tempy_m = float(tempy.args[0])  # Negative f
                tempy_p = float(tempy.args[1])  # Positive f

                f_m.append(tempy_m)
                f_p.append(tempy_p)

                temp_dout_m = float(
                    eq_dout_p.subs({M: M_, f: tempy_m, f0: tempx}).rhs
                )  # finding d_out (equation corresponds to same sign as d_in: +ve) from the calculated -ve f
                temp_dout_p = float(
                    eq_dout_p.subs({M: M_, f: tempy_p, f0: tempx}).rhs
                )  # finding d_out (equation corresponds to same sign as d_in: +ve) from the calculated +ve f

                d_out_m.append(temp_dout_m)
                d_out_p.append(temp_dout_p)

                soln.append(l)
                freq.append(l)

            except:
                no_soln.append(l)

        f0_ = np.array(f0_)
        f_m = np.array(f_m)
        f_p = np.array(f_p)
        d_out_m = np.array(d_out_m)
        d_out_p = np.array(d_out_p)

        print(f"f0: {f0_}")
        print(f"Frequencies: {lam}")
        print(f"\n+ve focus: {f_p}\n-ve focus: {f_m}\n")
        print(f"\nd_out_+ve: {d_out_p}\nd_out_-ve: {d_out_m}\n")

        self.focal_list1.options = f_p
        self.focal_list2.options = f_m

        if len(no_soln) > 0 and len(soln) > 0:
            print(f"\nSolution: Exists for {soln[0]}-{soln[-1]} GHz\n")
            print(f"nSolution does not exist for: {no_soln[0]}-{no_soln[-1]} GHz\n")
        elif len(soln) > 0:
            print(f"Solution: Exists for {soln[0]}-{soln[-1]} GHz")

        # self.fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(width, height), dpi=dpi)
        self.fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(width, height), dpi=dpi)
        plt.suptitle(
            title
            + "\n$w^{in}_{0}$: %.2f cm, $w^{out}_{0}$: %.3f cm, $d_{in}$: %.2f cm; [%i-%i] GHz"
            % (w0_in, w0_out, d1, lam[0], lam[-1])
        )

        plt.subplots_adjust(
            top=top, bottom=bottom, left=left, right=right, wspace=wspace
        )
        ax0, ax1 = ax  # , ax.twinx()

        # Subplot 1: Positve focus

        ax0.plot(freq, f_p, "--", color="C0", label="+ve f")

        ax0.tick_params(axis="y", labelcolor="C0")
        ax0.grid()
        ax0.set_xlabel(xlabel)
        ax0.set_ylabel(ylabel_0, color="C0")

        ax0_1 = ax0.twinx()
        ax0_1.plot(freq, d_out_p, "o", color="C1", label="+ $d_{out}$")
        ax0_1.plot(freq, -d_out_m, "x", color="C2", label="|-$d_{out}$|")

        ax0_1.grid()
        ax0_1.tick_params(axis="y", labelcolor="C1")
        ax0_1.set_ylabel("+ve $d_{out} (cm)$", color="C1")

        # Plotting f0
        ax0_2 = ax0.twiny()
        ax0_2.plot(f0_, f_p, ".", color="C0")
        ax0_2.set_xlabel("$f_0$ (cm)")

        # Subplot 2: Negative focus

        ax1.plot(freq, f_m, "--", color="C3", label=" -ve f")

        ax1.set_ylabel(ylabel_1, color="C3")
        ax1.set_xlabel(xlabel)
        ax1.grid()
        ax1.tick_params(axis="y", labelcolor="C3")

        ax1_1 = ax1.twinx()
        ax1_1.plot(freq, d_out_m, "o", color="C4", label=" - $d_{out}$")

        ax1_1.grid()
        ax1_1.tick_params(axis="y", labelcolor="C4")
        ax1_1.set_ylabel("-ve $d_{out} (cm)$", color="C4")

        # Plotting f0
        ax1_2 = ax1.twiny()
        ax1_2.plot(f0_, f_m, ".", color="C3")
        ax1_2.set_xlabel("$f_0$ (cm)")

        # Placing Legends
        ax0.legend(bbox_to_anchor=(lg_x, lg_y))
        ax0_1.legend(bbox_to_anchor=(lg_x, lg_y - 0.1))

        ax1.legend(bbox_to_anchor=(1 - lg_x, lg_y))
        ax1_1.legend(bbox_to_anchor=(1 - lg_x, lg_y - 0.1))

        # plt.tight_layout()

        self.save_name.value = f"f_for_{lam[0]}-{lam[-1]}GHz"

    def plot2(
        self,
        focal_list1,
        focal_list2,
        d1,
        width,
        height,
        dpi,
        lam,
        w0_in,
        w0_out,
        lg_x,
        lg_y,
        f_auto,
        f_manual1,
        f_manual2,
        theta,
    ):

        if not f_auto:
            focal_list1, focal_list2 = f_manual1, f_manual2

        M_ = w0_out / w0_in
        lam = np.linspace(*lam, 10)
        effective_f = float(f_ellipsoid.subs({R1: focal_list1, R2: focal_list2}))

        d_out_f0 = []
        d_out_values1, d_out_values2 = [], []
        freq, error_freq = [], []

        for wn in lam:
            try:

                temp_f0 = (np.pi * w0_in * w0_out) / ghz_cm(wn)

                temp_dout1 = float(
                    eq_dout_p.subs({M: M_, f0: temp_f0, f: focal_list1}).rhs
                )
                temp_dout2 = float(
                    eq_dout_p.subs({M: M_, f0: temp_f0, f: focal_list2}).rhs
                )

                d_out_values1.append(temp_dout1)
                d_out_values2.append(temp_dout2)

                freq.append(wn)
                d_out_f0.append(temp_f0)

            except Exception as e:
                print(f"Error: {e}\n")
                error_freq.append(wn)

        d_out_values1 = np.array(d_out_values1)  # +ve
        d_out_values2 = np.array(d_out_values2)  # -ve
        d_out_f0 = np.array(d_out_f0)

        # Making Plots

        self.fig1, axs = plt.subplots(figsize=(width, height), dpi=dpi)

        # Plot 1 +ve d_out

        axs.plot(
            freq, d_out_values1, color="C0", label=f"f(+ve) = {focal_list1:.2f} cm"
        )

        axs.grid()
        axs.tick_params(axis="y", labelcolor="C0")
        axs.set_ylabel("$d_{out}$ from +ve focus (cm)", color="C0")
        axs.set_xlabel("Frequency (GHz)")
        axs.set_title(
            "Determining $d_{out}$ from a fixed $f$ and $d_{in}$\n$w^{in}_{0}$: %.2f cm, $w^{out}_{0}$: %.3f cm, $d_{in}$: %.2f cm, [%i-%i] GHz"
            % (w0_in, w0_out, d1, lam[0], lam[-1])
        )

        # Plot 2 (twinx) -ve d_out
        axs1 = axs.twinx()
        axs1.plot(
            freq, d_out_values2, "x", color="C1", label=f"f(-ve) = {focal_list2:.2f} cm"
        )

        axs1.grid()
        axs1.tick_params(axis="y", labelcolor="C1")
        axs1.set_ylabel("$d_{out}$ from -ve focus (cm)", color="C1")

        # Placing Legends
        axs.legend(bbox_to_anchor=(lg_x, lg_y))
        axs1.legend(bbox_to_anchor=(lg_x, lg_y - 0.1))

        self.save_name1.value = f"d_out_for_{lam[0]}-{lam[-1]}GHz"

    def new_beams_fn(
        self,
        f_auto,
        focal_list1,
        focal_list2,
        f_manual1,
        f_manual2,
        win_new,
        w_out_use,
        d1_in,
        width,
        height,
        dpi,
    ):

        d_trap = d1_in-21.4

        colors = "C0, C1, C2, C3, C4, C5, C6, C7, C8, C9".split(", ")

        if f_auto:
            f_manual1, f_manual2 = focal_list1, focal_list2

        print(f"Focii: f_+ = {f_manual1:.2f}cm; f_- = {f_manual2:.2f}cm\n")

        wn_size = 5 # no. of wn
        size = 20

        # Defining figure:

        self.fig2 = plt.figure(figsize=(width + 5, height + 5), dpi=dpi)
        nrows, ncols = 3, 3

        ax3 = self.fig2.add_subplot(nrows, ncols, 1)
        ax4 = self.fig2.add_subplot(nrows, ncols, 2, sharey=ax3)
        bx0 = self.fig2.add_subplot(nrows, ncols, 3)

        ax5 = self.fig2.add_subplot(nrows, ncols, 4, sharex=ax3)
        ax6 = self.fig2.add_subplot(nrows, ncols, 5, sharey=ax5)
        bx1 = self.fig2.add_subplot(nrows, ncols, 6)

        ax7 = self.fig2.add_subplot(nrows, ncols, 7, sharex=ax3)
        ax8 = self.fig2.add_subplot(nrows, ncols, 8, sharey=ax7)
        bx2 = self.fig2.add_subplot(nrows, ncols, 9)

        # 3-sets of wavelength range for 3-different horn
        wn_set1 = np.linspace(325, 500, wn_size)
        wn_set2 = np.linspace(500, 750, wn_size)
        wn_set3 = np.linspace(750, 1100, wn_size)
        wn_list = np.array([wn_set1, wn_set2, wn_set3])

        # 3-sets of horns waist for different wavelength
        w0_out_list = [0.15, 0.099, 0.066]

        # In-use
        f_in_use = f_manual1

        # Making empty array
        d1 = []
        d2 = []
        f0_full = []

        for num, w0_out_temp in enumerate(w0_out_list):

            d1_temp = []
            d2_temp = []
            f0_full_temp = []

            wn_new_list = wn_list[num]
            M_ = w0_out_temp / win_new

            print(f"w_out: {w0_out_temp}\nw_in: {win_new}\nM={M_}\nwn: {wn_new_list}\n")

            for wn_new in wn_new_list:
                temp_f0 = (np.pi * win_new * w0_out_temp) / ghz_cm(wn_new)
                f0_full_temp.append(temp_f0)

                temp_dout = float(eq_dout_p.subs({M: M_, f0: temp_f0, f: f_in_use}).rhs)
                d2_temp.append(temp_dout)

                temp_din = float(eq_din_p.subs({M: M_, f0: temp_f0, f: f_in_use}).rhs)
                d1_temp.append(temp_din)

                ################################################################################################################
            ####################################################################################################################

            d1.append(d1_temp)
            d2.append(d2_temp)
            f0_full.append(f0_full_temp)

            print(f"d_out: {d2_temp}\nf0: {f0_full_temp}\nd_in: {d1_temp}\n")

            #####################################################################################################################
        #########################################################################################################################

        d1 = np.array(d1)
        d2 = np.array(d2)
        f0_full_range = np.array(f0_full_range)

        d1_beam = []
        d2_beam = []

        d1_beam_trap_optimal = []
        d1_beam_trap_required = []
        d2_beam_optimal = []
        d2_beam_from_trap_required = []
        d1_beam_from_required_d2 = []

        d2_required = []
        d1_from_required_d2 = []

        num = 0
        for w0_out_new in w0_out_list:

            wn_new_list = wn_list[num]
            d_in_new_list = d1[num]
            d_out_new_list = d2[num]

            mrkr = [".-", "x-", "--"]
            c = 0

            d1_beam_trap_optimal_temp = []
            d1_beam_trap_required_temp = []
            d2_beam_optimal_temp = []
            d2_beam_from_trap_required_temp = []
            d1_beam_from_required_d2_temp = []

            d2_req_temp = []
            d1_from_required_d2_temp = []

            for wn_new, d_out_new, d_in_new in zip(
                wn_new_list, d_out_new_list, d_in_new_list
            ):

                # Calculating d2 optimal beam from optimal d_out (d2)
                optimal_d2 = d2[0, 0]

                d2_optimal_array = np.linspace(0, d_out_new, size)
                d2_beam_optimal_calculation = np.array(
                    [
                        float(
                            w.subs({w0: w0_out_new, lam: ghz_cm(wn_new), z: i}).evalf()
                        )
                        for i in d2_optimal_array
                    ]
                )

                d2_plot = np.linspace(optimal_d2 - d_out_new, optimal_d2, size)

                ax3.plot(
                    d2_plot,
                    d2_beam_optimal_calculation,
                    colors[c] + mrkr[num],
                    zorder=1,
                    label=wn_new,
                )
                ax3.plot(
                    d2_plot,
                    -d2_beam_optimal_calculation,
                    colors[c] + mrkr[num],
                    zorder=1,
                )

                # Calculating d2 optimal beam from optimal d_in (d1)

                d1_optimal_array = np.linspace(
                    d_in_new - d1_in + d_trap, d_in_new, size
                )  # d1_in is the fixed distance from trap to mirror
                d1_beam_trap_optimal_calculation = np.array(
                    [
                        float(w.subs({w0: win_new, lam: ghz_cm(wn_new), z: i}).evalf())
                        for i in d1_optimal_array
                    ]
                )

                d1_plot = np.linspace(d_trap, d1_in, size)
                ax4.plot(
                    d1_plot, d1_beam_trap_optimal_calculation, colors[c] + mrkr[num], zorder=1
                )
                ax4.plot(
                    d1_plot, -d1_beam_trap_optimal_calculation, colors[c] + mrkr[num], zorder=1
                )

                # Making d1 propagates the way we want (i.e., 0 to d1_in, the fixed distance from trap to mirror)
                d1_fixed = np.linspace(d_trap, d1_in, size)

                d1_beam_trap_required_calculation = np.array(
                    [
                        float(w.subs({w0: win_new, lam: ghz_cm(wn_new), z: i}).evalf())
                        for i in d1_fixed
                    ]
                )

                ax6.plot(
                    d1_fixed, d1_beam_trap_required_calculation, colors[c] + mrkr[num], zorder=1
                )
                ax6.plot(
                    d1_fixed, -d1_beam_trap_required_calculation, colors[c] + mrkr[num], zorder=1
                )

                # Now from the position in mirror d_in, we calculate the required d_out distance (NOTE: This is not from optimal solution).
                wz_at_mirror_for_d2 = d1_beam_trap_required_calculation[-1]
                zc_w_out = (np.pi * w0_out_new ** 2) / ghz_cm(wn_new)

                d2_wz_required_position = float(
                    solveset(
                        eq_Wz.subs(
                            {Wz: wz_at_mirror_for_d2, w0: w0_out_new, Z_c: zc_w_out}
                        ),
                        z,
                    ).args[-1]
                )

                d2_required_array = np.linspace(0, d2_wz_required_position, size)

                d2_beam_from_trap_required_calculation = np.array(
                    [
                        float(
                            w.subs({w0: w0_out_new, lam: ghz_cm(wn_new), z: i}).evalf()
                        )
                        for i in d2_required_array
                    ]
                )

                d2_new_pos = optimal_d2 - d2_wz_required_position
                d2_plot = np.linspace(d2_new_pos, optimal_d2, size)

                ax5.plot(
                    d2_plot, d2_beam_from_trap_required_calculation, colors[c] + mrkr[num], zorder=1
                )
                ax5.plot(
                    d2_plot, -d2_beam_from_trap_required_calculation, colors[c] + mrkr[num], zorder=1
                )

                # Calculating d_in from the required d_out

                ax7.plot(
                    d2_plot, d2_beam_from_trap_required_calculation, colors[c] + mrkr[num], zorder=1
                )
                ax7.plot(
                    d2_plot, -d2_beam_from_trap_required_calculation, colors[c] + mrkr[num], zorder=1
                )

                M_ = w0_out_new/win_new
                d1_distance_from_required_d2 = float(
                    solveset(
                        eq_dout_din.subs(
                            {d_out: d2_wz_required_position, f: f_in_use, M: M_}
                        ), d_in
                    ).args[0]
                )

                d1_from_required_d2_array = np.linspace(d1_in - d1_distance_from_required_d2 + d_trap, d1_in, size)
                d1_beam_from_required_d2_calculation = np.array(
                    [
                        float(
                            w.subs({w0: win_new, lam: ghz_cm(wn_new), z: i}).evalf()
                        )
                        for i in d1_from_required_d2_array
                    ]
                )

                ax8.plot(
                    d1_plot, d1_beam_from_required_d2_calculation, colors[c] + mrkr[num], zorder=1
                )
                ax8.plot(
                    d1_plot, -d1_beam_from_required_d2_calculation, colors[c] + mrkr[num], zorder=1
                )

                # Gaussian beams calculated (temp)
                d1_beam_trap_optimal_temp.append(d1_beam_trap_optimal_calculation)
                d1_beam_trap_required_temp.append(d1_beam_trap_required_calculation)
                d2_beam_optimal_temp.append(d2_beam_optimal_calculation)
                d2_beam_from_trap_required_temp.append(d2_beam_from_trap_required_calculation)
                d1_beam_from_required_d2_temp.append(d1_beam_from_required_d2_calculation)

                d2_req_temp.append(d2_wz_required_position)
                d1_from_required_d2_temp.append(d1_distance_from_required_d2)
                c += 1

                #####################################################################################################################
            #########################################################################################################################

            # Gaussian beams calculated
            d1_beam_trap_optimal.append(d1_beam_trap_optimal_temp)
            d1_beam_trap_required.append(d1_beam_trap_required_temp)
            d2_beam_optimal.append(d2_beam_optimal_temp)
            d2_beam_from_trap_required.append(d2_beam_from_trap_required_temp)
            d1_beam_from_required_d2.append(d1_beam_from_required_d2_temp)

            d2_required.append(d2_req_temp)
            d1_from_required_d2.append(d1_from_required_d2_temp)
            num += 1

            #####################################################################################################################
        #########################################################################################################################

        # Making every list into numpy array
        d1_beam_trap_optimal = np.array(d1_beam_trap_optimal)
        d1_beam_trap_required = np.array(d1_beam_trap_required)
        d2_beam_optimal = np.array(d2_beam_optimal)
        d2_beam_from_trap_required = np.array(d2_beam_from_trap_required)
        d1_beam_from_required_d2 = np.array(d1_beam_from_required_d2)

        d2_required = np.array(d2_required)
        d1_from_required_d2 = np.array(d1_from_required_d2)

        #########################################################################################################################

        ## Plot configuration

        # ax3
        ax3.set(title="THz source to Mirror (optimal $d_{out}$ from calculation)", ylabel="w(z)-Beam Radius (cm)")
        ax3.legend(ncol=3)
        ax3.set_xlim(right=d2[0, 0])

        # ax4
        ax4.set(
            title="Mirror to Trap ($d_{in}$=%.1fcm) calculated beam path" % d1_in
        )
        ax4.legend(
            title=f"$f_+$ = {f_manual1:.2f}cm\n$f_-$ = {f_manual2:.2f}cm",
            loc="upper right",
        )

        # ax5
        ax5.set(title="Required $d_{out}$ (not from calculation)", xlabel="$d_{out}$ (cm)", ylabel="w(z)-Beam Radius (cm)")
        ax5.set_xlim(right=d2[0, 0])

        # ax6
        ax6.set(
            title="Required beam path(not calculated)"
        )
        ax6.set(xlabel="$d_{in}$ (cm)")

        # ax7
        ax7.set(title="Required $d_{out}$ (not from calculation)", xlabel="$d_{out}$ (cm)", ylabel="w(z)-Beam Radius (cm)")
        ax7.set_xlim(right=d2[0, 0])

        # ax8
        ax8.set(
            title="Required beam path(calculated)"
        )
        ax8.set(xlabel="$d_{in}$ (cm)")

        #########################################################################################################################

        for axes in (ax3, ax4, ax5, ax6, ax7, ax8):
            axes.grid()
            axes.xaxis.set_minor_locator(AutoMinorLocator())

        for axes in (ax4, ax6, ax8):
            axes.vlines(d_trap, ymin=-0.43, ymax=0.43)
            axes.vlines(0, ymin=-win_new, ymax=win_new)
            axes.annotate(f"$w_0$:{win_new}", xy=(3, 0))
            axes.set_xlim(right=d1_in)
            axes.invert_xaxis()

        #########################################################################################################################
        # New plot (for cheking the beam radius at trap):
        beam_waists = (d1_beam_trap_optimal, d1_beam_trap_required, d1_beam_from_required_d2)

        for ax, beam in zip((bx0, bx1, bx2), beam_waists):

            for i, w0_plot in enumerate(w0_out_list):
                ax.plot(np.arange(wn_size), beam[i][:, 0], '.-', label=f"{w0_plot}:{wn_list[i]}")

        for ax in (bx0, bx1, bx2):
            ax.set(xlabel='Not defined', ylabel='w(z) at trap(0.43 cm)', title='Beam radius achieved at trap')
            ax.hlines(0.43, xmin=0, xmax=4)
            ax.grid()
            ax.legend()

        plt.tight_layout()

        #########################################################################################################################

        with open("./Simulations/thz_simulation.log", "w") as log_file:

            t = datetime.datetime.now()
            print("Saving log file")

            log_file.write(f"Gaussian beam simulations: {t.ctime()}\n")
            log_file.write("####################################################\n\n")
            log_file.write(f"Focii: f_+ve = {f_manual1:.2f}cm; f_-ve = {f_manual2:.2f}cm\nMirror size required >= {(d2_beam_optimal[:, :, -1].max()*3):.2f}cm\n\n")
            log_file.write(f"Frequency: {wn_list}\n\nw0_in_trap: {win_new}\nw0_out: {w0_out_list}\nf0: {f0_full_range}\n")
            log_file.write("\n\n####################################################\n\n")
            log_file.write(f"\nBeam radius(At mirror):\n{d2_beam_optimal[:, :, -1]}\n")
            log_file.write("\n\n####################################################\n\n")

            for log, w0_out_log in  enumerate(w0_out_list):

                log_file.write(f"w0_out: {w0_out_log}\n")
                log_file.write(f"Frequency: {wn_list[log]}\n")

                log_file.write(f"Optimal distances:\n")
                log_file.write(f"d_in_distance:\n{d1[log]}\n")
                log_file.write(f"d_out_distance:\n{d2[log]}\n\n")

                log_file.write(f"Optimal beam:\n")
                log_file.write(f"d_in_optimal_beam at trap:\n{d1_beam_trap_optimal[log][:, 0]}\n")
                log_file.write(f"d_in_beam:\n{d1_beam_trap_optimal[log]}\n")
                log_file.write(f"d_out_beam:\n{d2_beam_optimal[log]}\n\n")

                log_file.write(f"Required Beam ( not calculated):\n")
                log_file.write(f"d_in_required_beam at trap (not calculated):\n{d1_beam_trap_required[log][:, 0]}\n")
                log_file.write(f"d_in_required_beam (for {d1_in}cm trap to mirror distance):\n")
                log_file.write(f"{d1_beam_trap_required[log]}\n\n")

                log_file.write(f"d_out_required distances:\n{d2_required[log]}\n")
                log_file.write(f"d_out_required_beam:\n{d2_beam_from_trap_required[log]}\n\n")

                log_file.write(f"Required Beam (calculated):\n")
                log_file.write(f"d_in_required_beam at trap (calculated):\n{d1_beam_from_required_d2[log][:, 0]}\n")
                log_file.write(f"d_in_from_required_d2 distances:\n{d1_from_required_d2[log]}\n")
                log_file.write(f"d_in_beam_from_required_d2:\n{d1_beam_from_required_d2[log]}\n\n")
                log_file.write("####################################################\n\n")

            log_file.write("####################################################\n\n")

            log_file.write("Simulation over.\n\nEND OF THE DOCUMENT\n")

            print("thz_simulation.log file saved")

        self.save_name2.value = f"calibrated_d_out_for_325-1100 GHz"

    def save_plot(self, event):

        try:
            if not isdir("./Simulations"):
                os.mkdir("./Simulations")

            self.fig.savefig(f"./Simulations/{self.save_name.value}.png")
            self.fig.savefig(f"./Simulations/{self.save_name.value}.pdf")

            t = datetime.datetime.now()
            self.print_save.value = (
                f'File "{self.save_name.value}.png" Saved; {t.ctime()}'
            )

        except:
            self.print_save2.value = "Error Occured."

    def save_plot1(self, event):

        try:
            if not isdir("./Simulations"):
                os.mkdir("./Simulations")

            self.fig1.savefig(f"./Simulations/{self.save_name1.value}.png")
            self.fig1.savefig(f"./Simulations/{self.save_name1.value}.pdf")

            t = datetime.datetime.now()
            self.print_save1.value = (
                f'File "{self.save_name1.value}.png" Saved; {t.ctime()}'
            )

        except:
            self.print_save2.value = "Error Occured."

    def save_plot2(self, event):

        try:
            if not isdir("./Simulations"):
                os.mkdir("./Simulations")

            self.fig2.savefig(f"./Simulations/{self.save_name2.value}.png")
            self.fig2.savefig(f"./Simulations/{self.save_name2.value}.pdf")

            t = datetime.datetime.now()
            self.print_save2.value = (
                f'File "{self.save_name2.value}.png" Saved; {t.ctime()}'
            )

        except:
            self.print_save2.value = "Error Occured."
```

```python
#%matplotlib widget
%matplotlib inline
start = f_case2()
```

    VBox(children=(HBox(children=(BoundedIntText(value=15, description='width', max=25, min=5), BoundedIntText(val…

```python
eq_dout_din
```

```python
eq_dout_din.subs({M:0.15/0.34, f:6.29, d_in:20}), eq_dout_din.subs({M:0.099/0.34, f:6.29, d_in:20}), eq_dout_din.subs({M:0.066/0.34, f:6.29, d_in:20})
```

```python
eq_din2, eq_dout2
```

```python
wn1 = np.linspace(325, 500, 10)
M1 = 0.15

for wn in wn1:
    zr = (np.pi*0.34**2)/ghz_cm(wn)
    temp_d1 = float(eq_din2.subs({f:6.29, M:M1/0.34, Z_r:zr}).rhs)
    temp_d2 = float(eq_dout2.subs({f:6.29, M:M1/0.34, d_in:temp_d1}).rhs)
    print(f"wn: {wn:.2f}\nd_in: {temp_d1:.1f} d_out: {temp_d2:.5f}\n")
```

```python
wn2 = np.linspace(500, 750, 10)
M2 = 0.099

for wn in wn2:
    zr = (np.pi*0.34**2)/ghz_cm(wn)
    temp_d1 = float(eq_din2.subs({f:6.29, M:M2/0.34, Z_r:zr}).rhs)
    temp_d2 = float(eq_dout2.subs({f:6.29, M:M2/0.34, d_in:temp_d1}).rhs)
    temp_d2_ = float(eq_dout2.subs({f:6.29, M:M2/0.34, d_in:20}).rhs)
    print(f"wn: {wn:.2f}\nd_in: {temp_d1:.1f}\nd_out1: {temp_d2:1f} d_out2: {temp_d2_:.1f}\n")
```

    wn: 500.00
    d_in: 27.0
    d_out1: 8.048031 d_out2: 7.5

    wn: 527.78
    d_in: 26.9
    d_out1: 8.039445 d_out2: 7.5

    wn: 555.56
    d_in: 26.8
    d_out1: 8.030349 d_out2: 7.5

    wn: 583.33
    d_in: 26.7
    d_out1: 8.020735 d_out2: 7.5

    wn: 611.11
    d_in: 26.6
    d_out1: 8.010594 d_out2: 7.5

    wn: 638.89
    d_in: 26.5
    d_out1: 7.999917 d_out2: 7.5

    wn: 666.67
    d_in: 26.3
    d_out1: 7.988694 d_out2: 7.5

    wn: 694.44
    d_in: 26.2
    d_out1: 7.976914 d_out2: 7.5

    wn: 722.22
    d_in: 26.0
    d_out1: 7.964565 d_out2: 7.5

    wn: 750.00
    d_in: 25.9
    d_out1: 7.951634 d_out2: 7.5

```python
wn3 = np.linspace(750, 1100, 10)
M3 = 0.066

for wn in wn3:
    zr = (np.pi*0.34**2)/ghz_cm(wn)
    temp_d1 = float(eq_din2.subs({f:6.29, M:M3/0.34, Z_r:zr}).rhs)
    temp_d2 = float(eq_dout2.subs({f:6.29, M:M3/0.34, d_in:temp_d1}).rhs)
    temp_d2_ = float(eq_dout2.subs({f:6.29, M:M3/0.34, d_in:20}).rhs)
    print(f"wn: {wn:.2f}\nd_in: {temp_d1:.1f}\nd_out1: {temp_d2:.1f} d_out2: {temp_d2_:.1f}\n")
```

    wn: 750.00
    d_in: 37.4
    d_out1: 7.5 d_out2: 6.8

    wn: 788.89
    d_in: 37.3
    d_out1: 7.5 d_out2: 6.8

    wn: 827.78
    d_in: 37.1
    d_out1: 7.5 d_out2: 6.8

    wn: 866.67
    d_in: 36.9
    d_out1: 7.4 d_out2: 6.8

    wn: 905.56
    d_in: 36.8
    d_out1: 7.4 d_out2: 6.8

    wn: 944.44
    d_in: 36.6
    d_out1: 7.4 d_out2: 6.8

    wn: 983.33
    d_in: 36.4
    d_out1: 7.4 d_out2: 6.8

    wn: 1022.22
    d_in: 36.2
    d_out1: 7.4 d_out2: 6.8

    wn: 1061.11
    d_in: 36.0
    d_out1: 7.4 d_out2: 6.8

    wn: 1100.00
    d_in: 35.8
    d_out1: 7.4 d_out2: 6.8

```python
zc = np.pi*0.34**2/ghz_cm(325)
M_ = 0.15/0.34

fp.subs({Z_c:zc, E:(1-M_**-2), d_in:20}), fm.subs({Z_c:zc, E:(1-M_**-2), d_in:20})
```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```


## Symbolic analysis of models

In this notebook, we use symbolic mathematics to study energy-based PySB model. We derive steady-state analytical solutions to study reaction networks without costly ODE simulations. First, we load the previously developed models for RAF and RAF inhibition:

```python
import toy_example_RAF_RAFi as model
from pysb.bng import generate_equations
from util_display import display_model_info

#generate the model equations
model=model.model
generate_equations(model)

#display model informations
display_model_info(model)
```

    Model information
    Species: 6
    Parameters: 12
    Expressions: 8
    Observables: 8
    Total Rules: 2
    Energy Rules: 2
    Non-energy Rules: 0
    Energy Patterns: 4
    Reactions: 12

Next, we define the sympy systems of equations corresponding to the ODE system but with the left-hand side (the derivative definition) set to zero, meaning that the system is considered to be at steady-state:

```python
import sympy
import scipy

#create a list of expressions to be substituted into kinetic rates of reactions
species = sympy.Matrix([sympy.Symbol(f"s_{i}", nonnegative=True) for i in range(len(model.species))])
subs = {e: e.expand_expr() for e in model.expressions | model._derived_expressions}
subs.update({sympy.Symbol(f"__s{i}"): s for i, s in enumerate(species)})
kinetics = sympy.Matrix([r['rate'] for r in model.reactions]).xreplace(subs)
#simplyfy kinetic
kinetics.simplify()
sm = sympy.SparseMatrix(*model.stoichiometry_matrix.shape, model.stoichiometry_matrix.todok())
obs_matrix = scipy.sparse.lil_matrix(
    (len(model.observables), len(model.species)), dtype=int
)
for i, obs in enumerate(model.observables):
    obs_matrix[i, obs.species] = obs.coefficients
om = sympy.SparseMatrix(*obs_matrix.shape, obs_matrix.todok())
odes = sm * kinetics
observables = om * species
```

The following cell currently needs to be customized to your specific model. Define conservation of mass expressions for all monomers, and an expression you would like to solve for.

```python
# Define conservation of mass expressions (each equal to zero).
conservation = sympy.Matrix([
    model.parameters["R_0"] - observables[0],
    model.parameters["I_0"] - observables[1],
])

system = sympy.Matrix.vstack(odes)
# This is just R_BRAFmut_active_obs, but it could be any expression.
#R_active = sympy.Symbol('R_active')
#expression = sympy.Matrix([R_active - observables[2]])
system = sympy.Matrix.vstack(odes, conservation)
display(system)
```

$\displaystyle \left[\begin{matrix}f^{1 - \phi_{RR}} kr_{RR} s_{4} - kf_{RI} s_{0} s_{1} - 1.0 kf_{RR} s_{0}^{2} + kr_{RI} s_{3} + 2 kr_{RR} s_{2} - f^{- \phi_{RR}} kf_{RR} s_{0} s_{3}\\f^{1 - \phi_{RI}} kr_{RI} s_{4} + 2.0 g^{1 - \phi_{RI}} kr_{RI} s_{5} - kf_{RI} s_{0} s_{1} + kr_{RI} s_{3} - g^{- \phi_{RI}} kf_{RI} s_{1} s_{4} - 2.0 f^{- \phi_{RI}} kf_{RI} s_{1} s_{2}\\f^{1 - \phi_{RI}} kr_{RI} s_{4} + 0.5 kf_{RR} s_{0}^{2} - kr_{RR} s_{2} - 2.0 f^{- \phi_{RI}} kf_{RI} s_{1} s_{2}\\f^{1 - \phi_{RR}} kr_{RR} s_{4} + kf_{RI} s_{0} s_{1} - 1.0 kf_{RR} s_{3}^{2} \left(\frac{1}{f g}\right)^{\phi_{RR}} - kr_{RI} s_{3} + 2 kr_{RR} s_{5} \left(f g\right)^{1 - \phi_{RR}} - f^{- \phi_{RR}} kf_{RR} s_{0} s_{3}\\- f^{1 - \phi_{RI}} kr_{RI} s_{4} - f^{1 - \phi_{RR}} kr_{RR} s_{4} + 2.0 g^{1 - \phi_{RI}} kr_{RI} s_{5} - g^{- \phi_{RI}} kf_{RI} s_{1} s_{4} + f^{- \phi_{RR}} kf_{RR} s_{0} s_{3} + 2.0 f^{- \phi_{RI}} kf_{RI} s_{1} s_{2}\\- 2.0 g^{1 - \phi_{RI}} kr_{RI} s_{5} + 0.5 kf_{RR} s_{3}^{2} \left(\frac{1}{f g}\right)^{\phi_{RR}} - kr_{RR} s_{5} \left(f g\right)^{1 - \phi_{RR}} + g^{- \phi_{RI}} kf_{RI} s_{1} s_{4}\\R_{0} - s_{0} - 2 s_{2} - s_{3} - 2 s_{4} - 2 s_{5}\\I_{0} - s_{1} - s_{3} - s_{4} - 2 s_{5}\end{matrix}\right]$

Solve the combined system of the ODEs and conservation expressions for the list of symbols used in our desired expression. There may be multiple solutions

```python
#solve the symbolic systems
#solutions = sympy.solve(system, list(species), force=True, manual=True)

#unfortunately the sympy solver seems not to be able to solve even this simple system
#working on implementin a new approach tailored for mass-action kinetic systems with specific characteristics
```

We evaluate the expressions or observable as function of the steady state species concentrations (S_x) to calculate the amount of active RAFs:

```python
# to be done
```