import numpy as np
import matplotlib.pyplot as plt
# make sure we see it on this notebook
%matplotlib inline

def eigenfunction1D(n, L, x):
    """ Normalized eigenfunction for the 1D particle in a box.
    
    n -- the quantum number
    L -- the size of the box
    x -- the NumPy array with the x values
    """
    raise NotImplementedError('You need to write this function!')


x = np.linspace(0, 10.0, 1000)
eig = eigenfunction1D(1, 10.0, x)
prob = eigenfunction1D(1, 10.0, x)**2

plt.plot(x, eig)
plt.plot(x, prob)

"""
Explanation: <figure>
  <IMG SRC="gfx/Logo_norsk_pos.png" WIDTH=100 ALIGN="right">
</figure>

Particle in a box
Roberto Di Remigio, Luca Frediani
The particle in a box model is among the simplest, exactly solvable models in quantum mechanics.
In the one-dimensional case, we assume a particle of mass $m$ to be confined into a box of length $L$.
The confinement is achieved by means of a potential energy operator that is zero inside the box and infinite outside, as in the Figure below.
<figure>
  <IMG SRC="gfx/Infinite_potential_well-en.svg">
</figure>

In practice, this means that the particle cannot escape the box: a result that we would have obtained also
from classical mechanics.
How does the quantum particle behave? We need to find the eigenfunctions and eigenvalues of the Hamiltonian operator, that is we have to solve the following ordinary differential equation:
\begin{equation}
-\frac{\hbar^2}{2m}\frac{\mathrm{d}^2}{\mathrm{d}x^2} \psi_n(x) = E_n\psi_n(x)
\end{equation}
with boundary conditions:
\begin{equation}
\begin{aligned}
  \psi_n(0) &= 0 \
  \psi_n(L) &= 0
\end{aligned}
\end{equation}
Thus acceptable solutions are of the form:
\begin{equation}
\psi_n(x) = \sin(\frac{n\pi x}{L}) \quad\quad \forall n \neq 0
\end{equation}
with energies:
\begin{equation}
E_n = \frac{h^2n^2}{8mL^2}  \quad\quad \forall n \neq 0
\end{equation}
Exercise 1: Normalization
The wavefunction(s) given above as solution to the particle in a box problem are not normalized:
\begin{equation}
\left\langle \psi_n | \psi_n \right\rangle = \int \mathrm{d}x \psi_n^*(x)\psi_n(x) = |A|^2 \neq 1
\end{equation}
Find the normalization constant.
Exercise 2: Ground-state and probabilities
Given the normalized ground-state wavefunction for a particle confined to a one-dimensional box of length $L$, suppose the box is $10.0\,\,\mathrm{nm}$ long. What is the probability that the particle is:
1. Between $a = 4.95\,\,\mathrm{nm}$ and $b=5.05\,\,\mathrm{nm}$
2. Between $a = 1.95\,\,\mathrm{nm}$ and $b=2.05\,\,\mathrm{nm}$
3. Between $a = 9.90\,\,\mathrm{nm}$ and $b=10.0\,\,\mathrm{nm}$
4. In the right half of the box
5. In the central third of the box
How can we generalize to the excited states?
Exercise 3: Eigenfunctions and probability distributions
Write a Python function to plot the eigenfunctions of the particle in a box model. The function should accept the quantum number $n$, the length of the box $L$ and a NumPy array of $x$ values as arguments:
```Python
def eigenfunction1D(n, L, x):
    """ Normalized eigenfunction for the 1D particle in a box.
n -- the quantum number
L -- the size of the box
x -- the NumPy array with the x values
"""

Once this function is defined, we can obtain the respective probability distribution by taking its square:Python
x = np.linspace(0, 10.0, 1000)
eig = eigenfunction1D(1, 10.0, x)
prob = eigenfunction1D(1, 10.0, x)**2
and plot both of them with:Python
plt.plot(x, eig)
plt.plot(x, prob)
```
End of explanation
"""


import numpy as np
import matplotlib.pyplot as plt
# make sure we see it on this notebook
%matplotlib inline

def normalize(coeffs):
    """ Normalization constant for a linear combination of 1D particle in a box eigenfunctions

    coeffs -- a NumPy array with the coefficients
    """
    raise NotImplementedError('You need to write this function!')
    
    
coeffs = np.array([0.5, 0.5])
normalize(coeffs)

"""
Explanation: Exercise 4: Normalization of linear combinations of eigenfunctions
A linear combination is defined as:
\begin{equation}
\Psi(x) = \sum_{i=0}^{N} c_i \psi_i(x)
\end{equation}
where $c_i$ are the coefficients and $\psi_i(x)$ are the eigenfunctions of the particle in a box.
Is $\Psi(x)$ normalized? If not, how can we normalize it? Write a Python function that returns the normalization constant given a vector of coefficients $c_i$ in the linear combination.
```Python
def normalize(coeffs):
    """ Normalization constant for a linear combination of 1D particle in a box eigenfunctions
coeffs -- a NumPy array with the coefficients
"""

```
End of explanation
"""


def energy(coeffs, M, L):
    """ Return energy of a linear combination of 1D particle in a box eigenfunctions.

    coeffs -- the coefficients of the linear combination
    M -- the mass of the particle
    L -- the size of the box
    """
    raise NotImplementedError('You have to write this function')

coeffs = np.array([0.5, 0.5])
energy(coeffs, 1.0, 10.0)

"""
Explanation: Exercise 5: Energy of linear combinations
The expectation value of an operator $O$ on the wavefunction $\Psi$ is defined as:
\begin{equation}
\bar{O} = \frac{\langle\Psi|O|\Psi\rangle}{\langle\Psi|\Psi\rangle} = \langle\Psi|O|\Psi\rangle
\end{equation}
where the last equality holds only if $\Psi$ is normalized.
The energy is the expectation value of the Hamiltonian operator:
\begin{equation}
\bar{H} = E = \langle\Psi|H|\Psi\rangle
\end{equation}
Given $\Psi$ a linear combination of particle in a box eigenfunctions, calculate its respective energy.
Hint: remember that $H\psi_i(x) = E_i\psi_i(x)$ and that the $\psi_i(x)$ are normalizd. You can also assume $\Psi$ to be already normalized.
Write a Python function to calculate the energy of such linear combinations. The function should accept the coefficients $c_i$ of the linear combination, the mass $M$ of the particle and the length $L$ of the box as arguments.
```Python
def energy(coeffs, M, L):
    """ Return energy of a linear combination of 1D particle in a box eigenfunctions.
coeffs -- the coefficients of the linear combination
M -- the mass of the particle
L -- the size of the box
"""

```
End of explanation
"""


