```python
import numpy as np
import astropy.units as u
import matplotlib.pyplot as plt
import seaborn as sns
import scipy.integrate

from astropy import io
from galpy import potential
from galpy.orbit import Orbit
from galpy.util import bovy_plot
from scipy.optimize import curve_fit,root

sns.set()
```

### Problem 1: Dark matter and scaling relations in disk galaxies.

In one of the recent Q&A sessions, we played with rotation curves from the SPARC
compilation (Lelli et al. 2016) and decomposed them into disk and halo contributions. Let’s
see what else we can learn about dark matter and galaxy structure from these data!

Important note: You can use the functions from the class notebook to load the and
process the data, but be aware that the table entries as read by the code do not appear to
have proper astropy units set. This means that when you naively pass them to, e.g., galpy
functions, galpy will interpret the input as having no units. You can get around this issue by
explicitly casting the table to the units that you want, for example, data[’R’].to(u.kpc).
If you use galpy for this problem (which is probably useful. . . ), it may also help to add a
configuration file1 with astropy-units = True (this is not the default) to the directory in
which you are working. That way, galpy functions will always return astropy Quantities
with units attached. https://docs.galpy.org/en/latest/installation.html#configuration-file

(a) When we did the exercise in class, we used disk rotation curves that were unconstrained by other data. But we in fact have surface-brightness measurements for all of the
galaxies in the sample (included in the table as SBdisk), which we can use to determine
an appropriate disk rotation curve and its contribution to the total rotation curve. For the
three galaxies that we studied in class, re-fit the rotation curve as a razor-thin disk plus
NFW halo, but first finding good disk parameters by fitting the surface brightness profile
(assume M/L = 1 as is appropriate in the mid-IR used in the SPARC compilation). You
can just fit by hand and determine parameters that roughly fit by eye.

```python
raw = io.ascii.read('MassModels_Lelli2016c.mrt')
ugc05 = raw[raw['ID'] == 'UGC05918']
ugc08 = raw[raw['ID'] == 'UGC08490']
ngc65 = raw[raw['ID'] == 'NGC6503']
```

```python
ugc05
```

Table length=8

ID	D	R	Vobs	e_Vobs	Vgas	Vdisk	Vbul	SBdisk	SBbul

	Mpc	kpc	km / s	km / s	km / s	km / s	km / s	Lsun / pc2	Lsun / pc2

str11	float64	float64	float64	float64	float64	float64	float64	float64	float64

UGC05918	7.66	0.56	16.9	3.93	3.56	6.92	0.0	13.38	0.0

UGC05918	7.66	1.11	27.4	3.93	3.93	11.04	0.0	9.28	0.0

UGC05918	7.66	1.67	33.5	3.93	3.56	13.76	0.0	6.42	0.0

UGC05918	7.66	2.23	36.4	3.93	8.55	16.36	0.0	4.67	0.0

UGC05918	7.66	2.79	39.8	3.93	12.27	16.13	0.0	2.61	0.0

UGC05918	7.66	3.34	41.2	3.93	12.47	16.4	0.0	2.22	0.0

UGC05918	7.66	3.9	42.6	3.93	13.27	17.86	0.0	2.07	0.0

UGC05918	7.66	4.46	44.5	3.93	15.66	16.94	0.0	0.0	0.0

```python
color = sns.color_palette("crest", 2) # It's snowing idk

disk_surfdens = [20.,430,2000]*u.Msun/u.pc**2
disk_scalelen = [1.5,0.59,.9]*u.kpc
gal = [ugc05,ugc08,ngc65]
galname = ['UGC05918','UGC08490','NGC6503']

f, axes = plt.subplots(1,3,figsize=(25,5))

for i,ax in enumerate(axes):

    R = np.array(gal[i]['R']) * u.kpc

    disk = potential.RazorThinExponentialDiskPotential(amp=disk_surfdens[i],hr=disk_scalelen[i])

    # Set vertical height to be 0 so surface density will be in plane of the disk.
    surface_density = disk.surfdens(R,0)

    ax.plot(R,surface_density,label="Razor Thin Exponential Disk",color=color[0],lw=10,alpha=.6)

    ax.scatter(R,gal[i]['SBdisk'],label="SBdisk",color=color[1],edgecolor="black",s=100,zorder=10)

    ax.set_title(galname[i])
    ax.set_xlabel('R (kpc)')
    ax.set_ylabel(r'Surface Brightness ($L_\odot / pc^2 $)')
axes[2].legend()
plt.suptitle("Surface Brightness Profile");
```

```python
color = sns.color_palette("viridis", 4)

# My 'by-eye' guesses for the virial mass and concentration of the dark matter halo.
halo_virial_mass = [5e10,1e11,1e12] * u.Msun
halo_conc = [10,21,8]

for i in range(3):
    # Set up Plot
    plt.figure(figsize=(10,5))
    if i ==0:
        plt.suptitle("Rotation Curve as a Razor-Thin Disk plus NFW Halo");
    # Fit the disk by a Razor Thin Exponential Disk
    disk = potential.RazorThinExponentialDiskPotential(amp=disk_surfdens[i],hr=disk_scalelen[i])
    # Fit the dark matter halo with an NFW potential
    halo= potential.NFWPotential(mvir=halo_virial_mass[i].to_value(1e12*u.Msun),conc=halo_conc[i],ro=8.,vo=220.)
    # Add these potentials together.
    disk_halo = disk + halo

    plt.xlim(0.,1.1*np.max(gal[i]['R']))
    plt.ylim(0.,1.3*np.max(gal[i]['Vobs']))

    R = np.array(gal[i]['R']) * u.kpc
    R_range = [.5 * R.min(),1.1*R.max()]
    V = np.array(gal[i]['Vobs']) * u.km / u.s
    V_er = np.array(gal[i]['e_Vobs']) * u.km / u.s

    plt.errorbar(R.value,V.value,yerr=V_er.value,ls='none',color=color[0])
    disk.plotRotcurve(Rrange=R_range,overplot=True,label=r'$\mathrm{disk}$',color=color[1])
    halo.plotRotcurve(Rrange=R_range,overplot=True,label=r'$\mathrm{halo}$',color=color[2])
    potential.plotRotcurve(disk+halo,Rrange=R_range,overplot=True,label=r'$\mathrm{total}$',color=color[3])
    plt.title(galname[i])
    plt.ylabel(r'$V_c\,(\mathrm{km\,s}^{-1})$')
    plt.xlabel(r'$R\,(\mathrm{kpc})$')
    plt.legend(loc='lower right',frameon=False,fontsize=18.)

```

(b) What do you learn about the dark matter contribution in these
galaxies?

- In all cases, dark matter contributes the most at large radii though in the first case, UGC05918, dark matter is the primary contributor at all radii. The amount of dark matter present exceeds that of ordinary matter.
- In the last case, NGC6503, fitting the surface brightness with a Razor Thin Exponential disk may have not been the best potential since the disk-model contributes excessively to the rotation curve at short radii despite the model fitting well with the data.
- UGC08490 seems to have the highest concentration, where the concentration of dark matter seems to affect the shape of the model where the virial mass affects some of the shape but mostly the height.

(c) The Tully-Fisher relation is a famous scaling relation for disk
galaxies that relates a galaxies asymptotic velocity $V_\infty$ to its
luminosity. It can be used, for example, to obtain distances to
galaxies, because the rotation velocity can be measured relatively
easily and the luminosity obtained from the relation can then be
combined with the observed magnitude to obtain the distance. It's also
a useful constraint on galaxy formation models. Let's use the
bulgeless SPARC galaxies to measure the Tully-Fisher relation
ourselves! To obtain a quantitative determination of $V_\infty$, we can
fit the rotation curves with a gravitational potential of the form

\begin{equation}
  \phi(r) = {v_0^2\over 2}\,\ln\left(r^2 + r_0^2\right)\,,
\end{equation}

with two free parameters, $v_0$ and $r_0$. Determine the circular
velocity as a function of radius $v_c(r)$ for this potential. What is
$V_\infty$ in terms of $v_0$ and $r_0$?

The centripetal acceleration is given by

\begin{equation}
    a_r = \frac{d\Phi}{dr} \tag{8.14}
\end{equation}

\begin{equation}
    a_r = \frac{d\Phi}{dr} = \frac{d}{dr}\Big({v_0^2\over 2}\,\ln\left(r^2 + r_0^2\right)\Big) = v_o^2 \frac{r}{r^2 + r_0^2}
\end{equation}

\begin{equation}
    a_r = \frac{v_c^2}{r} \tag{8.14}
\end{equation}

\begin{equation}
   \frac{v_c^2}{r} = v_o^2 \frac{r}{r^2 + r_0^2}
\end{equation}

The circular velocity as a function of radius:

\begin{equation}
   v_c(r) = v_o \frac{r}{\sqrt{r^2 + r_0^2}}
\end{equation}

$V_\infty$ in terms of $v_0$ and $r_0$

\begin{equation}
   v_c(r\rightarrow\infty) = v_o \frac{\infty}{\sqrt{\infty^2 + r_0^2}} = v_o
\end{equation}

(d) Now fit this model to all bulgeless SPARC galaxies. Fit the model rotation velocity by minimizing $\chi^2$ defined as

\begin{equation}
  \chi^2 = \sum_i {\left[v_c(\texttt{R}_i) - \texttt{Vobs}_i\right]^2 \over \texttt{e\_Vobs}_i^2}\,,
\end{equation}

where $i$ indexes the $N$ data points for each galaxy $\{\texttt{R}_i,\texttt{Vobs}_i,\texttt{e\_Vobs}_i\}$.

You can use scipy optimization functions for this. Compare the model fit to the observed rotation curve for the three galaxies from part (a) and do the fit for all 143 bulgeless galaxies.

```python
class rotcurve:

    def __init__(self,galaxy,step_size):

        # For organization gather observed values into the class
        self.R = galaxy['R']
        self.Vobs = galaxy['Vobs']
        self.eVobs = galaxy['e_Vobs']
        self.galaxy_name = galaxy['ID'][0]
        self.L_per_area = galaxy['SBdisk']

        # Create grid of parameters to search through
        self.v0_range = np.arange(0,np.max(self.Vobs)+100,step_size)
        self.r0_range = np.arange(0,np.max(self.R)+50,step_size)

        # Meshgrid combines all possible variations of v0_range and r0_range
        self.v0_grid,self.r0_grid = np.meshgrid(self.v0_range,self.r0_range)
        self.v0_grid,self.r0_grid = self.v0_grid.flatten(),self.r0_grid.flatten()

        # Get chi squared value for each combination of parameters
        self.sols = self.search(galaxy,self.v0_grid,self.r0_grid)

        # Search for the minimum chi sq solution
        self.ind = np.where(self.sols==np.min(self.sols))
        self.r0 = self.r0_grid[self.ind]
        self.v0 = self.v0_grid[self.ind]
        self.vc = self.v_c(self.R,self.v0,self.r0)

    def v_c(self,r,v0,r0):
        # Equation derived above for the circular velocity
        return v0 * r / (np.sqrt(r**2+r0**2))

    def chisq(self,vc,vobs,e_vobs):
        # Chi squared test
        return np.sum((vc - vobs)**2 / e_vobs**2)

    def search(self,galaxy,v0_grid,r0_grid):
        sols = []
        for v0,r0 in zip(v0_grid,r0_grid):
            # Run parameter grid through model of circular velocity (vc) and calculate the chi squared (chisq) value for each combination
            sols.append(self.chisq(self.v_c(galaxy['R'],v0,r0),galaxy['Vobs'],galaxy['e_Vobs']))
        return sols

```

```python
galaxy = [ugc05,ugc08,ngc65]

f,axes = plt.subplots(1,3,figsize=(20,5))
color = sns.color_palette("viridis", 2)
for ax,gal in zip(axes,galaxy):

    # Get class object rotcurve for each of the three galaxies
    RC = rotcurve(gal,1)
    # Plot the observed data
    ax.errorbar(RC.R,RC.Vobs,yerr=RC.eVobs,ls='none',color=color[0],label='data')
    # Plot the model
    ax.plot(RC.R,RC.vc,color=color[1],label='model')
    # Format
    ax.set_title(RC.galaxy_name)
    ax.set_ylabel(r'$V_c\,(\mathrm{km\,s}^{-1})$')
    ax.set_xlabel(r'$R\,(\mathrm{kpc})$')
axes[2].legend()
```

```python
# Stolen from class
def bulgeless():
    raw= io.ascii.read('MassModels_Lelli2016c.mrt')
    allgals= set(raw['ID'])
    out= ()
    for gal in allgals:
        data= raw[raw['ID'] == gal]
        if np.all(np.fabs(data['Vbul']) < 1e-10):
            out= out+(gal,)
    return out
# Get the names of galaxies that contain no bulge component
nobulge_name = bulgeless()
# Get a class object for each of these
NoBulgeRC = []
for name in nobulge_name:
    galaxy = raw[raw['ID'] == name]
    NoBulgeRC.append(rotcurve(galaxy,1))
```

(e) Now derive the Tully-Fisher relation by computing the total luminosity of each galaxy and plotting $V_\infty$ versus the total
luminosity. Discuss what you see.

The Tully-Fisher relation is generally well fit as $L \propto V_\infty^\alpha$; if you fit this form to your obtained relation, what value for $\alpha$ do you get (you can fit without taking errors into account).

What $\alpha$ do you get if you only include $L > 10^{10}\,L_\odot$ galaxies?

Which of these two determinations do you think is more accurate and why?

```python
color = sns.color_palette("viridis", 3)
def L(RC):
    # Integrate 2 * pi * R * Surface Brightness per Area * dR to get full luminosity
    R = RC.R.to(u.pc).value # Handle Units
    return np.trapz(2*np.pi*R*RC.L_per_area,x=R)

# Get Vinf and Luminosity from the functions we've written.
Vinf = np.array([RC.v0[0] for RC in NoBulgeRC])
L = np.array([L(RC) for RC in NoBulgeRC])

def plot(Vinf,L,plot_title):
    # Power law function relation luminosity to Vinf
    def TF(V,alpha):
        return V ** alpha

    # Fit function to data
    popt, pcov = curve_fit(TF, Vinf, L)

    # Create range over which to commute fit
    Vinf_range = np.arange(25,350)

    plt.figure(figsize=(15,5))
    # Computed Vinf and L
    plt.scatter(np.log10(L),np.log10(Vinf),color=color[0],label='Fit')
    # Power Law Fit
    plt.plot(np.log10(TF(Vinf_range,popt[0])),np.log10(Vinf_range),color=color[2],lw=5,alpha=0.6)
    # Formatting
    plt.xlabel(r'Log$_{10} (L_\odot)$',fontsize=20)
    plt.ylabel(r'Log$_{10}(V_\infty)$',fontsize=20)
    plt.title(plot_title+f'  alpha={popt[0]:.8}',fontsize=25);

plot(Vinf,L,'Power Law Model, Fit over Entire Range')
```

```python

bright_V = Vinf[L > 1e10]
bright_L = L[L > 1e10]

plot(bright_V,bright_L,'Power Law Model, Bright Galaxies ($L > 10^{10} L_\odot$)')
plt.scatter(np.log10(L),np.log10(Vinf),color=color[1],zorder=0,label='Not Fit')
plt.legend();
```

```python
dim_V = Vinf[L < 1e10]
dim_L = L[L < 1e10]

plot(dim_V,dim_L,'Power Law Model, Bright Galaxies ($L < 10^{10} L_\odot$)')
plt.scatter(np.log10(L),np.log10(Vinf),color=color[1],zorder=0,label='Not Fit')
plt.legend();
```

```python

small_V = Vinf[Vinf < 90]
small_L = L[Vinf < 90]

plot(small_V,small_L,'Power Law Model, Small $V_\infty$ ($V_\infty < 90$ km s$^{-1}$)')
plt.scatter(np.log10(L),np.log10(Vinf),color=color[1],zorder=0,label='Not Fit')
plt.legend();
```

```python
large_V = Vinf[Vinf > 90]
large_L = L[Vinf > 90]

plot(large_V,large_L,'Power Law Model, Small $V_\infty$ ($V_\infty > 90$ km s$^{-1}$)')
plt.scatter(np.log10(L),np.log10(Vinf),color=color[1],zorder=0,label='Not Fit')
plt.legend();
```

```python
alphas =[4.6201564, 4.6201395, 4.6367407, 4.7571277, 4.620151]
np.mean(alphas),np.std(alphas)
```

    (4.650863060000001, 0.0535194945850052)

Fitting brighter galaxies $L > 10^{10}\,L_\odot$ versus the full range did not seem to change the power law exponent by much. The values were the same up to a few significant figures. I took a look at McGaugh et al., 2000 and they discussed a break in the Tully-Fisher relation around $V_c \sim 90$ km s$^{-1}$ so I played around with fitting different ranges of the data. You can see the break by eye as low luminosity (or low mass since we use $L\sim M$) galaxies have more spread. Given this you would expect fitting brighter galaxies to give a better fit since the relation is tighter. Oddly it seems to me that, by eye and without factoring in errors, fitting $V_c < 90$ gives the best fit. Regardless the power law slope is given by $4.65\pm 0.05$ so the variation between fits is not extremely high. If it is true, that bright galaxies give the best fit, that would be quite lucky! As it would enable distance measurements at farther distances which does seem to be the case given the cosmic distance ladder.

### Problem 2: The zero-velocity curve.

When we looked at orbits in disk galaxies, we discussed the
zero-velocity curve: the curve in the meridional plane $(R,z)$ where
$v_R = v_z = 0$. Let's explore this important curve further!

(a) For the example orbit in Chapter 10.1, we showed
the zero-velocity curve as a constant energy contour, but we could
also explicitly compute $z$ as a function of $(R,E,L_z)$ on the
zero-velocity curve by demanding that $v_R = v_z = 0$. Do this for the
first example orbit in Chapter 10.1 and compare to the zero-velocity
curve that you get from the contouring method. (note that you may not
be able to solve for $z$ analytically and might have to resort to
numerical methods)

```python
def plot_potential(E,Lz,pot):
    # Stolen from you (:
    levels = np.linspace(-1.25,0.,11)
    cntrcolors = ['k' for l in levels]
    cntrcolors[np.arange(len(levels))\
       [np.fabs(levels-E) < 0.01][0]]= '#ff7f0e'
    potential.plotPotentials(pot,
                     rmin=0.3,rmax=1.05,
                     zmin=-0.3,zmax=0.3,
                     nrs=101,nzs=101,
                     justcontours=True,
                     levels=levels,
                     cntrcolors=cntrcolors,
                     effective=True,
                     Lz=Lz)
```

```python
R_range=np.arange(0.42,0.93,0.01)
guess = np.repeat(0.1,len(R_range))

def effPot(z,R=R_range):
    # When does E = effPot ?
    # effPot is equation 10.8
    pot = potential.MWPotential2014
    E = -1.25
    Lz = 0.60
    Pot = potential.evaluatePotentials(pot,R,z)
    return Pot + (Lz/R)**2./2 - E

# Use FindRoot method to get a value for z that will make E equal the effective potential for a range of R values.
Zs = root(effPot,guess)

# Plot the contour method
plot_potential(E=-1.25,Lz=0.60,pot=potential.MWPotential2014)

# Plot my method
plt.scatter(R_range,Zs.x)

# Flip it and reverse it
plt.scatter(R_range,-Zs.x)

# Make the plot bigger
f = plt.gcf()
f.set_size_inches(10, 5)
```

(b) In the notes, we discussed how orbits touch the zero-velocity curve at four points, but there are orbits at a given $(E,L_z)$ that only touch the zero-velocity curve at two points.

These orbits are called the thin tubes, because in 3D they look like a tube (that's thin!).

- Devise an algorithm to find this orbit for a given $(E,L_z)$ and apply it to the $(E,L_z)$ of the example orbit in (a)
- [there are various, qualitatively different ways to do this; explain why your algorithm works in general].
- Show the orbit in $(R,z)$ and $(x,y,z)$
- Discuss what you think is interesting about this orbit.

```python
R, E, Lz = 0.8, -1.25, 0.6
pot=potential.MWPotential2014

deltaE = 2.*(E-potential.evaluatePotentials(pot,R,0.) - (Lz/R)**2./2)

vR = np.sqrt(deltaE)

vT = Lz/R

z = 0

vZ = np.sqrt(2.*(E-potential.evaluatePotentials(pot,R,0.) -(Lz/R)**2./2.-vR**2./2))

phi = 0

orbit = Orbit([R,vR,vT,z,vZ,phi])

ts = np.linspace(0.,100.,2001)
orbit.integrate(ts, pot)

# Plot the contour method
plot_potential(E,Lz,pot=pot)

# Plot the orbit
orbit.plot(overplot=True)

# Make the plot bigger
f = plt.gcf()
f.set_size_inches(10, 5)
```

```python
x,y,z = [],[],[]

for t in ts:
    x.append(orbit(t).x())
    y.append(orbit(t).y())
    z.append(orbit(t).z())
plt.clf()
```

```python
#Z is 0!

plt.scatter(x,y,s=1)
plt.xlabel('x')
plt.ylabel('y')
```

To find the orbit that just touches two points of the zero-velocity curve, I set $vR = \sqrt{\Delta E}$ where $\Delta E = 2(E-\Phi(z,R) - \frac{1}{2}(Lz/R)^2)$. This essentially amounts to allowing the other velocity component, vZ to be 0, that way all the energy gets put into vR, maximizing the amount it can be.

That's pretty much it, I don't know if I should have done more of an algorithm than that. One of the things that I noticed when testing this out on other potentials is that if $\Delta E$ is negative this can be kind of an issue since then vR is unreal. In the case of this orbit, using something basic like a Kepler potential gives a slightly negative $\Delta E$ but increasing E a bit takes care of this problem. This solution may not work for every single orbit but I do think it will work in general. The original method I played around with down below in the graveyard was to treat it like an optimization problem. I let vR vary from 0 to $\Delta E$ and then took the points where the orbit came close to the zero-velocity curve and tried to minimize $\frac{1}{N}\sum_i |z_i|$. I thought this would work since clearly the orbit with two points is the one that touches the inner and outer corner of the "eye" like curve which sits around z = 0. I quickly realized though that the answer just ended up being $vR = \sqrt{\Delta E}$ every time.

What's interesting about these "tubes" is in x,y,z they look like a flat torus. These orbits completely avoid the center while oscillating from -1 to 1 in either direction.

```python

```

### Graveyard

At first I think I misinterpretted question 2 and wrote a ton of code thinking I was supposed to find the four z points at the corner and now I'm attached so I didn't want to delete it.

```python
class zero_orbit:

    def __init__(self,R,E,Lz,pot=potential.MWPotential2014):

        self.potential = pot

        self.R = R

        self.E = E

        self.Lz = Lz

        orbit, time_points = self.compute_zero_velocity_z()

        self.orbit = orbit

        self.zero_t = np.array(time_points)[:,0]

        self.zero_d = np.array(time_points)[:,1]

        self.zero_r,self.zero_z = self.rZ(orbit,time_points)

        self.optimized_value = np.min(np.sum(self.zero_z)/len(self.zero_z))

        self.deltaE = 2.*(E-potential.evaluatePotentials(pot,R,0.) - (Lz/R)**2./2)

    def compute_zero_velocity_z(self,vR=0):
        vT = self.Lz / self.R
        vZ = np.sqrt(2.*(self.E-potential.evaluatePotentials(self.potential,self.R,0.) -(self.Lz/self.R)**2./2.-vR**2./2))
        z = 0
        phi = 0

        # How close to get to the corner is acceptable
        tol = 0.05 # 10 ???

        # Calculate the orbit, I took physical units off for now since I want things to be on the same scale
        orbit = Orbit([self.R,vR,vT,z,vZ,phi])
        ts = np.linspace(0.,100.,2001)
        orbit.integrate(ts, self.potential)

        # On the zero-velocity curve when vR = vZ = 0
        zero_velocity_curve_points = []
        for t in ts :
            vr = orbit(t).vR()
            vz = orbit(t).vz()

            # How close are we to the distance velocity curve
            distance_to_curve = np.sqrt(vr**2 + vz**2)

            # If we're close to 0 save the time step and the distance
            if np.isclose(distance_to_curve, 0, atol = tol):
                zero_velocity_curve_points.append([t, distance_to_curve])

        return orbit, zero_velocity_curve_points

    def rZ(self,orbit,time_points):
        # At times in the orbit where the sqrt of vr^2 vz^2 is closest to zero, collect the R and Z.
        rs,zs = [],[]
        for t in time_points :
            rs.append(orbit(t[0]).R())#.value)
            zs.append(orbit(t[0]).z())#.value)

        return np.array(rs),np.array(zs)

    def plot_potential(self):
        levels = np.linspace(-1.25,0.,11)
        cntrcolors = ['k' for l in levels]
        cntrcolors[np.arange(len(levels))\
           [np.fabs(levels-self.E) < 0.01][0]]= '#ff7f0e'
        potential.plotPotentials(self.potential,
                         rmin=0.3,rmax=1.05,
                         zmin=-0.3,zmax=0.3,
                         nrs=101,nzs=101,
                         justcontours=True,
                         levels=levels,
                         cntrcolors=cntrcolors,
                         effective=True,
                         Lz=self.Lz)

    def optimize(self):

        deltaE = 2.*(self.E-potential.evaluatePotentials(self.potential,self.R,0.) - (self.Lz/self.R)**2./2)

        print(deltaE)

        vR_range = np.arange(0,np.sqrt(self.deltaE),0.01)

        print(len(vR_range))

        self.optimized_vals = []
        for vR in vR_range:
            vT = self.Lz/self.R
            vZ = np.sqrt(2.*(self.E-potential.evaluatePotentials(self.potential,self.R,0.) -(self.Lz/self.R)**2./2.-vR**2./2))
            z = 0
            phi = 0

            orbit, time_points = self.compute_zero_velocity_z(vR)

            zero_r,zero_z = self.rZ(orbit,time_points)

            optimized_value = np.min(np.sum(np.abs(zero_z))/len(zero_z))

            print(optimized_value)

            self.optimized_vals.append(optimized_value)

```

```python
zOrb = zero_orbit(R = 0.8, E = -1.25, Lz = 0.60)

zOrb.plot_potential()
zOrb.orbit.plot(overplot=True,use_physical=False)
plt.scatter(zOrb.zero_r,zOrb.zero_z, c='r', zorder=10)
plt.scatter(0.8,0.16966632)
f = plt.gcf()
f.set_size_inches(10, 5)
```

```python
vR = np.sqrt(zOrb.deltaE)

orbit, time_points = zOrb.compute_zero_velocity_z(vR)
```

```python
zOrb.plot_potential()
orbit.plot(overplot=True,use_physical=False)
```

```python
zOrb.plot_potential()
orbit.plot(overplot=True,use_physical=False)
```

```python
np.sqrt(E)
```

    /home/bethany/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:1: RuntimeWarning: invalid value encountered in sqrt
      """Entry point for launching an IPython kernel.

    nan

```python
vR_range = np.arange(0,np.sqrt(E))

nOrb = zero_orbit(R = 0.8, E = -1.25, Lz = 0.60,vR=0.2)
nOrb.optimized_value
```

```python
zOrb = zero_orbit(R = 0.8, E = -1.25, Lz = 0.60)

zOrb.plot_potential()
zOrb.orbit.plot(overplot=True,use_physical=False)
plt.scatter(zOrb.zero_r,zOrb.zero_z, c='r', zorder=10)

f = plt.gcf()
f.set_size_inches(10, 5)
```

```python
plt.clf()
%matplotlib inline

R, E, Lz = 0.8, -.25, 0.6

pot=potential.KeplerPotential()

deltaE = 2.*(E-potential.evaluatePotentials(pot,R,0.) - (Lz/R)**2./2)

vR = np.sqrt(np.abs(deltaE))

vT = Lz/R

z = 0

vZ = np.sqrt(2.*(E-potential.evaluatePotentials(pot,R,0.) -(Lz/R)**2./2.-vR**2./2))

phi = 0
print([R,vR,vT,z,vZ,phi])
orbit = Orbit([R,vR,vT,z,vZ,phi])

ts = np.linspace(0.,100.,1001)
orbit.integrate(ts, pot)

# Plot the contour method
#plot_potential(E,Lz,pot=pot)

# Plot the orbit
orbit.plot(overplot=True)

# Make the plot bigger
f = plt.gcf()
f.set_size_inches(10, 5)
ax = plt.gca()
ax.set_ylim(-1,1)
ax.set_xlim(1,100)
```


# Variational Principle  using Symbolic Mathematics in Python

## 1. Introduction

The variational principle tells us that we can use a trial wavefunction to solve the Schrodinger equation using the following theorem:

$${{\int {{\Psi ^*}\hat H{\rm{ }}\Psi } d\tau } \over {\int {{\Psi ^*}\Psi } d\tau }} \ge {E_0}$$

We will use Sympy to solve the particle in a box problem by guessing a trial wavefunction using variational principle

```python
import sympy as sym
```

This exercise is a bit more self-guided than the other notebooks we have done. One of the most useful things you can do is **open last week's notebook to remember the commands in sympy**. Also, remember that google is your friend:

1. [Sympy tutorial](https://docs.sympy.org/latest/tutorial/index.html)
2. [Stack Overflow](https://stackoverflow.com/search?q=sympy+)
3. [Stack Exchange](https://stackexchange.com/)

## 2. Particle in a box

The wave function that we pick for a particle in a box needs to have the following properties

1. single valued
1. normalizable
1. function and its first derivative are continuous
1. boundary condition that the wave function goes to zero at the ends of the box

Particle in a box: a is a classical particle, red is real part, blue is imaginary part.

This particle only expericnes kinetic energy between the box, so the Hamiltonian for this system is

$$\hat H = {{ - {\hbar ^2}} \over {2m}}{{{d^2}} \over {d{x^2}}} + \left\{ {\matrix{{V(x) = 0} & {0 < x < a}  \cr {V(x) = \infty } & {x < 0\text{ }{\rm{ or}}\;x > a}  \cr } } \right.$$

For our purposes, that means we can consider the Hamiltonian to be

$$\hat H = {{ - {\hbar ^2}} \over {2m}}{{{d^2}} \over {d{x^2}}}$$

as long as we keep the limits of integration to be $(0,a)$

### 2.1 Trial Wave function

Although the particle in box has a well known solution

[https://en.wikipedia.org/wiki/Particle_in_a_box](https://en.wikipedia.org/wiki/Particle_in_a_box)

(or check your favorite pchem book)

We are going to guess a trial wave function:

$$\Phi (x) = \left( {{x \over a} - {{{x^3}} \over a}} \right) + \alpha \left( {{{{x^5}} \over {{a^5}}} - {1 \over 2}\left( {{{{x^7}} \over {{a^7}}} + {{{x^7}} \over {{a^7}}}} \right)} \right)$$

### 2.2 Exercise: Variational Theorem

We are going to follow the following plan:

1. Solve for the energy of the trial wave function above

$${E_{trial}} = {{\int\limits_0^a {\Phi (x){{ - {\hbar ^2}} \over {2m}}{{{d^2}} \over {d{x^2}}}\Phi (x)dx} } \over {\int\limits_0^a {\Phi {{(x)}^2}dx} }}$$

Your answer will be a function of $ m,a,\text{and } \alpha$ We will use $\alpha$ as the parameter we vary to minimize the energy and make a new trial wave function.

2. Minimize the trial energy
We will use a first derivative of the trial energy $${d \over {d\alpha }}{E_{trial}}(\alpha )$$ to find the value of $\alpha$ that gives you the lowest energy

3. Plot your new wavefunction compared to the ground state particle in a box: $${\psi _{true}}(x) = {\left( {{2 \over a}} \right)^{1/2}}\sin {{n\pi x} \over a}$$ Plot as a function of $x/a$ from $0$ to $1$. Assuming this has $m=m_e$, and $a=a_0$ use atomic (theorist) units to plot the function.

4. Compare your trial energy to the actual energy (using atomic units)
$${E_{true}}(n = 1) = {{{\hbar ^2}{\pi ^2}} \over {2m{a^2}}}$$

```python
import sympy as sym
a,x,m=sym.symbols('a,x,m')
sym.Rational(1/2)
from sympy.physics.units import hbar
sym.init_printing()
Phi,phi,alpha,hbar=sym.symbols("Phi,phi,alpha,hbar")
alpha,phi,Phi,hbar
```

```python
import sympy as sym
phi = ((x/a)-(x**3/a**3)+ alpha*((x**5/a**5)-1/2*(x**9/a**9+x**9/a**9)))
top1 = (phi*sym.Rational(-1/2))
top2 = sym.diff(sym.diff(phi,x),x)
bottom1 = (phi**2)
```

```python
import sympy as sym
top = sym.integrate((top1*top2),(x,0,a))
bottom = sym.integrate((bottom1),(x,0,a))
top/bottom
Etrial = (top/bottom).subs(a,1)
print(Etrial)
```

    (0.309703368526898*alpha**2 + 0.311688311688312*alpha + 0.4)/(0.010207336523126*alpha**2 + 0.0355200355200355*alpha + 0.0761904761904761)

```python
import sympy as sym
sym.solveset(sym.diff(Etrial,alpha))
```

```python
Etrial.subs(alpha, -0.257747652303185)
```

```python
Etrial.subs(alpha, -4.73345207773701)
```

```python
NewWave = phi.subs(alpha, 5.02467668289965)
print(NewWave)
```

    x/a - x**3/a**3 + 5.02467668289965*x**5/a**5 - 5.02467668289965*x**9/a**9

```python
NewWave2 = NewWave.subs(a,1)
print(NewWave2)
```

    -5.02467668289965*x**9 + 5.02467668289965*x**5 - x**3 + x

```python
sym.plot(NewWave2,(x,0,1))
```

```python
sym.plot((2**1/2)*sym.sin(sym.pi*x),(x,0,1))
```

```python
sym.plot((2**1/2)*sym.sin(sym.pi*x),NewWave2,(x,0,1))
```

```python
Etrue = sym.pi**2/(2*m)
print(Etrue)
```

    pi**2/(2*m)

Your descriptions/explanations here

### 2.3 Exercise: New trial wavefunction

Determine the minimum energy of the particle in a box using a new trial wavefunction $$x^\alpha(x-a)^\alpha$$

1. Find the minimum energy, $E_{trial}$
2. Plot the new trial wavefunction and compare it to the true solution and the wavefunction you found above
3. Compare you new energy to the trial energy you found above
4. Which wavefunction is better? How do you know?

```python
trialwave = (x**alpha)*(x-a)**alpha
print(trialwave)
```

    x**alpha*(-a + x)**alpha

```python
top11 = (trialwave*sym.Rational(1/2))
top21 = sym.diff(sym.diff(trialwave,x),x)
bottom11 = (trialwave**2)
top2 = sym.integrate((top11*top21),(x,0,a))
print(top2)
```

    alpha*Integral(x**(2*alpha)*(-a + x)**(2*alpha)*(a**2*alpha - a**2 - 4*a*alpha*x + 2*a*x + 4*alpha*x**2 - 2*x**2)/(x**2*(-a + x)**2), (x, 0, a))/2

```python
bottom2 = sym.integrate((bottom11),(x,0,a))
print(bottom2)
```

    a*a**(4*alpha)*exp(2*I*pi*alpha)*gamma(2*alpha + 1)*hyper((-2*alpha, 2*alpha + 1), (2*alpha + 2,), 1)/gamma(2*alpha + 2)

```python
Etrial2 = (top2/bottom2).subs(a,1)
print(Etrial2)
```

    alpha*exp(-2*I*pi*alpha)*gamma(2*alpha + 2)*Integral(x**(2*alpha)*(x - 1)**(2*alpha)*(4*alpha*x**2 - 4*alpha*x + alpha - 2*x**2 + 2*x - 1)/(x**2*(x - 1)**2), (x, 0, 1))/(2*gamma(2*alpha + 1)*hyper((-2*alpha, 2*alpha + 1), (2*alpha + 2,), 1))

```python
sym.solveset(sym.diff(Etrial2,alpha))
```

$\displaystyle \left\{\alpha \mid \alpha \in \mathbb{C} \wedge \left(- 2 \alpha {{}_{2}F_{1}\left(\begin{matrix} - 2 \alpha, 2 \alpha + 1 \\ 2 \alpha + 2 \end{matrix}\middle| {1} \right)} \operatorname{polygamma}{\left(0,2 \alpha + 1 \right)} \int\limits_{0}^{1} \frac{x^{2 \alpha} \left(x - 1\right)^{2 \alpha} \left(4 \alpha x^{2} - 4 \alpha x + \alpha - 2 x^{2} + 2 x - 1\right)}{x^{2} \left(x - 1\right)^{2}}\, dx + 2 \alpha {{}_{2}F_{1}\left(\begin{matrix} - 2 \alpha, 2 \alpha + 1 \\ 2 \alpha + 2 \end{matrix}\middle| {1} \right)} \operatorname{polygamma}{\left(0,2 \alpha + 2 \right)} \int\limits_{0}^{1} \frac{x^{2 \alpha} \left(x - 1\right)^{2 \alpha} \left(4 \alpha x^{2} - 4 \alpha x + \alpha - 2 x^{2} + 2 x - 1\right)}{x^{2} \left(x - 1\right)^{2}}\, dx + \alpha {{}_{2}F_{1}\left(\begin{matrix} - 2 \alpha, 2 \alpha + 1 \\ 2 \alpha + 2 \end{matrix}\middle| {1} \right)} \int\limits_{0}^{1} \frac{x^{2 \alpha} \left(x - 1\right)^{2 \alpha} \left(4 x^{2} - 4 x + 2 \left(4 \alpha x^{2} - 4 \alpha x + \alpha - 2 x^{2} + 2 x - 1\right) \log{\left(x \right)} + 2 \left(4 \alpha x^{2} - 4 \alpha x + \alpha - 2 x^{2} + 2 x - 1\right) \log{\left(x - 1 \right)} + 1\right)}{x^{2} \left(x - 1\right)^{2}}\, dx - 2 i \pi \alpha {{}_{2}F_{1}\left(\begin{matrix} - 2 \alpha, 2 \alpha + 1 \\ 2 \alpha + 2 \end{matrix}\middle| {1} \right)} \int\limits_{0}^{1} \frac{x^{2 \alpha} \left(x - 1\right)^{2 \alpha} \left(4 \alpha x^{2} - 4 \alpha x + \alpha - 2 x^{2} + 2 x - 1\right)}{x^{2} \left(x - 1\right)^{2}}\, dx - \alpha \frac{d}{d \alpha} {{}_{2}F_{1}\left(\begin{matrix} - 2 \alpha, 2 \alpha + 1 \\ 2 \alpha + 2 \end{matrix}\middle| {1} \right)} \int\limits_{0}^{1} \frac{x^{2 \alpha} \left(x - 1\right)^{2 \alpha} \left(4 \alpha x^{2} - 4 \alpha x + \alpha - 2 x^{2} + 2 x - 1\right)}{x^{2} \left(x - 1\right)^{2}}\, dx + {{}_{2}F_{1}\left(\begin{matrix} - 2 \alpha, 2 \alpha + 1 \\ 2 \alpha + 2 \end{matrix}\middle| {1} \right)} \int\limits_{0}^{1} \frac{x^{2 \alpha} \left(x - 1\right)^{2 \alpha} \left(4 \alpha x^{2} - 4 \alpha x + \alpha - 2 x^{2} + 2 x - 1\right)}{x^{2} \left(x - 1\right)^{2}}\, dx\right) \Gamma\left(2 \left(\alpha + 1\right)\right) = 0 \right\} \setminus \left\{\alpha \mid \alpha \in \mathbb{C} \wedge e^{2 i \pi \alpha} \Gamma\left(2 \alpha + 1\right) {{{}_{2}F_{1}\left(\begin{matrix} - 2 \alpha, 2 \alpha + 1 \\ 2 \alpha + 2 \end{matrix}\middle| {1} \right)}}^{2} = 0 \right\}$

Your descriptions/explanations here

### 2.4 Exercise: Design your own wavefunction!

**Now you get to make your own wavefunction!**

The only guidance I would give you is that it make sense mathematically and that it include $\alpha$ so that you can minimize the energy.

Remember that $a$ and $x$ are both length units, and that trigonometric, logarithmic, and exponential functions are all unitless

Using your new wavefunction:

1. Find the minimum energy, $E_{trial}$
2. Plot the new trial wavefunction and compare it to the true solution and the wavefunction you found above
3. Compare you new energy to the trial energy you found above
4. Which wavefunction is better? How do you know?

```python
mywave = (2*x**alpha)*(x+a)**alpha
print(mywave)
```

    2*x**alpha*(a + x)**alpha

```python
top1 = (mywave*sym.Rational(-1/2))
top2 = sym.diff(sym.diff(mywave,x),x)
bottom1 = (mywave**2)
top = sym.integrate((top1*top2),(x,0,a))
bottom = sym.integrate((bottom1),(x,0,a))
top/bottom
Etrial = (top/bottom).subs(a,1)
print(Etrial)
```

```python
import sympy as sym
sym.solveset(sym.diff(Etrial,alpha))
```

```python
Etrial.subs(alpha,
```

```python
Etrial.subs(alpha,
```

```python
NewWave = phi.subs(alpha,
print(NewWave)
```

```python
NewWave2 = NewWave.subs(a,1)
print(NewWave2)
```

Your descriptions/explanations here

```python

```

# Reading Homework

Read the following sections in Kramer

- 4.2.3 Born-Oppenheimer approximation
- 4.3.2 Secular equation
- All of 4.5

For each subsection
- write down the subchapter name
- what was the most important idea
- draw an idea digram of the main idea

**Make sure to upload this to the assignment repository**

Example idea diagram:

```python

```


```python
import numpy as np
import cv2
import matplotlib.pyplot as plt
import os
import argparse
import glob
import torch
import torch.nn as nn
from torch.autograd import Variable
from DnCNN.models import DnCNN
from DnCNN.utils import *
```

# Padding

Instead of padding an input image of dimensions $(w, h)$ so that it's padded output has dimensions $(2w-1, 2h-1)$, it's preferred to find the next power of 2 that is larger than $(2w-1, 2h-1)$. This is because the discrete Fourier Transform (DFT) performs better on inputs with dimensions that factor into small primes (e.g. 2).

"The convolution of two functions with non-zero values at indices from 0 to N−1, i.e. N values, will have non-zero values at indices 0 to 2(N−1), i.e. 2N−1 values. Thus, it suffices to use a DFT of length 2N−1 to avoid cyclical overlap. However, if N is chosen so as to factor into small primes, this property is lost if a transform with length 2N−1 is performed; in this case, it is far more efficient to pad to 2N instead."

```python
""" Returns the nearest power of 2 that is larger than 2N. This optimization relies on the fact that
the Discrete Fourier Transform leverages small prime factors for speed up.
" n is the input dimension that we'd like to pad
"""
def compute_padding(n):
    return int(np.power(2, np.ceil(np.log2(2 * n - 1))))

test_n = 3
assert compute_padding(test_n) == 8, 'The padding calcuations are incorrect.'

""" Returns the input image padded so that its dimensions are the next power of 2 larger than 2w and 2h.
" input is the image to pad
"""
def pad(input):
    h, w = compute_padding(input.shape[0]), compute_padding(input.shape[1])
    padded = cv2.copyMakeBorder(input, h // 2, h // 2, w // 2, w // 2, cv2.BORDER_CONSTANT)
    return padded

test_img = cv2.imread('./img/psf.tif', cv2.IMREAD_GRAYSCALE)
assert pad(test_img).shape == (5296, 5696), 'The padded dimensions are incorrect.'

""" Returns the cropped version of a padded image
" input is the padded image
" original_shape is the original dimensions of an image padded with pad() """
def crop(input, original_shape):
    diff = np.subtract(input.shape, original_shape)
    assert(np.all(diff >= 0))
    h, w = original_shape
    h_pad, w_pad = diff // 2
    return input[h_pad:(input.shape[0] - h_pad), w_pad:(input.shape[1] - w_pad)]

assert crop(pad(test_img), test_img.shape).shape == test_img.shape, 'Cropped dimensions are incorrect.'

""" Resizes the input image
" img is the image to resize
" factor is the resizing factor
"""
def resize(img, factor):
    num = int(-np.log2(factor))
    for i in range(num):
        img = 0.25 * (img[::2,::2,...] + img[1::2,::2,...] + img[::2,1::2,...] + img[1::2,1::2,...])
    return img
```

# Denoiser
Plug and play alternating direction method of multipliers (PnP ADMM) uses a denoising algorithm as the proximal update to the dual variable. Here, the denoising convolution neural network (DnCNN) is used as the denoiser. This is preferred over the block matching 3D (BM3D) algorithm for speed and also performance. However, it lacks explicit control over the denoising strength like BM3D. Additionally, the network can make use of GPU resources thanks for PyTorch.

```python
class Denoiser():
    def __init__(self):
        n_layers = 17
        net = DnCNN(channels=1, num_of_layers=n_layers)

        device_ids = [0]
        self.model = nn.DataParallel(net, device_ids=device_ids).cuda()
        model_path = './DnCNN/logs/DnCNN-S-15/net.pth'
        self.model.load_state_dict(torch.load(model_path))
        self.model.eval()

    def denoise(self, img):
        img = np.float32(img[:, :, 0])
        img = np.expand_dims(img, 0)
        img = np.expand_dims(img, 1)
        noisy = Variable(torch.Tensor(img).cuda())
        with torch.no_grad():
            result = torch.clamp(noisy - self.model(noisy), 0., 1.)
        result = torch.squeeze(result.data.cpu(), 0).permute(1, 2, 0)
        return result.numpy()

# Run this to visualize the denoiser's performance.
# dncnn = Denoiser()
# noisy = cv2.imread('noisy.jpg')
# dncnn.denoise(noisy)
```

# PnP ADMM

Image reconstruction from DiffuserCam images and a measured PSF can be stated as the following optimization problem:

\begin{equation}
\mathcal{L}_{p}(x,z,\lambda) = \underset{x,z}{argmin}\frac{1}{2}\|Ax-b\|^{2}_{2} - log(p(z)) + \frac{\rho}{2}\|x-z+\frac{\lambda}{\rho}\|^{2}_{2}
\end{equation}
, where $\mathcal{L}_{p}(x,z,\lambda)$ is the scaled form of the augmented Lagrangian, $A$ is the point squared function (PSF), x is the reconstruction we'd like to recover, b is the DiffuserCam image, $\lambda$ is the Langrange multiplier, and $\rho$ is an arbitrarily chosen constant.

## Proximal operator updates
The iterative updates to $x, z, u$ are computed below:

\begin{equation}
x^{k+1} := 0 = \nabla{x} \mathcal{L}_{p}(x,z,\lambda) := (A^{T}A + \rho)^{-1}(A^{T}b + \rho{}z^{k} - \lambda^{k})
\end{equation}

\begin{equation}
z^{k+1} = \mathcal{D}(x^{k+1} + u^{k}), u = \frac{\lambda}{\rho}
\end{equation}

\begin{equation}
u^{k+1} := 0 = \nabla_{u} \mathcal{L}_{p}(x,z,\lambda) := u + x^{k+1} - z^{k+1}
\end{equation}

## Optimization

Explicitly computing $$A^{T}A, A^{T}b$$ in the proximal update to $x$ can be expensive, especially when working on large images (and hence matrices).

Because the forward operator $A$ can be expressed as a linear convolution, we can efficiently compute the $x$ update as below:

\begin{equation}
x^{k+1} = \mathcal{F}^{-1}(\frac{\overline{\mathcal{F}(A)}\mathcal{F}(b) + \rho \mathcal{F}(z - u)}{|\mathcal{F}(A)|^{2} + \rho})
\end{equation}

Note that the image $b$ needs to be padded (and FFT shifted) before convolution. The final output must also be cropped to remove the excess padding.

```python
""" Solve argminx Ax = b.
" A is the 2D PSF image
" b is the measurement
"""
def pnp_admm(b, A, rho=1, num_iter=5, debug=False):
    orig_dims = A.shape
    dncnn = Denoiser()
    # Sets the initial state of optimization variables x, z, u.
    x = pad(np.ones(b.shape) * 0.5)
    z = x.copy()
    u = np.zeros(x.shape)

    # Precomputes matrix math.
    AtA = np.abs(np.fft.fft2(pad(A))) ** 2
    Atb = cv2.filter2D(pad(b), -1, A, borderType=cv2.BORDER_WRAP)

    for i in np.arange(num_iter):
        # Proximal update for x := x^{t+1} = (AtA + rho)^{-1} (Atb + rho * z - lambda)
        num = np.fft.fft2(Atb + rho * (z - u))
        denom = (AtA + rho)
        x = np.real((np.fft.ifft2(num / denom)))
        # Proximal update for z := D(x+u)
        z = dncnn.denoise(np.expand_dims(x + u, axis=2))
        z = np.squeeze(z)
        # Proximal update for u.
        u += (x - z)
        if debug:
            plt.figure()
            plt.title(f'PnP ADMM iteration {i}')
            plt.imshow(crop(np.clip(x, 0, 1) * 255, b.shape), cmap='gray')
            plt.show()
    # Clamps the RGB image between 0 and 1.
    x = np.clip(x, 0, 1)
    return crop(x, b.shape)
```

# Reconstruction

The code below loads the PSF and DiffuserCam image and outputs the PnP ADMM reconstructed image.

```python
""" Utility function to threshold the background in the PSF image. This is helpful because the PSF measurements in our lab setup are noisy.
" psf is the PSF image to filter
"""
def filter_psf(psf, threshold=0.3):
    psf[psf < (threshold * 255)] = 0
    return psf

""" Reconstructs the scene.
" psf is the filename of the PSF iamge
" imgs is an array of DiffuserCam images
" downsample_factor is the amount to downsample the PSF and DiffuserCam images
" filter determines whether to filter the input PSF image
" show_plots determines whether image plots will be displayed
"""
def reconstruct(psf, imgs, downsample_factor=(1.0/8.0), rho=1, num_iter=3, filter=True, show_plots=False, rgb=True):
    # Read in the PSF image.
    # psf = cv2.imread(psf, cv2.IMREAD_GRAYSCALE).astype(float)
    if filter:
        psf = filter_psf(psf)
    psf = resize(psf, downsample_factor) / 255

    outputs = []
    # Iterate through the images and reconstruct them one at a time.
    for img in imgs:
        # img = cv2.imread(img).astype(float) if rgb else cv2.imread(img, cv2.IMREAD_GRAYSCALE).astype(float)
        img = resize(img, downsample_factor) / 255
        if rgb:
            r = pnp_admm(img[:,:,0], psf, rho, num_iter)
            g = pnp_admm(img[:,:,1], psf, rho, num_iter)
            b = pnp_admm(img[:,:,2], psf, rho, num_iter)
            reconstruction = np.dstack((r, g, b))
        else:
            reconstruction = pnp_admm(img, psf, rho, num_iter)
        # reconstruction *= 255
        outputs.append(reconstruction)
        if show_plots:
            plt.figure(figsize=(30, 12))
            plt.subplot(1, 3, 1)
            plt.title('PSF')
            plt.imshow(psf, cmap='gray')
            plt.subplot(1, 3, 2)
            plt.title('DiffuserCam image')
            plt.imshow(img, cmap='gray') if not rgb else plt.imshow(img)
            plt.subplot(1, 3, 3)
            plt.title(f'Reconstruction (rho={rho}, n_iter={num_iter})')
            plt.imshow(reconstruction, cmap='gray') if not rgb else plt.imshow(reconstruction)
        # plt.show()
    return outputs

# silence = reconstruct('img/1210/psf_1.tiff', ['img/1210/marker.tiff', 'img/1210/textbook.tiff', 'img/1210/textbook2.tiff', 'img/1210/box.tiff', 'img/1210/thor.tiff', 'img/1210/hand.tiff'], filter=True, show_plots=True, rgb=True)

## TAPE PSF:
## Marker images (OK, but not great)
# c = reconstruct('img/1210/psf_1.tiff', ['img/1210/double_marker.tiff', 'img/1210/marker.tiff', 'img/1210/marker_close.tiff'], filter=True, show_plots=True, rgb=True)
## Textbook images (Not great)
# c = reconstruct('img/1210/psf_1.tiff', ['img/1210/textbook2.tiff', 'img/1210/textbook.tiff', 'img/1210/textbook_last.tiff'], filter=True, show_plots=True, rgb=True)
## Hand images (terrible)
# c = reconstruct('img/1210/psf_1.tiff', ['img/1210/hand.tiff', 'img/1210/hand_best.tiff', 'img/1210/hand2.tiff', 'img/1210/hand_last.tiff', 'img/1210/hand_bright.tiff'], filter=True, show_plots=True, rgb=True)
## Thor images (barely recognizable)
# c = reconstruct('img/1210/psf_1.tiff', ['img/1210/thor.tiff', 'img/1210/thor2.tiff', 'img/1210/thor_close.tiff', 'img/1210/thor_last.tiff', 'img/1210/thor_last2.tiff'], filter=True, show_plots=True, rgb=True)
## Chart images (Terrible)
# c = reconstruct('img/1210/psf_1.tiff', ['img/1210/chart.tiff', 'img/1210/chart_far.tiff'], filter=True, show_plots=True, rgb=True)
## Misc
# c = reconstruct('img/1210/psf_1.tiff', ['img/1210/small_text.tiff', 'img/1210/box.tiff', 'img/1210/bright.tiff', 'img/1210/dim.tiff'], filter=True, show_plots=True, rgb=True)

# b = reconstruct('img/1210/psf.tiff', ['img/1210/marker.tiff'], filter=True, show_plots=True, rgb=True)
# a = reconstruct('img/1210/psf_10.tiff', ['img/1210/marker.tiff'], filter=True, show_plots=True, rgb=True)
# a = reconstruct('img/1210/psf_1_close.tiff', ['img/1210/marker.tiff', 'img/1210/marker_close.tiff'], filter=True, show_plots=True, rgb=True)
# reconstruct('img/batch2/psf.tiff', ['img/batch2/marker.tiff', 'img/batch2/grapes.tiff'], show_plots=True)
# reconstruct('img/psf.tif', ['img/measurement.tif'], filter=False, show_plots=True, rgb=False)
# reconstruct('img/batch2/psf.tiff', ['img/batch2/marker.tiff'], filter=False, show_plots=True, rgb=True)
# reconstruct('img/batch2/psf.tiff', ['img/batch2/marker.tiff'], filter=True, show_plots=True, rgb=True, downsample_factor=(1.0/4))
# reconstruct('img/batch2/psf.tiff', ['img/batch2/marker_edit.tiff'], filter=True, show_plots=True, rgb=True)
# reconstruct('img/batch1/psf.tiff', ['img/batch1/thor_box.tiff'], filter=True, show_plots=True)
# reconstruct('img/batch2/psf.tiff', ['img/batch2/thor.tiff'], filter=True, show_plots=True)
```

```python
def loadData(show_im=True):
    psf = cv2.imread('img/psf.tif', cv2.IMREAD_GRAYSCALE).astype(float)
    data = cv2.imread('img/measurement.tif', cv2.IMREAD_GRAYSCALE).astype(float)
    print(np.max(psf), np.max(data))

    """In the picamera, there is a non-trivial background
    (even in the dark) that must be subtracted"""
    bg = np.mean(psf[5:15,5:15])
    psf -= bg
    data -= bg
    print(np.max(psf), np.max(data))

    """Resize to a more manageable size to do reconstruction on.
    Because resizing is downsampling, it is subject to aliasing
    (artifacts produced by the periodic nature of sampling). Demosaicing is an attempt
    to account for/reduce the aliasing caused. In this application, we do the simplest
    possible demosaicing algorithm: smoothing/blurring the image with a box filter"""

    def resize(img, factor):
        num = int(-np.log2(factor))
        for i in range(num):
            img = 0.25*(img[::2,::2,...]+img[1::2,::2,...]+img[::2,1::2,...]+img[1::2,1::2,...])
        return img

    f = 0.25
    psf = resize(psf, f)
    data = resize(data, f)

    """Now we normalize the images so they have the same total power. Technically not a
    necessary step, but the optimal hyperparameters are a function of the total power in
    the PSF (among other things), so it makes sense to standardize it"""

    psf /= np.linalg.norm(psf.ravel())
    data /= np.linalg.norm(data.ravel())
    print(np.max(psf), np.max(data))

    if show_im:
        fig1 = plt.figure()
        plt.imshow(psf, cmap='gray')
        plt.title('PSF')
#         display.display(fig1)
        fig2 = plt.figure()
        plt.imshow(data, cmap='gray')
        plt.title('Raw data')
#         display.display(fig2)
    return psf, data

# Simulation.
downsample = 1.0/8.0

# Loads image and PSF
psf, x = loadData(show_im=False)
y = pnp_admm(x, psf, 1, 3)
y = cv2.resize(y, (x.shape[1], x.shape[0]))
print(psf.shape, x.shape, y.shape)
plt.figure()
plt.title("Clear image")
plt.imshow(y)
```

```python
# Simulation.
downsample = 1.0/8.0

# Loads image and PSF
psf = cv2.imread('img/psf.tif', cv2.IMREAD_GRAYSCALE).astype(float)
x = cv2.imread('img/measurement.tif', cv2.IMREAD_GRAYSCALE).astype(float)
y = reconstruct(psf, [x], downsample_factor=downsample, filter=False, show_plots=False, rgb=False)
y[0] = cv2.resize(y[0], (x.shape[1], x.shape[0]))
print(psf.shape, x.shape, y[0].shape)
plt.figure()
plt.title("Clear image")
plt.imshow(y[0])
# plt.figure()
# plt.title("PSF")
# plt.imshow(psf)

# Applies PSF kernel to blur the image---this generates synthetic DiffuserCam image.
# psf = resize(psf, downsample)
x_tilde = cv2.filter2D(y[0], -1, psf, borderType=cv2.BORDER_WRAP)
# x_tilde = cv2.resize(x_tilde, (x.shape[1], x.shape[0]))
plt.figure()
plt.title("Artificially blurred image")
plt.imshow(x_tilde)
plt.figure()
plt.title("Original blurred image")
plt.imshow(x)

# TODO: Figure out the difference between these two. Might need to clamp?

# Applies added white Gaussian noise to the image.

# Runs PnP-ADMM to perform blind deconvolution. This tests the performance of the algorithm.
print(psf.shape, x_tilde.shape)
# y_tilde = reconstruct(psf, [x_tilde], filter=False, show_plots=False, rgb=False)
# plt.figure()
# plt.title("Artificial reconstruction")
# plt.imshow(y_tilde[0])
```

```python
""" Artificially blurs a clear image using the PSF, then does blind deconvolution to retrieve the clear image back """
def simulation(psf, img):
    blur = cv2.filter2D(img, -1, psf, borderType=cv2.BORDER_WRAP)
    print(blur.shape)
    plt.figure()
    plt.title("Synthetically blurred image")
    plt.imshow(blur)
    y = reconstruct(psf, [blur], filter=False, show_plots=False, rgb=False)
    plt.figure()
    plt.title("Reconstructed image")
    plt.imshow(y[0])
    plt.figure()
    plt.title("Original image")
    plt.imshow(img)
    print(img.shape, y[0].shape)

psf = cv2.imread('img/psf.tif', cv2.IMREAD_GRAYSCALE).astype(float)
y = cv2.imread('img/test/USAF.jpg', cv2.IMREAD_GRAYSCALE).astype(float)
y = cv2.resize(y, (psf.shape[1], psf.shape[0]))
simulation(psf, y)
```


# Implementation details: deriving expected moment dynamics
$$
\def\n{\mathbf{n}}
\def\x{\mathbf{x}}
\def\N{\mathbb{\mathbb{N}}}
\def\X{\mathbb{X}}
\def\NX{\mathbb{\N_0^\X}}
\def\C{\mathcal{C}}
\def\Jc{\mathcal{J}_c}
\def\DM{\Delta M_{c,j}}
\newcommand\diff{\mathop{}\!\mathrm{d}}
\def\Xc{\mathbf{X}_c}
\newcommand{\muset}[1]{\dot{\{}#1\dot{\}}}
$$
This notebook walks through what happens inside `compute_moment_equations()`.
We restate the algorithm outline, adding code snippets for each step.
This should help to track down issues when, unavoidably, something fails inside `compute_moment_equations()`.

```python
# initialize sympy printing (for latex output)
from sympy import init_printing
init_printing()

# import functions and classes for compartment models
from compartor import *
from compartor.compartments import ito, decomposeMomentsPolynomial, getCompartments, getDeltaM, subsDeltaM, get_dfMdt_contrib
```

We only need one transition class.
We use "coagulation" from the coagulation-fragmentation example

```python
D = 1 # number of species
x = Content('x')
y = Content('y')

transition_C = Transition(Compartment(x) + Compartment(y), Compartment(x + y), name = 'C')
k_C = Constant('k_C')
g_C = 1
Coagulation = TransitionClass(transition_C, k_C, g_C)

transition_classes = [Coagulation]
display_transition_classes(transition_classes)
```

$\displaystyle \begin{align} \left[x\right] + \left[y\right]&\overset{h_{C}}{\longrightarrow}\left[x + y\right] && h_{C} = \frac{k_{C} \left(n{\left(y \right)} - \delta_{x y}\right) n{\left(x \right)}}{\delta_{x y} + 1} \end{align}$

$$
\def\n{\mathbf{n}}
\def\x{\mathbf{x}}
\def\N{\mathbb{\mathbb{N}}}
\def\X{\mathbb{X}}
\def\NX{\mathbb{\N_0^\X}}
\def\C{\mathcal{C}}
\def\Jc{\mathcal{J}_c}
\def\DM{\Delta M_{c,j}}
\newcommand\diff{\mathop{}\!\mathrm{d}}
\def\Xc{\mathbf{X}_c}
\newcommand{\muset}[1]{\dot{\{}#1\dot{\}}}
$$

For a compartment population $\n \in \NX$ evolving stochastically according to stoichiometric equations from transition classes $\C$, we want to find an expression for
$$
\frac{\diff}{\diff t}\left< f(M^\gamma, M^{\gamma'}, \ldots) \right>
$$
in terms of expectations of population moments $M^\alpha, M^{\beta}, \ldots$

```python
fM = Moment(0)**2
display(fM)
```

### (1)
From the definition of the compartment dynamics, we have
$$
\diff M^\gamma = \sum_{c \in \C} \sum_{j \in \Jc} \DM^\gamma \diff R_{c,j}
$$
We apply Ito's rule to derive
$$
\diff f(M^\gamma, M^{\gamma'}, \ldots) = \sum_{c \in \C} \sum_{j \in \Jc}
    \left(
        f(M^\gamma + \DM^\gamma, M^{\gamma'} + \DM^{\gamma'}, \ldots)
        - f(M^\gamma, M^{\gamma'}, \ldots)
    \right) \diff R_{c,j}
$$

Assume, that $f(M^\gamma, M^{\gamma'}, \ldots)$ is a polynomial in $M^{\gamma^i}$ with $\gamma^i \in \N_0^D$.

Then $\diff f(M^\gamma, M^{\gamma'}, \ldots)$ is a polynomial in $M^{\gamma^k}, \DM^{\gamma^l}$ with $\gamma^k, \gamma^l \in \N_0^D$, that is,
$$
\diff f(M^\gamma, M^{\gamma'}, \ldots) = \sum_{c \in \C} \sum_{j \in \Jc}
    \sum_{q=1}^{n_q} Q_q(M^{\gamma^k}, \DM^{\gamma^l})
    \diff R_{c,j}
$$
where $Q_q(M^{\gamma^k}, \DM^{\gamma^l})$ are monomials in $M^{\gamma^k}, \DM^{\gamma^l}$.

```python
dfM = ito(fM)
dfM
```

### (2)
Let's write $Q_q(M^{\gamma^k}, \DM^{\gamma^l})$ as
$$
Q_q(M^{\gamma^k}, \DM^{\gamma^l}) = k_q \cdot \Pi M^{\gamma^k} \cdot \Pi M^{\gamma^k}
$$
where $k_q$ is a constant,
  $\Pi M^{\gamma^k}$ is a product of powers of $M^{\gamma^k}$, and
  $\Pi \DM^{\gamma^l}$ is a product of powers of $\DM^{\gamma^l}$.

Analogous to the derivation in SI Appendix S.3, we arrive at the expected moment dynamics
$$
\frac{\diff\left< f(M^\gamma, M^{\gamma'}, \ldots) \right>}{\diff t} =
    \sum_{c \in \C} \sum_{q=1}^{n_q} \left<
        \sum_{j \in \Jc} k_q \cdot \Pi M^{\gamma^k} \cdot \Pi \DM^{\gamma^k} \cdot h_{c,j}(\n)
    \right>
$$

```python
monomials = decomposeMomentsPolynomial(dfM)
monomials
```

### (3)
Analogous to SI Appendix S.4, the contribution of class $c$, monomial $q$ to the expected dynamics of $f(M^\gamma, M^{\gamma'}, \ldots)$ is
$$
\begin{align}
\frac{\diff\left< f(M^\gamma, M^{\gamma'}, \ldots) \right>}{\diff t}
    &= \left<
        {\large\sum_{j \in \Jc}} k_q \cdot \Pi M^{\gamma^k} \cdot \Pi \DM^{\gamma^l} \cdot h_{c,j}(\n)
    \right>
    \\
    &= \left<
        {\large\sum_{\Xc}} w(\n; \Xc) \cdot k_c \cdot k_q \cdot \Pi M^{\gamma^k} \cdot g_c(\Xc) \cdot
        \left<
            \Pi \DM^{\gamma^l} \;\big|\; \Xc
        \right>
    \right>
\end{align}
$$

```python
c = 0 # take the first transition class
q = 1 # ... and the second monomial

tc = transition_classes[c]
transition, k_c, g_c, pi_c = tc.transition, tc.k, tc.g, tc.pi
(k_q, pM, pDM) = monomials[q]
```

First we compute the expression
$$
l(\n; \Xc) = k_c \cdot k_q \cdot \Pi(M^{\gamma^k}) \cdot g_c(\Xc) \cdot
        \left<
            \Pi \DM^{\gamma^l} \;\big|\; \Xc
        \right>
$$
We start by computing the $\DM^{\gamma^l}$ from reactants and products of the transition ...

```python
reactants = getCompartments(transition.lhs)
products = getCompartments(transition.rhs)
DM_cj = getDeltaM(reactants, products, D)
DM_cj
```

... and then substituting this expression into every occurence of $\DM^\gamma$ in `pDM` (with the $\gamma$ in `DM_cj` set appropriately).

```python
pDMcj = subsDeltaM(pDM, DM_cj)
print('pDM = ')
display(pDM)
print('pDMcj = ')
display(pDMcj)
```

Then we compute the conditional expectation of the result.

```python
cexp = pi_c.conditional_expectation(pDMcj)
cexp
```

Finally we multiply the conditional expectation with the rest of the terms:
* $k_c$, and $g_c(\Xc)$ from the specification of `transition[c]`, and
* $k_q$, and $\Pi(M^{\gamma^k})$ from `monomials[q]`.

```python
l_n_Xc = k_c * k_q * pM * g_c * cexp
l_n_Xc
```

### (4)
Let's consider the expression $A = \sum_{\Xc} w(\n; \Xc) \cdot l(\n; \Xc)$ for the following cases of reactant compartments:
$\Xc = \emptyset$,
$\Xc = \muset{\x}$, and
$\Xc = \muset{\x, \x'}$.

(1) $\Xc = \emptyset$:

Then $w(\n; \Xc) = 1$, and
$$
    A = l(\n)
$$

(2) $\Xc = \muset{\x}$:

Then $w(\n; \Xc) = \n(\x)$, and
$$
    A = \sum_{\x \in \X} \n(\x) \cdot l(\n; \muset{\x})
$$

(3) $\Xc = \muset{\x, \x'}$:

Then
$$
    w(\n; \Xc) = \frac{\n(\x)\cdot(\n(\x')-\delta_{\x,\x'})}
    {1+\delta_{\x,\x'}},
$$
and
$$
\begin{align}
    A &= \sum_{\x \in \X} \sum_{\x' \in \X}
            \frac{1}{2-\delta_{\x,\x'}}
            \cdot w(\n; \Xc) \cdot l(\n; \muset{\x, \x'}) \\
      &= \sum_{\x \in \X} \sum_{\x' \in \X}
            \frac{\n(\x)\cdot(\n(\x')-\delta_{\x,\x'})}{2}
            \cdot l(\n; \muset{\x, \x'}) \\
      &= \sum_{\x \in \X} \sum_{\x' \in \X}
            \n(\x)\cdot\n(\x') \cdot \frac{1}{2}l(\n; \muset{\x, \x'})
         \: - \:
         \sum_{\x \in \X}
            \n(\x) \cdot \frac{1}{2}l(\n; \muset{\x, \x})
\end{align}
$$

### (5)
Now let
$$
l(\n; \Xc) = k_c \cdot k_q \cdot \Pi(M^{\gamma^k}) \cdot g_c(\Xc) \cdot
        \left<
            \Pi \DM^{\gamma^l} \;\big|\; \Xc
        \right>
$$

Plugging in the concrete $\gamma^l$ and expanding, $l(\n; \Xc)$ is a polynomial in $\Xc$.

Monomials are of the form $k \x^\alpha$ or $k \x^\alpha \x'^\beta$ with $\alpha, \beta \in \N_0^D$.
(Note that occurences of $\Pi M^{\gamma^k}$ are part of the constants $k$.)

Consider again the different cases of reactant compartments $\Xc$:

(1) $\Xc = \emptyset$:
$$
\frac{\diff\left< f(M^\gamma, M^{\gamma'}, \ldots) \right>}{\diff t}
    = \left<l(\n)\right>
$$

(2) $\Xc = \muset{\x}$:
$$
\frac{\diff\left< f(M^\gamma, M^{\gamma'}, \ldots) \right>}{\diff t}
    = \left<R(l(\n; \muset{\x})\right>
$$
where $R$ replaces all $k \x^\alpha$ by $k M^\alpha$.

(3) $\Xc = \muset{\x, \x'}$:
$$
\frac{\diff\left< f(M^\gamma, M^{\gamma'}, \ldots) \right>}{\diff t}
    = \frac{1}{2}\left<R'(l(\n; \muset{\x, \x'})\right>
    \: - \:
    \frac{1}{2}\left<R(l(\n; \muset{\x, \x})\right>
$$
where $R'$ replaces all $k \x^\alpha \X'^\beta$ by $k M^\alpha M^\beta$,
and again $R$ replaces all $k \x^\alpha$ by $k M^\alpha$.

All this (the case destinction and replacements) is done in the function `get_dfMdt_contrib()`.

```python
dfMdt = get_dfMdt_contrib(reactants, l_n_Xc, D)
dfMdt
```

### (6)
Finally, sum over contributions from all $c$, $q$ for the total
$$
\frac{\diff\left< f(M^\gamma, M^{\gamma'}, \ldots) \right>}{\diff t}
$$


### 1. Vector Spaces

MML textbook에서는, Vector spaces를 다음과 같이 정의하고 있습니다.

"A real-valued vector space $V = (v,+, \cdot)$ is a set $v$ with two operations"
$$ + : V x V = V $$
$$ \cdot : R x V = V $$

Space(공간)에 대한 정의를 Group(군)으로 나타냈습니다.

Group은 **'closure', 'associativity', 'neutral element', 'inverse'** 4가지 성질을 만족해야 하므로
벡터 공간 내의 벡터들도 이러한 성질을 만족해야 합니다!

조금 더 쉽게 풀어보자면,

**Vector spaces** : N차원 내 존재하는 벡터들의 선형결합연산( $+$ : Addition, $\cdot$ : Multiplication by scalar )으로 만들어낼 수 있는 조합들의 Group으로 이해할 수 있을 것 같네요!

정말 간단하게는.. n차원 내 존재하는 모든 벡터들이 벡터공간을 이룬다고 쉽게 생각할 수 있겠습니다. 개념적으로는 선형결합연산으로 이룰 수 있는 군(group)으로 표현된 것일 뿐이죠. **물론, zero-vector도 벡터 공간 내에 포함되겠지요.** 선형결합연산으로 zero-vector를 포함하고 있으니까요!

### 2. Vector subspaces

Vector subspaces는 다음과 같이 정의합니다.

"$V = (v, +, \cdot)$, $U \subseteq V$ 일 때, $U(u, +, \cdot)$ 은 **Vector subspace이다.**"

"**Vector subspaces** : 벡터 공간 내 벡터들의 선형결합연산( Addition, Multiplication by scalar )으로 만들어낼 수 있는 작은 벡터공간" 정도로 이해할 수 있을 것 같습니다.

Vector subspaces는 정의에서도 느낄 수 있듯,

예를 들어, 2차원 실좌표공간($R^{2}$) 내에선,

1. 원점(origin, (0,0)).
2. 원점을 지나는 직선
3. $R^{2}$ 상 표현 가능한 2-dimensional vector

이 Vector subspace가될 수 있습니다. **주의할 것은 원점은 모든 부분공간의 부분집합이 됩니다. 원점이 포함되지 않는 부분공간(subspace)란 존재하지 않습니다.** 원점이 존재하지 않는다면, 그 부분공간은 원점에 대해 닫혀있지 않아 Group이 될 수 없기 때문입니다.(Group의 "Closure" 성질 위반, 부분공간이 되기 위해선, $+$, $\cdot$ 연산의 결과가 그 부분공간 내에 포함되어 있어야 하는 데, zero-vector가 없다면, 그 부분공간이 zero-vector가 연산의 결과가 되는 경우를 품지 못하는 경우가 발생하기 때문입니다.)

다음과 같은 직선은 부분공간이 아닙니다. 물론 2차원 공간 내 존재하는 벡터의 선형결합으로 이루어진 '공간(space)'이지만, zero-vector를 지나지 않기(포함하지 않기) 때문에, 모든 선형결합 연산의 결과에 닫혀(closed)있지 않아 Vector subspace로 볼 수 없습니다!

```python
%matplotlib inline

plt.xlim(-1.4, 4.4)
plt.ylim(-0.6, 5.2)
plt.plot([-2,-1,0,1,2,3,4],[0,1,2,3,4,5,6])
plt.plot(0, 0, 'ro', ms=10)
plt.title("Not Subspace",fontsize=15)
plt.show()
```

또 다른 예로, 3차원 실좌표공간($R^{3})을 생각해보자면,

1. 원점(origin, (0,0)
2. 원점을 지나는 직선
3. 원점을 갖는 평면
4. $R^{3}$상 표현 가능한 3-dimensional vector

가 Vector subspace가 될 수 있습니다.

### 3. Matrix subspace (column space)

벡터 외, 행렬에도 역시 부분공간(subspace)가 존재합니다.

벡터처럼, 선형결합 연산(Addition, Multiplication by scalar)으로 조합해낼 수 있는 공간입니다. 단지, 행렬의 column Vector가 만들어낸다는 차이점이 있을 뿐입니다.

행렬 내 column들의 선형결합으로 만들 수 있는 subspace가 바로 "column space"이자 Matrix의 subspace입니다!

예를 들면, 아래와 같은 행렬의 경우

$$
\begin{align}
\begin{bmatrix}
10 & 11 & 12\\
11 & 12 & 13\\
12 & 13 & 14\\
\end{bmatrix}
\end{align}
$$

$$
\begin{align}
c_{1}
\begin{bmatrix}
10\\
11\\
12\\
\end{bmatrix}
+
c_{2}
\begin{bmatrix}
11\\
12\\
13\\
\end{bmatrix}
+
c_{3}
\begin{bmatrix}
12\\
13\\
14\\
\end{bmatrix}
\end{align}
$$

와 같이 각 열벡터들의 선형결합으로 나타낼 수 있고, 이것이 곧 column space, subspace of Matrix를 이루게 된다.

1.Study/2. with computer/1.Math_code/5. Mathematics for Machine Learning/2. Linear Algebra_4. Vector Spaces.ipynb


```python
import os, sys
import h5py
import numpy as np
from scipy.io import loadmat
import cv2
import matplotlib
%matplotlib inline
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from numpy import matrix as mat
from sympy import *
from numpy import linalg as la
```

```python
def getFx(para, frame):   # para中为(13*frame+3)的一套参数，frame传进来是为了确定循环次数
    #先写出参数表达式，ABDCDEF六个点的齐次坐标
    K = Matrix([[1149.67569986785, 0.0, 508.848621645943],
                     [0.0, 1147.59161666764, 508.064917088557],
                     [0.0, 0.0, 1.0]])
    r11, r12, r13, r14, r21, r22, r23, r24, r31, r32, r33 = symbols('r11 r12 r13 r14 r21 r22 r23 r24 r31 r32 r33')
    Rt = Matrix([[r11, r12, r13, r14], [r21, r22, r23, r24], [r31, r32, r33, 1]])
    a, b, c, th, al = symbols('a b c th al')
    ua, va, wa, ub, vb, wb, uc, vc, wc, ud, vd, wd, ue, ve, we, uf, vf, wf = symbols('ua va wa ub vb wb uc vc wc ud vd wd ue ve we uf vf wf')
    f = Symbol('f')
    XA = Matrix([[-a * c * cos(th) * cos(al)], [c-a * c * sin(th)], [-a * c *  cos(th) * sin(al)], [1]])
    XB = Matrix([[0], [c], [0], [1]])
    XC = Matrix([[a * c * cos(th) * cos(al)], [c+a * c * sin(th)], [a * c * cos(th) * sin(al)], [1]])
    XD = Matrix([[-b * c], [0], [0], [1]])
    XE = Matrix([[0], [0], [0], [1]])
    XF = Matrix([[b * c], [0], [0], [1]])
    ua, va, wa = K[0,:] * (Rt * XA), K[1,:] * (Rt * XA), K[2,:] * (Rt * XA)
    ub, vb, wb = K[0,:] * (Rt * XB), K[1,:] * (Rt * XB), K[2,:] * (Rt * XB)
    uc, vc, wc = K[0,:] * (Rt * XC), K[1,:] * (Rt * XC), K[2,:] * (Rt * XC)
    ud, vd, wd = K[0,:] * (Rt * XD), K[1,:] * (Rt * XD), K[2,:] * (Rt * XD)
    ue, ve, we = K[0,:] * (Rt * XE), K[1,:] * (Rt * XE), K[2,:] * (Rt * XE)
    uf, vf, wf = K[0,:] * (Rt * XF), K[1,:] * (Rt * XF), K[2,:] * (Rt * XF)
    #根据每一帧的循环，提取出Rt的参数，K是公用的，代入参数写出3D坐标，并计算出u/w，v/w
    #写成f的形式，即按照六个点分块，每块里面有M帧
    getfx = mat(np.zeros((6*frame*2,1)))
    for i in range(6):
        for j in range(frame):
            if i == 0 :
                f = Matrix([ua/wa, va/wa])
            elif i == 1 :
                f = Matrix([ub/wb, vb/wb])
            elif i == 2 :
                f = Matrix([uc/wc, vc/wc])
            elif i == 3 :
                f = Matrix([ud/wd, vd/wd])
            elif i == 4 :
                f = Matrix([ue/we, ve/we])
            else:
                f = Matrix([uf/wf, vf/wf])
            f_value = f.subs({r11:para[13*j], r12:para[13*j+1], r13:para[13*j+2], r14:para[13*j+3],
                        r21:para[13*j+4], r22:para[13*j+5], r23:para[13*j+6], r24:para[13*j+7],
                        r31:para[13*j+8], r32:para[13*j+9], r33:para[13*j+10], th:para[13*j+11],
                        al:para[13*j+12], a:para[-3], b:para[-2], c:para[-1]})
            getfx[i*frame*2+j*2] = f_value[0]
            getfx[i*frame*2+j*2+1] = f_value[1]
    #返回getfx值，2*frame*6 by 1
    return getfx
```

```python
def getJacobian(point, frame, para):
    # 用参数表示K，R矩阵
    focalx, focaly, px, py = symbols('focalx focaly px py')
    r11, r12, r13, r14, r21, r22, r23, r24, r31, r32, r33 = symbols('r11 r12 r13 r14 r21 r22 r23 r24 r31 r32 r33')
    Rt = Matrix([[r11, r12, r13, r14], [r21, r22, r23, r24], [r31, r32, r33, 1]])
    K = Matrix([[focalx, 0, px], [0, focaly, py], [0, 0, 1]])
    # KRt = K * Rt
    # 用参数表示ABCDEF六个点坐标
    a, b, c, th, al = symbols('a b c th al')
    ua, va, wa, ub, vb, wb, uc, vc, wc, ud, vd, wd, ue, ve, we, uf, vf, wf = symbols('ua va wa ub vb wb uc vc wc ud vd wd ue ve we uf vf wf')
    f = Symbol('f')
    if point == 0 :
        XA = Matrix([[-a * c * cos(th) * cos(al)], [c-a * c * sin(th)], [-a * c *  cos(th) * sin(al)], [1]])
        ua, va, wa = K[0,:] * (Rt * XA), K[1,:] * (Rt * XA), K[2,:] * (Rt * XA)
        f = Matrix([ua/wa, va/wa])
    elif point == 1 :
        XB = Matrix([[0], [c], [0], [1]])
        ub, vb, wb = K[0,:] * (Rt * XB), K[1,:] * (Rt * XB), K[2,:] * (Rt * XB)
        f = Matrix([ub/wb, vb/wb])
    elif point == 2 :
        XC = Matrix([[a * c * cos(th) * cos(al)], [c+a * c * sin(th)], [a * c * cos(th) * sin(al)], [1]])
        uc, vc, wc = K[0,:] * (Rt * XC), K[1,:] * (Rt * XC), K[2,:] * (Rt * XC)
        f = Matrix([uc/wc, vc/wc])
    elif point == 3 :
        XD = Matrix([[-b * c], [0], [0], [1]])
        ud, vd, wd = K[0,:] * (Rt * XD), K[1,:] * (Rt * XD), K[2,:] * (Rt * XD)
        f = Matrix([ud/wd, vd/wd])
    elif point == 4 :
        XE = Matrix([[0], [0], [0], [1]])
        ue, ve, we = K[0,:] * (Rt * XE), K[1,:] * (Rt * XE), K[2,:] * (Rt * XE)
        f = Matrix([ue/we, ve/we])
    elif point == 5:
        XF = Matrix([[b * c], [0], [0], [1]])
        uf, vf, wf = K[0,:] * (Rt * XF), K[1,:] * (Rt * XF), K[2,:] * (Rt * XF)
        f = Matrix([uf/wf, vf/wf])
    args = Matrix([r11, r12, r13, r14, r21, r22, r23, r24, r31, r32, r33, th, al, a, b, c])
    f_X1 = f[0,:].jacobian(args)
    f_X2 = f[1,:].jacobian(args)
    JA = Matrix([f_X1, f_X2])   # 2 by 16 matrix
    JA_value = JA.subs({focalx:1149.676, focaly:1147.592, px:508.849, py:508.065, r11:para[13*frame], r12:para[13*frame+1],
                                 r13:para[13*frame+2], r14:para[13*frame+3], r21:para[13*frame+4], r22:para[13*frame+5],
                                 r23:para[13*frame+6], r24:para[13*frame+7], r31:para[13*frame+8], r32:para[13*frame+9],
                                 r33:para[13*frame+10], th:para[13*frame+11], al:para[13*frame+12], a:para[-3], b:para[-2], c:para[-1]})
    #JA_value = JA_value.subs({f:1149.68})
    return JA_value
```

```python
def getJ(para, frame):
    getj = mat(np.zeros((6*frame*2, 13*frame+3)))
    for m in range(6):
        for n in range(frame):
            JA_value = getJacobian(m, n, para)
            #print(JA_value)
            getj[2*(m*frame+n):2*(m*frame+n+1), 13*n:13*n+13] = JA_value[:, 0:13]
            getj[2*(m*frame+n):2*(m*frame+n+1), -3:] = JA_value[:, -3:]
    return getj
```

```python
def getE(getfx, frame, k, xkp):
    E = mat(np.zeros((6*frame*2,1)))
    for i in range(6):
        for j in range(frame):
            if i==0 :
                E[(i*frame+j)*2] = getfx[i*frame*2+j*2] - xkp[i*3+k,0]
                E[(i*frame+j)*2+1] = getfx[i*frame*2+j*2+1] - xkp[i*3+k, 1]
            elif i==1 :
                E[(i*frame+j)*2] = getfx[i*frame*2+j*2] - xkp[i*3+k,0]
                E[(i*frame+j)*2+1] = getfx[i*frame*2+j*2+1] - xkp[i*3+k, 1]
            elif i==2 :
                E[(i*frame+j)*2] = getfx[i*frame*2+j*2] - xkp[i*3+k,0]
                E[(i*frame+j)*2+1] = getfx[i*frame*2+j*2+1] - xkp[i*3+k, 1]
            elif i==3 :
                E[(i*frame+j)*2] = getfx[i*frame*2+j*2] - xkp[i*3+k,0]
                E[(i*frame+j)*2+1] = getfx[i*frame*2+j*2+1] - xkp[i*3+k, 1]
            elif i==4 :
                E[(i*frame+j)*2] = getfx[i*frame*2+j*2] - xkp[i*3+k,0]
                E[(i*frame+j)*2+1] = getfx[i*frame*2+j*2+1] - xkp[i*3+k, 1]
            elif i==5 :
                E[(i*frame+j)*2] = getfx[i*frame*2+j*2] - xkp[i*3+k,0]
                E[(i*frame+j)*2+1] = getfx[i*frame*2+j*2+1] - xkp[i*3+k, 1]
    return E
```

```python
def LM_opti(frame, ith, x2d, x_para, u=1, v=2, step_max=500):
    J = mat(np.zeros((6*frame*2, 13*frame+3)))
    E = mat(np.zeros((6*frame*2,1)))   # E = f(X) - b ;
    E_temp = mat(np.zeros((6*frame*2,1)))   # E_temp compare with E in L-M
    x_k = mat(x_para.copy())   #parameter initialization
    step = 0   # iteration steps
    mse_last = 0   # mse value after iteration each time

    # L-M Algorithm obtain optimal parameters
    while(step < step_max):
        step += 1
        mse, mse_temp = 0, 0

        # generate Jacobian Matrix and calculate E
        getfx = mat(np.zeros((6*frame*2,1)))
        getfx = getFx(x_k, frame)
        E = getE(getfx, frame, ith, x2d)   #单帧模式，第几帧，GT

        for i in range(6*frame*2):
            mse += E[i]**2
        mse /= 6*frame*2

        # get new J
        J = mat(np.zeros((6*frame*2, 13*frame+3)))
        J = getJ(x_k, frame)

        # delta X = ...
        #print(J.T * J)
        dx = mat(np.zeros((13*frame+3,1)))
        LM = u * mat(np.eye(13*frame+3))
        dx = -(J.T * J + LM).I * J.T * E
        x_k_temp = x_k.copy()
        x_k_temp += dx

        #R的更新不能简单赋值
        #get R meet R.T*R=I
        #U * D * V.T = R --> R+ = U * V.T
        '''
        R_old = mat([[x_k_temp[0,0], x_k_temp[1,0], x_k_temp[2,0]],
                     [x_k_temp[4,0], x_k_temp[5,0], x_k_temp[6,0]],
                     [x_k_temp[8,0], x_k_temp[9,0], x_k_temp[10,0]]])
        U, sigma, VT = la.svd(R_old)
        R_new = U * VT
        x_k_temp[0,0], x_k_temp[1,0], x_k_temp[2,0] = R_new[0,0], R_new[0,1], R_new[0,2]
        x_k_temp[4,0], x_k_temp[5,0], x_k_temp[6,0] = R_new[1,0], R_new[1,1], R_new[1,2]
        x_k_temp[8,0], x_k_temp[9,0], x_k_temp[10,0] = R_new[2,0], R_new[2,1], R_new[2,2]
        '''
        ###########

        # calculate E_temp with x_k_temp
        # copy from E with x_k
        getfx_temp = mat(np.zeros((6*frame*2,1)))
        getfx_temp = getFx(x_k_temp, frame)
        E_temp = getE(getfx_temp, frame, ith, x2d)

        for i in range(6*frame*2):
            mse_temp += E_temp[i]**2
        mse_temp /= 6*frame*2

        # segma value to choose optimization model
        segma = (mse - mse_temp)/((dx.T * (u * dx - J.T * E))[0,0])

        # calculate new u
        if segma > 0:
            s = 1.0/3.0
            v = 2
            x_k = x_k_temp
            mse = mse_temp
            u = u * max(s, 1-pow(2*segma,3))
            u = u[0,0]
        else:
            u = u * v
            v = v * 2
            x_k = x_k_temp

        print("step = %d, abs(mse-mse_last) = %.8f" %(step, abs(mse-mse_last)))
        if abs(mse-mse_last)<0.000001:
            break
        mse_last = mse

    print("step = ", step)
    print("mse = ", mse_last)
    #print("parameter = ", x_k)

    return x_k
```

```python
frame = 3
m = loadmat("valid.mat")

# camera intrinsic matrix
K = m["annot"][0][0][4]
K_cam = K[0][0].tolist()

# key point 3D groundtruth
gt = m["annot"][0][0][3]
img_gt = np.zeros((frame, 3, 17))
img1_gt = gt[135]   # array 3 by 17
img2_gt = gt[136]   # array 3 by 17
img3_gt = gt[137]   # array 3 by 17
img_gt[0] = img1_gt
img_gt[1] = img2_gt
img_gt[2] = img3_gt
#print(img_gt)

# key point 2D groundtruth
kp = np.zeros((frame, 17, 2))
for m in range(frame):
    for i in range(17):
        u = K_cam[0] * mat([img_gt[m][0][i], img_gt[m][1][i], img_gt[m][2][i]]).T
        v = K_cam[1] * mat([img_gt[m][0][i], img_gt[m][1][i], img_gt[m][2][i]]).T
        w = K_cam[2] * mat([img_gt[m][0][i], img_gt[m][1][i], img_gt[m][2][i]]).T
        kp[m][i][0] = u/w
        kp[m][i][1] = v/w
#print(kp)

# show image
img1 = cv2.imread("S9_Posing_1.55011271_000676.jpg")
img2 = cv2.imread("S9_Posing_1.55011271_000681.jpg")
img3 = cv2.imread("S9_Posing_1.55011271_000686.jpg")
plt.figure(figsize = (15, 45)) # 图像窗口名称
plt.subplot(131)
plt.imshow(img1[:,:,[2,1,0]])
plt.subplot(132)
plt.imshow(img2[:,:,[2,1,0]])
plt.subplot(133)
plt.imshow(img3[:,:,[2,1,0]])
plt.axis('on') # 关掉坐标轴为 off
plt.title('image') # 图像题目
plt.show()

# visualize key points
"""
txt = ['1','2','3','4','5','6','7','8','9','10','11','12','13','14','15','16','17']
plt.figure()
plt.subplot(131)
img_kp = plt.scatter(kp[0][:,0], kp[0][:,1], s = 80, c = 'g', marker = 'X')
for i in range(17):
    plt.annotate(txt[i], xy = (kp[0][i,0], kp[0][i,1]), xytext = (kp[0][i,0]+0.1, kp[0][i,1]+0.1))
plt.subplot(132)
img_kp = plt.scatter(kp[1][:,0], kp[1][:,1], s = 80, c = 'g', marker = 'X')
for i in range(17):
    plt.annotate(txt[i], xy = (kp[1][i,0], kp[1][i,1]), xytext = (kp[1][i,0]+0.1, kp[1][i,1]+0.1))
plt.subplot(133)
img_kp = plt.scatter(kp[2][:,0], kp[2][:,1], s = 80, c = 'g', marker = 'X')
for i in range(17):
    plt.annotate(txt[i], xy = (kp[2][i,0], kp[2][i,1]), xytext = (kp[2][i,0]+0.1, kp[2][i,1]+0.1)) # 这里xy是需要标记的坐标，xytext是对应的标签坐标
plt.axis('on') # 关掉坐标轴为 off
plt.title('keypoint of image') # 图像题目
"""

# visualize torso joints A-F
plt.figure(figsize = (5,15))
plt.subplot(311)
for i in {0,1,4,8,11,14}:
    img_kp = plt.scatter(kp[0][i,0], kp[0][i,1], s = 80, c = 'g', marker = 'X')
plt.subplot(312)
for i in {0,1,4,8,11,14}:
    img_kp = plt.scatter(kp[1][i,0], kp[1][i,1], s = 80, c = 'g', marker = 'X')
plt.subplot(313)
for i in {0,1,4,8,11,14}:
    img_kp = plt.scatter(kp[2][i,0], kp[2][i,1], s = 80, c = 'g', marker = 'X')
plt.axis('on') # 关掉坐标轴为 off
plt.title('torso keypoint of image') # 图像题目
plt.show()

# save 2D keypoint coordinate
x2d = np.zeros((6 * frame,2))   #[AAABBBCCCDDDEEEFFF].T
for i in range(6):
    for j in range(frame):
        if i==0 :
            x2d[i*frame+j, :] = kp[j][14, :]
        elif i==1 :
            x2d[i*frame+j, :] = kp[j][8, :]
        elif i==2 :
            x2d[i*frame+j, :] = kp[j][11, :]
        elif i==3 :
            x2d[i*frame+j, :] = kp[j][4, :]
        elif i==4 :
            x2d[i*frame+j, :] = kp[j][0, :]
        elif i==5 :
            x2d[i*frame+j, :] = kp[j][1, :]
print(x2d)
```

```python
# parameter initialization for first frame
x_para = np.zeros((13+3,1))
x_para[0] = -1   # r11
x_para[5] = 1   # r22
x_para[10] = -1   # r33
x_para[3] = 0.0047   #tx(cal)
x_para[7] = -0.0997   #ty(cal)
x_para[11] = 0   # th
x_para[12] = 0   #al
x_para[-3] = 0.35   # a
x_para[-2] = 0.25   # b
x_para[-1] = -0.095  # c(cal)
getfx_ini = getFx(x_para, 1)
print(getfx_ini)
```

    [[476.02538041]
     [284.62882932]
     [514.25209744]
     [284.62882932]
     [552.47881446]
     [284.62882932]
     [486.94729956]
     [393.65003291]
     [514.25209744]
     [393.65003291]
     [541.55689531]
     [393.65003291]]

```python
# visualize first frame after initilization
plt.figure()
img_kp = plt.scatter(kp[0][14,0], kp[0][14,1], s = 80, c = 'g', marker = 'X')
img_kp = plt.scatter(getfx_ini[0,0], getfx_ini[1,0], s = 80, c = 'r', marker = 'X')
img_kp = plt.scatter(kp[0][8,0], kp[0][8,1], s = 80, c = 'g', marker = 'X')
img_kp = plt.scatter(getfx_ini[2,0], getfx_ini[3,0], s = 80, c = 'r', marker = 'X')
img_kp = plt.scatter(kp[0][11,0], kp[0][11,1], s = 80, c = 'g', marker = 'X')
img_kp = plt.scatter(getfx_ini[4,0], getfx_ini[5,0], s = 80, c = 'r', marker = 'X')
img_kp = plt.scatter(kp[0][4,0], kp[0][4,1], s = 80, c = 'g', marker = 'X')
img_kp = plt.scatter(getfx_ini[6,0], getfx_ini[7,0], s = 80, c = 'r', marker = 'X')
img_kp = plt.scatter(kp[0][0,0], kp[0][0,1], s = 80, c = 'g', marker = 'X')
img_kp = plt.scatter(getfx_ini[8,0], getfx_ini[9,0], s = 80, c = 'r', marker = 'X')
img_kp = plt.scatter(kp[0][1,0], kp[0][1,1], s = 80, c = 'g', marker = 'X')
img_kp = plt.scatter(getfx_ini[10,0], getfx_ini[11,0], s = 80, c = 'r', marker = 'X')
plt.axis('on') # 关掉坐标轴为 off
plt.title('image_kp_ABCDEF_original_opt') # 图像题目
plt.show()
```

```python
# optimization for each frame
X_para = np.zeros((13*frame+3,1))
X_para[-3] = 0.35   # a
X_para[-2] = 0.25   # b
X_para[-1] = -0.095  # c(cal)
para_data = np.zeros((frame,16))
x_k_para = x_para.copy()
for i in range(frame):
    x_k = LM_opti(1, i, x2d, x_k_para)   #单帧优化，第几帧，GT，初始化参数

    # R matrix adjustment

    R_old = mat([[x_k[0,0], x_k[1,0], x_k[2,0]],
             [x_k[4,0], x_k[5,0], x_k[6,0]],
             [x_k[8,0], x_k[9,0], x_k[10,0]]])
    U, sigma, VT = la.svd(R_old)
    R_new = U * VT
    x_k[0,0], x_k[1,0], x_k[2,0] = R_new[0,0], R_new[0,1], R_new[0,2]
    x_k[4,0], x_k[5,0], x_k[6,0] = R_new[1,0], R_new[1,1], R_new[1,2]
    x_k[8,0], x_k[9,0], x_k[10,0] = R_new[2,0], R_new[2,1], R_new[2,2]

    print()
    print("For the No.", i+1, " frame : x_k = ", x_k)
    para_data[i] = mat(x_k).T
    # for next frame initilization
    x_k_para = x_k
    X_para[13*i:13*i+13] = x_k[0:13,0]
print(para_data)
print(X_para)
```

    step = 1, abs(mse-mse_last) = 9.25453983
    step = 2, abs(mse-mse_last) = 0.05628926
    step = 3, abs(mse-mse_last) = 0.00023892
    step = 4, abs(mse-mse_last) = 0.00003550
    step = 5, abs(mse-mse_last) = 0.00000600
    step = 6, abs(mse-mse_last) = 0.00000140
    step = 7, abs(mse-mse_last) = 0.00000059
    step =  7
    mse =  [[9.19796875]]

    For the No. 1  frame : x_k =  [[-9.97798014e-01]
     [-6.58262456e-02]
     [ 8.12577467e-03]
     [ 6.47967046e-03]
     [-6.58158758e-02]
     [ 9.97830600e-01]
     [ 1.53732561e-03]
     [-9.96904449e-02]
     [-8.20934299e-03]
     [ 9.99135460e-04]
     [-9.99965804e-01]
     [ 7.40954851e-02]
     [ 5.27034975e-02]
     [ 3.63567383e-01]
     [ 2.55433593e-01]
     [-8.78429471e-02]]
    step = 1, abs(mse-mse_last) = 10.19603605
    step = 2, abs(mse-mse_last) = 0.26428436
    step = 3, abs(mse-mse_last) = 0.01386093
    step = 4, abs(mse-mse_last) = 0.00019781
    step = 5, abs(mse-mse_last) = 0.00013194
    step = 6, abs(mse-mse_last) = 0.00008733
    step = 7, abs(mse-mse_last) = 0.00006271
    step = 8, abs(mse-mse_last) = 0.00004550
    step = 9, abs(mse-mse_last) = 0.00003305
    step = 10, abs(mse-mse_last) = 0.00002395
    step = 11, abs(mse-mse_last) = 0.00001730
    step = 12, abs(mse-mse_last) = 0.00001246
    step = 13, abs(mse-mse_last) = 0.00000895
    step = 14, abs(mse-mse_last) = 0.00000641
    step = 15, abs(mse-mse_last) = 0.00000457
    step = 16, abs(mse-mse_last) = 0.00000326
    step = 17, abs(mse-mse_last) = 0.00000231
    step = 18, abs(mse-mse_last) = 0.00000164
    step = 19, abs(mse-mse_last) = 0.00000116
    step = 20, abs(mse-mse_last) = 0.00000082
    step =  20
    mse =  [[9.91725041]]

    For the No. 2  frame : x_k =  [[-9.95464362e-01]
     [-5.19786722e-02]
     [ 7.96801251e-02]
     [ 2.82750046e-04]
     [-5.33888922e-02]
     [ 9.98450838e-01]
     [-1.56700384e-02]
     [-9.97721556e-02]
     [-7.87421799e-02]
     [-1.98529984e-02]
     [-9.96697310e-01]
     [ 1.62911934e-01]
     [ 5.18197754e-01]
     [ 4.42342393e-01]
     [ 2.63307728e-01]
     [-9.00014453e-02]]
    step = 1, abs(mse-mse_last) = 17.24330917
    step = 2, abs(mse-mse_last) = 4.81374060
    step = 3, abs(mse-mse_last) = 0.36139776
    step = 4, abs(mse-mse_last) = 1.72911107
    step = 5, abs(mse-mse_last) = 0.00319213
    step = 6, abs(mse-mse_last) = 0.00002203
    step = 7, abs(mse-mse_last) = 0.00000689
    step = 8, abs(mse-mse_last) = 0.00000398
    step = 9, abs(mse-mse_last) = 0.00000258
    step = 10, abs(mse-mse_last) = 0.00000168
    step = 11, abs(mse-mse_last) = 0.00000109
    step = 12, abs(mse-mse_last) = 0.00000071
    step =  12
    mse =  [[10.33582936]]

    For the No. 3  frame : x_k =  [[-0.97778353]
     [-0.05453156]
     [ 0.20239981]
     [-0.00734245]
     [-0.06022446]
     [ 0.99794082]
     [-0.02207125]
     [-0.10002235]
     [-0.20077945]
     [-0.03377033]
     [-0.97905423]
     [ 0.24425288]
     [ 0.87964921]
     [ 0.69640617]
     [ 0.2562235 ]
     [-0.09460461]]
    [[-9.97798014e-01 -6.58262456e-02  8.12577467e-03  6.47967046e-03
      -6.58158758e-02  9.97830600e-01  1.53732561e-03 -9.96904449e-02
      -8.20934299e-03  9.99135460e-04 -9.99965804e-01  7.40954851e-02
       5.27034975e-02  3.63567383e-01  2.55433593e-01 -8.78429471e-02]
     [-9.95464362e-01 -5.19786722e-02  7.96801251e-02  2.82750046e-04
      -5.33888922e-02  9.98450838e-01 -1.56700384e-02 -9.97721556e-02
      -7.87421799e-02 -1.98529984e-02 -9.96697310e-01  1.62911934e-01
       5.18197754e-01  4.42342393e-01  2.63307728e-01 -9.00014453e-02]
     [-9.77783527e-01 -5.45315558e-02  2.02399810e-01 -7.34244781e-03
      -6.02244577e-02  9.97940817e-01 -2.20712528e-02 -1.00022346e-01
      -2.00779452e-01 -3.37703262e-02 -9.79054226e-01  2.44252876e-01
       8.79649208e-01  6.96406175e-01  2.56223500e-01 -9.46046054e-02]]
    [[-9.97798014e-01]
     [-6.58262456e-02]
     [ 8.12577467e-03]
     [ 6.47967046e-03]
     [-6.58158758e-02]
     [ 9.97830600e-01]
     [ 1.53732561e-03]
     [-9.96904449e-02]
     [-8.20934299e-03]
     [ 9.99135460e-04]
     [-9.99965804e-01]
     [ 7.40954851e-02]
     [ 5.27034975e-02]
     [-9.95464362e-01]
     [-5.19786722e-02]
     [ 7.96801251e-02]
     [ 2.82750046e-04]
     [-5.33888922e-02]
     [ 9.98450838e-01]
     [-1.56700384e-02]
     [-9.97721556e-02]
     [-7.87421799e-02]
     [-1.98529984e-02]
     [-9.96697310e-01]
     [ 1.62911934e-01]
     [ 5.18197754e-01]
     [-9.77783527e-01]
     [-5.45315558e-02]
     [ 2.02399810e-01]
     [-7.34244781e-03]
     [-6.02244577e-02]
     [ 9.97940817e-01]
     [-2.20712528e-02]
     [-1.00022346e-01]
     [-2.00779452e-01]
     [-3.37703262e-02]
     [-9.79054226e-01]
     [ 2.44252876e-01]
     [ 8.79649208e-01]
     [ 3.50000000e-01]
     [ 2.50000000e-01]
     [-9.50000000e-02]]

```python
def visual2D(kp, getfx):
    plt.figure()
    img_kp = plt.scatter(kp[14,0], kp[14,1], s = 80, c = 'g', marker = 'X')
    img_kp = plt.scatter(getfx[0,0], getfx[1,0], s = 80, c = 'r', marker = 'X')
    img_kp = plt.scatter(kp[8,0], kp[8,1], s = 80, c = 'g', marker = 'X')
    img_kp = plt.scatter(getfx[2,0], getfx[3,0], s = 80, c = 'r', marker = 'X')
    img_kp = plt.scatter(kp[11,0], kp[11,1], s = 80, c = 'g', marker = 'X')
    img_kp = plt.scatter(getfx[4,0], getfx[5,0], s = 80, c = 'r', marker = 'X')
    img_kp = plt.scatter(kp[4,0], kp[4,1], s = 80, c = 'g', marker = 'X')
    img_kp = plt.scatter(getfx[6,0], getfx[7,0], s = 80, c = 'r', marker = 'X')
    img_kp = plt.scatter(kp[0,0], kp[0,1], s = 80, c = 'g', marker = 'X')
    img_kp = plt.scatter(getfx[8,0], getfx[9,0], s = 80, c = 'r', marker = 'X')
    img_kp = plt.scatter(kp[1,0], kp[1,1], s = 80, c = 'g', marker = 'X')
    img_kp = plt.scatter(getfx[10,0], getfx[11,0], s = 80, c = 'r', marker = 'X')
    plt.axis('on') # 关掉坐标轴为 off
    plt.title('image_kp_ABCDEF_original_opt') # 图像题目
    plt.show()
```

```python
visual = np.zeros((13+3,1))
for i in range(frame):
    visual = mat(para_data[i]).T
    getfx_visual = getFx(visual, 1)
    visual2D(kp[i], getfx_visual)
```

```python
def visual3D(img_gt, para):
    K = Matrix([[1149.67569986785, 0.0, 508.848621645943],
                 [0.0, 1147.59161666764, 508.064917088557],
                 [0.0, 0.0, 1.0]])
    r11, r12, r13, r14, r21, r22, r23, r24, r31, r32, r33 = symbols('r11 r12 r13 r14 r21 r22 r23 r24 r31 r32 r33')
    Rt = Matrix([[r11, r12, r13, r14], [r21, r22, r23, r24], [r31, r32, r33, 1]])
    a, b, c, th, al = symbols('a b c th al')
    ua, va, wa, ub, vb, wb, uc, vc, wc, ud, vd, wd, ue, ve, we, uf, vf, wf = symbols('ua va wa ub vb wb uc vc wc ud vd wd ue ve we uf vf wf')
    XA = Matrix([[-a * c * cos(th) * cos(al)], [c-a * c * sin(th)], [-a * c *  cos(th) * sin(al)], [1]])
    XB = Matrix([[0], [c], [0], [1]])
    XC = Matrix([[a * c * cos(th) * cos(al)], [c+a * c * sin(th)], [a * c * cos(th) * sin(al)], [1]])
    XD = Matrix([[-b * c], [0], [0], [1]])
    XE = Matrix([[0], [0], [0], [1]])
    XF = Matrix([[b * c], [0], [0], [1]])
    A3D = Rt * XA
    B3D = Rt * XB
    C3D = Rt * XC
    D3D = Rt * XD
    E3D = Rt * XE
    F3D = Rt * XF
    j = 0
    #A3D_value = A3D.subs({r11:para[13*j], r12:para[13*j+1], r13:para[13*j+2], r14:para[13*j+3],
    #            r21:para[13*j+4], r22:para[13*j+5], r23:para[13*j+6], r24:para[13*j+7],
    #            r31:para[13*j+8], r32:para[13*j+9], r33:para[13*j+10], th:para[13*j+11],
    #            al:para[13*j+12], a:para[-3], b:para[-2], c:para[-1]})
    Y = mat([img_gt[0][0], img_gt[1][0], img_gt[2][0]]).T
    s = Y[2,0]
    print(s)
    #A3D_value = getValue(A3D, para) * s
    A3D_value = A3D.subs({r11:para[0], r12:para[1], r13:para[2], r14:para[3],
        r21:para[4], r22:para[5], r23:para[6], r24:para[7],
        r31:para[8], r32:para[9], r33:para[10], th:para[11],
        al:para[12], a:para[-3], b:para[-2], c:para[-1]}) * s
    B3D_value = B3D.subs({r11:para[0], r12:para[1], r13:para[2], r14:para[3],
        r21:para[4], r22:para[5], r23:para[6], r24:para[7],
        r31:para[8], r32:para[9], r33:para[10], th:para[11],
        al:para[12], a:para[-3], b:para[-2], c:para[-1]}) * s
    C3D_value = C3D.subs({r11:para[0], r12:para[1], r13:para[2], r14:para[3],
        r21:para[4], r22:para[5], r23:para[6], r24:para[7],
        r31:para[8], r32:para[9], r33:para[10], th:para[11],
        al:para[12], a:para[-3], b:para[-2], c:para[-1]}) * s
    D3D_value = D3D.subs({r11:para[0], r12:para[1], r13:para[2], r14:para[3],
        r21:para[4], r22:para[5], r23:para[6], r24:para[7],
        r31:para[8], r32:para[9], r33:para[10], th:para[11],
        al:para[12], a:para[-3], b:para[-2], c:para[-1]}) * s
    E3D_value = E3D.subs({r11:para[0], r12:para[1], r13:para[2], r14:para[3],
        r21:para[4], r22:para[5], r23:para[6], r24:para[7],
        r31:para[8], r32:para[9], r33:para[10], th:para[11],
        al:para[12], a:para[-3], b:para[-2], c:para[-1]}) * s
    F3D_value = F3D.subs({r11:para[0], r12:para[1], r13:para[2], r14:para[3],
        r21:para[4], r22:para[5], r23:para[6], r24:para[7],
        r31:para[8], r32:para[9], r33:para[10], th:para[11],
        al:para[12], a:para[-3], b:para[-2], c:para[-1]}) * s
    print(A3D_value,B3D_value,C3D_value,D3D_value,E3D_value,F3D_value)

    X1 = mat([img_gt[0][14], img_gt[1][14], img_gt[2][14]]).T
    X2 = mat([img_gt[0][8], img_gt[1][8], img_gt[2][8]]).T
    X3 = mat([img_gt[0][11], img_gt[1][11], img_gt[2][11]]).T
    X4 = mat([img_gt[0][4], img_gt[1][4], img_gt[2][4]]).T
    X5 = mat([img_gt[0][0], img_gt[1][0], img_gt[2][0]]).T
    X6 = mat([img_gt[0][1], img_gt[1][1], img_gt[2][1]]).T
    print(X1, X2, X3, X4, X5, X6)

    # 绘制散点图
    fig = plt.figure()
    ax = Axes3D(fig)
    ax.scatter(X1[0], X1[1], X1[2], s = 80, c = 'g', marker = 'X')
    ax.scatter(X2[0], X2[1], X2[2], s = 80, c = 'g', marker = 'X')
    ax.scatter(X3[0], X3[1], X3[2], s = 80, c = 'g', marker = 'X')
    ax.scatter(X4[0], X4[1], X4[2], s = 80, c = 'g', marker = 'X')
    ax.scatter(X5[0], X5[1], X5[2], s = 80, c = 'g', marker = 'X')
    ax.scatter(X6[0], X6[1], X6[2], s = 80, c = 'g', marker = 'X')
    ax.scatter(A3D_value[0], A3D_value[1], A3D_value[2], s = 80, c = 'r', marker = 'X')
    ax.scatter(B3D_value[0], B3D_value[1], B3D_value[2], s = 80, c = 'r', marker = 'X')
    ax.scatter(C3D_value[0], C3D_value[1], C3D_value[2], s = 80, c = 'r', marker = 'X')
    ax.scatter(D3D_value[0], D3D_value[1], D3D_value[2], s = 80, c = 'r', marker = 'X')
    ax.scatter(E3D_value[0], E3D_value[1], E3D_value[2], s = 80, c = 'r', marker = 'X')
    ax.scatter(F3D_value[0], F3D_value[1], F3D_value[2], s = 80, c = 'r', marker = 'X')

    ax.set_zlabel('Z', fontdict={'size': 15, 'color': 'red'})
    ax.set_ylabel('Y', fontdict={'size': 15, 'color': 'red'})
    ax.set_xlabel('X', fontdict={'size': 15, 'color': 'red'})
    plt.show()

    # 计算3D error
    d1 = A3D_value-X1
    e1 = sum(map(lambda x:x**2, d1))
    d2 = B3D_value-X2
    e2 = sum(map(lambda x:x**2, d2))
    d3 = C3D_value-X3
    e3 = sum(map(lambda x:x**2, d3))
    d4 = D3D_value-X4
    e4 = sum(map(lambda x:x**2, d4))
    d5 = E3D_value-X5
    e5 = sum(map(lambda x:x**2, d5))
    d6 = F3D_value-X6
    e6 = sum(map(lambda x:x**2, d6))
    error = (e1**(0.5)+e2**(0.5)+e3**(0.5)+e4**(0.5)+e5**(0.5)+e6**(0.5))/6
    print(error)
    print(d1,d2,d3,d4,d5,d6)
```

```python
visual3d = np.zeros((13+3,1))
for i in range(frame):
    visual3d = mat(para_data[i]).T
    visual3D(img_gt[i], visual3d)
```

```python

```


```python
import numpy as np
from numpy import linalg as LA
import math
import random

#任意次元のベクトルを生成する関数
#nは生成したい次元
def qubit(n):

    e=np.zeros(n,dtype='float64')

    e[0]=random.random()
    unf=1-e[0]

    for i in range(1,n-1):
        e[i]=random.uniform(0,unf)
        unf=unf-e[i]
    e[n-1]=unf

    return e

#ランダムなベクトルを規格化し、状態ベクトルになる．
def normal(e):
    n=e.size      #配列の次元を確認する関数

    E=np.zeros(n,dtype='float64')

    for i in range(0,n):
        E[i]=np.sqrt(e[i])
    return E

#ベクトルの間のテンソル積
def tensor(a,b):
    n1=a.size
    n2=b.size

    P=np.zeros(n1*n2,dtype='float64')
    for i in range(0,n1):
        for j in range(0,n2):
            P[i*n2+j]=a[i]*b[j]
    return P

#xはn1次元のベクトル、yはn2次元のベクトルとおく、zの次元はn1*n2である
x=normal(qubit(2))
y=normal(qubit(2))
z=tensor(x,y)

print('状態１は')
print(x)
print('状態2は')
print(y)
print('合成系の状態は')
print(z)

#ノルムの確認

print(LA.norm(x,2))

print(LA.norm(y,2))

print(LA.norm(z,2))
```

    状態１は
    [0.37032596 0.92890187]
    状態2は
    [0.39147442 0.92018899]
    合成系の状態は
    [0.14497314 0.34076987 0.36364132 0.85476527]
    1.0
    1.0
    1.0

# 純粋状態の密度演算子

\begin{equation}
\rho_{1}=|n1\rangle\langle n1|,\rho_{2}=|n2\rangle\langle n2|,
\end{equation}

合成系の密度演算子は

\begin{equation}
\rho=|n1\otimes n2\rangle\langle n1\otimes n2|
\end{equation}

```python
rho=np.dot(z,z.T)
print(rho)
print(z.T)
```

    1.0000000000000002
    [0.05849676 0.96047723 0.01654388 0.27163932]

```python
print(z.T)
```

    [0.05849676 0.96047723 0.01654388 0.27163932]

```python

#密度演算子を生成する関数
def density(a):
    dimension=a.size

    b=np.zeros([dimension,dimension],dtype='float64')
    for i in range(0,dimension):
        for j in range(0,dimension):
            b[i][j]=a[i]*a[j]

    return b

#トレースを計算する関数
def trace(a):
    sum=0;
    n=a[0].size
    for i in range(0,n):
        sum+=a[i][i]

    return sum

rho=density(z)
print(rho)
print(trace(rho))
```

    [[3.42187097e-03 5.61848063e-02 9.67763412e-04 1.58900205e-02]
     [5.61848063e-02 9.22516506e-01 1.58900205e-02 2.60903385e-01]
     [9.67763412e-04 1.58900205e-02 2.73699981e-04 4.49396852e-03]
     [1.58900205e-02 2.60903385e-01 4.49396852e-03 7.37879226e-02]]
    1.0000000000000002

```python
m1=np.array([1,0])
m2=np.array([0,1])
print(density(m1))
```

    [[1. 0.]
     [0. 0.]]

# 計算基底における同時測定

計算基底について、

\begin{equation}
|00\rangle,  |01\rangle,  |10\rangle,  |11\rangle
\end{equation}

```python
print(tensor(m1,m2))
```

    [0. 1. 0. 0.]

```python
#00 ketの確率を計算する

m=tensor(m1,m2)

def prob(base,state):
    z=np.dot(base,state)
    return math.pow(z,2)

#zは合成系の状態ベクトルである
prob1=prob(tensor(m1,m1),z)
prob2=prob(tensor(m1,m2),z)
prob3=prob(tensor(m2,m1),z)
prob4=prob(tensor(m2,m2),z)

print(prob1,prob2,prob3,prob4)
print(prob1+prob2+prob3+prob4)

#print(prob(m,z))
```

    0.0034218709747689006 0.9225165064594494 0.00027369998131451556 0.07378792258446734
    1.0000000000000002

# 計算基底における部分測定

```python
print(np.dot(m1,x))
```

    0.9622569186211228

```python
parcial1=np.dot(m1,x)*y

print(parcial1)
print(math.pow(LA.norm(parcial1,2),2))
```

    [0.05849676 0.96047723]
    0.9259383774342183

# 密度演算子による部分測定

```python
rho_m=trace(density(y))
print(trace(rho_m*density(x)*density(m1)))
```

    0.7093372706653313

# 密度演算子による同時測定

状態$\rho$の下での同時測定は

\begin{equation}
    Pr(E=m,F=n | \rho)=Tr(\rho E_{m}\otimes F_{n})
\end{equation}

```python
#行列式同士のテンソルを作る関数
#P135

#n1=a.size
 #   n2=math.pow(n2,2)

 #   e=np.zeros([n2,n2],dtype='float64')

  #  for i in range(0,n1):
   #     for j in range(0,n1):

```

```python
print(x)
print(LA.norm(x,2))
rho1=density(x)
print(rho1)

```

    [0.96225692 0.27214265]
    1.0
    [[0.92593838 0.26187115]
     [0.26187115 0.07406162]]

```python
E1=0.4*density(m1)
print(E1)
```

    [[0.4 0. ]
     [0.  0. ]]

```python
E2=0.4*density(np.array([1/np.sqrt(2),1/np.sqrt(2)]))
print(E2)
```

    [[0.2 0.2]
     [0.2 0.2]]

```python
E3=np.array([[1,0],[0,1]])-E2-E1
print(E3)
```

    [[ 0.4 -0.2]
     [-0.2  0.8]]

```python
probi1=trace(np.dot(rho1,E1))
print(probi1)
probi2=trace(np.dot(rho1,E2))
probi3=trace(np.dot(rho1,E3))
print(probi2,probi3)
print(probi1+probi2+probi3)
```

    0.35556736235065745
    0.32569326064375625 0.31873937700558636
    1.0