import espressomd
import espressomd.magnetostatics

espressomd.assert_features(['DIPOLES', 'DP3M', 'LENNARD_JONES'])

%matplotlib inline
import matplotlib.pyplot as plt
plt.rcParams.update({'font.size': 18})
import numpy as np
import tqdm

"""
Explanation: Ferrofluid - Part 3
Table of Contents

Susceptibility with fluctuation formulas
Derivation of the fluctuation formula
Simulation


Magnetization curve of a 3D system

Remark: The equilibration and sampling times used in this tutorial would be not sufficient for scientific  purposes, but they are long enough to get at least a qualitative insight of the behaviour of ferrofluids. They have been shortened so we achieve reasonable computation times for the purpose of a tutorial.
Susceptibility with fluctuation formulas
In this part we want to calculate estimators for the initial susceptibility, i.e. the susceptibility at zero external magnetic field. One could carry out several simulations with different external magnetic field strengths and get the initial susceptibility by fitting a line to the results. We want to go a more elegant way by using fluctuation formulas known from statistical mechanics.
In three dimensions the initial susceptibility $\chi_{init}$ can be calculated with zero field simulations through
\begin{equation}
    \chi_\mathrm{init} = \frac{V \cdot \mu_0}{3 \cdot k_\mathrm{B} T} \left( \langle \boldsymbol{M}^2 \rangle - \langle \boldsymbol{M} \rangle^2 \right) = \frac{\mu_0}{3 \cdot k_\mathrm{B} T \cdot V} \left( \langle \boldsymbol{\mu}^2 \rangle - \langle \boldsymbol{\mu} \rangle^2 \right)
\end{equation}
where $\boldsymbol{M}$ is the magnetization vector and $\boldsymbol{\mu}$ is the total magnetic dipole moment of the system. In direction $i$ it reads
\begin{equation}
    M_i = \frac{1}{V} \Bigg\langle \sum_{j=1}^N \tilde{\mu}_j^i \Bigg\rangle
\end{equation}
where $\tilde{\mu}_j^i$ is the $j$ th dipole moment in direction $i$.
Derivation of the fluctuation formula
We want to derive the fluctuation formula. We start with the definition of the magnetic susceptibility. In general this reads
\begin{equation}
    \chi \equiv \frac{\partial}{\partial H} \langle M_{\boldsymbol{H}} \rangle
\end{equation}
with $\langle M_{\boldsymbol{H}} \rangle$ the ensemble averaged magnetization in direction of a homogeneous external magnetic field $\boldsymbol{H}$.
In thermal equilibrium the ensemble average of the magnetization reads
\begin{equation}
    \langle M_{\boldsymbol{H}} \rangle = \frac{1}{V Z_\mathrm{c}} \left \lbrack \sum_{\alpha} \mu_{\boldsymbol{H},\alpha} e^{ -\beta E_{\alpha}(H=0) + \beta\mu_0\mu_{\boldsymbol{H},\alpha}H }\right \rbrack
\end{equation}
with $Z_\mathrm{c}$ the canonical partition function, $E_{\alpha}(H=0)$ the energy without an external magnetic field $\boldsymbol{H}$, $\beta$ the inverse thermal energy $\frac{1}{k_\mathrm{B}T}$, $\mu_{\boldsymbol{H},\alpha}$ the total magnetic dipole moment of the system in direction of the external magnetic field $\boldsymbol{H}$ in microstate $\alpha$ and $V$ the system volume.
Now we insert the magnetization $\langle M_{\boldsymbol{H}} \rangle$ in the definition of the magnetic susceptibility $\chi$ and let the derivative operate on the ensemble average. We get the fluctuation formula
\begin{equation}
    \chi = \frac{\beta\mu_0}{V} \left \lbrack \frac{1}{Z_\mathrm{c}}\sum_{\alpha} \mu_{\alpha}^2~ e^{ -\beta E_{\alpha}(H=0) + \beta\mu_0\mu_{\boldsymbol{H},\alpha}H } - \frac{1}{Z_\mathrm{c}}\sum_{\alpha} \mu_{\alpha}~ e^{ -\beta E_{\alpha}(H=0) + \beta\mu_0\mu_{\boldsymbol{H},\alpha}H }~~ \frac{1}{Z_\mathrm{c}}\sum_{\alpha'}\mu_{\alpha'}~ e^{ -\beta E_{\alpha'}(H=0) + \beta\mu_0\mu_{\boldsymbol{H},\alpha}H }\right \rbrack = \frac{\beta\mu_0}{V} \left \lbrack \langle \mu_{\boldsymbol{H}}^2 \rangle - \langle \mu_{\boldsymbol{H}} \rangle^2 \right \rbrack = \frac{\beta\mu_0}{V} \left(\Delta \mu_{\boldsymbol{H}}\right)^2
\end{equation}
At zero external magnetic field ($H = 0$) there is no distinct direction for the system, so we can take the fluctuations $\Delta \mu$ in all directions and divide it by the dimension. Thus we can use more data points of our simulation for the average and get a more precise estimator for the susceptibility. Thus finally the fluctuation formula for the initial susceptibility in three dimensions reads
\begin{equation}
    \chi_\mathrm{init} = \frac{\beta\mu_0}{3V} \left \lbrack \langle \boldsymbol{\mu}^2 \rangle - \langle \boldsymbol{\mu} \rangle^2 \right \rbrack = \frac{V\beta\mu_0}{3} \left \lbrack \langle \boldsymbol{M}^2 \rangle - \langle \boldsymbol{M} \rangle^2 \right \rbrack
\end{equation}
where $\boldsymbol{\mu}$ and $\boldsymbol{M}$ are defined above.
Simulation
In this part we want to consider a three dimensional ferrofluid system and compare our result for the initial susceptibility $\chi_\mathrm{init}$ with them of Ref. <a href='#[1]'>[1]</a>.
First we import all necessary packages and check for the required ESPResSo features
End of explanation
"""


lj_sigma = 1
lj_epsilon = 1
lj_cut = 2**(1. / 6.) * lj_sigma

# magnetic field constant
mu_0 = 1.

# Particles
N = 1000

# Volume fraction
# phi = rho * 4. / 3. * np.pi * ( lj_sigma / 2 )**3.
phi = 0.0262

# Dipolar interaction parameter lambda = mu_0 m^2 /(4 pi sigma^3 kT)
dip_lambda = 3.

# Temperature
kT = 1.0

# Friction coefficient
gamma = 1.0

# Time step
dt = 0.02

# box size 3d
box_size = (N * np.pi * 4. / 3. * (lj_sigma / 2.)**3. / phi)**(1. / 3.)

"""
Explanation: Now we set up all necessary simulation parameters
End of explanation
"""


system = espressomd.System(box_l=(box_size, box_size, box_size))
system.time_step = dt

# Lennard-Jones interaction
system.non_bonded_inter[0, 0].lennard_jones.set_params(epsilon=lj_epsilon, sigma=lj_sigma, cutoff=lj_cut, shift="auto")

# Random dipole moments
np.random.seed(seed=1)
dip_phi = 2 * np.pi * np.random.random((N, 1))
dip_cos_theta = 2 * np.random.random((N, 1)) - 1
dip_sin_theta = np.sin(np.arccos(dip_cos_theta))
dip = np.hstack((
    dip_sin_theta * np.sin(dip_phi),
    dip_sin_theta * np.cos(dip_phi),
    dip_cos_theta))

# Random positions in system volume
pos = box_size * np.random.random((N, 3))

# Add particles
particles = system.part.add(pos=pos, rotation=N * [(True, True, True)], dip=dip)

# Remove overlap between particles by means of the steepest descent method
system.integrator.set_steepest_descent(
    f_max=0, gamma=0.1, max_displacement=0.05)

while system.analysis.energy()["total"] > 5 * kT * N:
    system.integrator.run(20)

# Switch to velocity Verlet integrator
system.integrator.set_vv()
system.thermostat.set_langevin(kT=kT, gamma=gamma, seed=1)

# tune verlet list skin
system.cell_system.skin = 0.8

# Setup dipolar P3M
accuracy = 5E-4
system.actors.add(espressomd.magnetostatics.DipolarP3M(accuracy=accuracy, prefactor=dip_lambda * lj_sigma**3 * kT))

"""
Explanation: Next we set up the system. As in part I, the orientation of the dipole moments is set directly on the particles, whereas the magnitude of the moments is taken into account when determining the prefactor of the dipolar P3M (for more details see part I).
Hint:
It should be noted that we seed both the Langevin thermostat and the random number generator of numpy. Latter means that the initial configuration of our system is the same every time this script is executed. As the time evolution of the system depends not solely on the Langevin thermostat but also on the numeric accuracy and DP3M (the tuned parameters are slightly different every time) it is only partly predefined. You can change the seeds to simulate with a different initial configuration and a guaranteed different time evolution.
End of explanation
"""


equil_rounds = 10
equil_steps = 1000
for i in tqdm.trange(equil_rounds):
    system.integrator.run(equil_steps)

"""
Explanation: Now we equilibrate for a while
End of explanation
"""


import espressomd.observables
dipm_tot_calc = espressomd.observables.MagneticDipoleMoment(ids=particles.id)

"""
Explanation: As we need the magnetization of our system, we use <tt>espressomd.observables.MagneticDipoleMoment</tt> to calculate the total dipole moment of the system which is the magnetization times the volume of the system.
End of explanation
"""


# Sampling
loops = 2000

"""
Explanation: Now we set the desired number of loops for the sampling
End of explanation
"""


# initialize array for hold the sampled dipole moments
dipms = np.full((loops, 3), np.nan)

# sample dipole moment
for i in tqdm.trange(loops):
    system.integrator.run(10)
    dipms[i, :] = dipm_tot_calc.calculate()

# calculate average first and second moment of total dipole moment
dipm_tot = np.mean(dipms, axis=0)
dipm_tot_2 = np.mean(dipms**2, axis=0)

"""
Explanation: and sample the first and second moment of the magnetization or total dipole moment, by averaging over all total dipole moments occurring during the simulation
End of explanation
"""


# dipole moment
dipm = np.sqrt(dip_lambda * 4 * np.pi * lj_sigma**3. * kT / mu_0)
print(f'dipm = {dipm:.4f}')

"""
Explanation: For the estimator of the initial susceptibility $\chi_\mathrm{init}$ we need the magnitude of one single dipole moment
End of explanation
"""


# susceptibility in 3d system
chi = mu_0 / (system.volume() * 3. * kT) * (np.sum(dipm_tot_2 * dipm**2.) - np.sum(np.square(dipm_tot * dipm)))

"""
Explanation: Now we can calculate $\chi_\mathrm{init}$ from our simulation data
End of explanation
"""


print(f'chi = {chi:.4f}')

"""
Explanation: and print the result
End of explanation
"""


chi_L = 8. * dip_lambda * phi
print(f'chi_L = {chi_L:.4f}')

"""
Explanation: Compared with the value $\chi = 0.822 \pm 0.017$ of Ref. <a href='#[1]'>[1]</a> (see table 1) it should be very similar.
Now we want to compare the result with the theoretical expectations.
At first with the simple Langevin susceptibility
End of explanation
"""


chi_I = chi_L * (1 + chi_L / 3. + chi_L**2. / 144.)
print(f'chi_I = {chi_I:.4f}')

"""
Explanation: and at second with the more advanced one (see Ref. <a href='#[1]'>[1]</a> eq. (6)) which has a cubic accuracy in $\chi_\mathrm{L}$ and reads
\begin{equation}
    \chi = \chi_\mathrm{L} \left( 1 + \frac{\chi_\mathrm{L}}{3} + \frac{\chi_\mathrm{L}^2}{144} \right)
\end{equation}
End of explanation
"""


alphas = np.array([0, 0.25, 0.5, 1, 2, 3, 4, 8])

"""
Explanation: Both of them should be smaller than our result, but the second one should be closer to our one. The deviation of the theoretical results to our simulation result can be explained by the fact that in the Langevin model there are no interactions between the particles incorporated at all and the more advanced (mean-field-type) one of Ref. <a href='#[1]'>[1]</a> do not take occurring cluster formations into account but assumes a homogeneous distribution of the particles. For higher values of the volume fraction $\phi$ and the dipolar interaction parameter $\lambda$ the deviations will increase as the cluster formation will become more pronounced. 
Magnetization curve of a 3D system
At the end of this tutorial we now want to sample the magnetization curve of a three dimensional system and compare the results with analytical solutions. Again we will compare with the Langevin function but also with the approximation of Ref. <a href='#[2]'>[2]</a> (see also Ref. <a href='#[1]'>[1]</a> for the right coefficients) which takes the dipole-dipole interaction into account. For this approximation, which is a modified mean-field theory based on the pair correlation function, the Langevin parameter $\alpha$ is replaced by
\begin{equation}
    \alpha' = \alpha + \chi_\mathrm{L}~L(\alpha) + \frac{\chi_\mathrm{L}^{2}}{16} L(\alpha) \frac{\mathrm{d} L(\alpha)}{\mathrm{d}\alpha}
\end{equation}
where $\chi_\mathrm{L}$ is the Langevin susceptibility
\begin{equation}
   \chi_\mathrm{L} = \frac{N}{V}\frac{\mu_0 \mu^2}{3k_\mathrm{B}T} = 8 \cdot \lambda \cdot \phi
\end{equation}
Analogous to part II we start at zero external magnetic field and increase the external field successively. At every value of the external field we sample the total dipole moment which is proportional to the magnetization as we have a fixed volume.
First we create a list of values of the Langevin parameter $\alpha$. As we already sampled the magnetization at zero external field in the last section, we take this value and continue with the sampling of an external field unequal zero
End of explanation
"""


# remove all constraints
system.constraints.clear()

# array for magnetizations in field direction
magnetizations = np.full_like(alphas, np.nan)

# use result for alpha=0 from previous chapter
magnetizations[0] = np.average(dipm_tot)

# number of loops for sampling
loops_m = 500

for ndx, alpha in enumerate(pbar := tqdm.tqdm(alphas)):
    pbar.set_description(f"Sampling for α={alpha:.2f}")
    if alpha == 0:
        continue
    # set magnetic field constraint
    H_dipm = (alpha * kT)
    H_field = [H_dipm, 0, 0]
    H_constraint = espressomd.constraints.HomogeneousMagneticField(H=H_field)
    system.constraints.add(H_constraint)

    # equilibration
    for i in range(equil_rounds):
        system.integrator.run(equil_steps)

    # sampling
    magn_temp = np.full(loops_m, np.nan)
    for i in range(loops_m):
        system.integrator.run(20)
        magn_temp[i] = dipm_tot_calc.calculate()[0]

    # save average magnetization
    magnetizations[ndx] = np.mean(magn_temp)

    # remove constraint
    system.constraints.clear()

"""
Explanation: Now for each value in this list we sample the total dipole moment / magnetization of the system for a while. Keep in mind that we only the current orientation of the dipole moments, i.e. the unit vector of the dipole moments, is saved in the particle list but not their magnitude. Thus we have to use $H\cdot \mu$ as the external magnetic field, where $\mu$ is the magnitude of a single magnetic dipole moment.
We will apply the field in x-direction using the class <tt>constraints</tt> of ESPResSo.
As in part II we use the same system for every value of the Langevin parameter $\alpha$. Thus we use that the system is already pre-equilibrated from the previous run so we save some equilibration time. For scientific purposes one would use a new system for every value for the Langevin parameter to ensure that the systems are independent and no correlation effects are measured. Also one would perform more than just one simulation for each value of $\alpha$ to increase the precision of the results.
End of explanation
"""


# Langevin function
def L(y):
    return np.cosh(y) / np.sinh(y) - 1 / y

#  second order mean-field-model from Ref. [2]
def alpha_mean_field(alpha, dip_lambda, phi):
    chi = 8. * dip_lambda * phi
    return alpha + chi * L(alpha) + chi**2. / 16. * L(alpha) * (1. / alpha**2. - 1. / np.sinh(alpha)**2.)

"""
Explanation: Now we define the Langevin function and the modified mean-field-approximation of the Langevin parameter of Ref. <a href='#[2]'>[2]</a>
End of explanation
"""


y = np.arange(0.01, 10, 0.1)

initial_susceptibility = system.volume() * kT * chi / (N * mu_0 * dipm**2)

plt.figure(figsize=(10, 10))
plt.ylim(0, 1.)
plt.xlabel(r'$\alpha$', fontsize=20)
plt.ylabel(r'$M^*$', fontsize=20)
plt.plot(y, L(y), label='Langevin function')
plt.plot(y, L(alpha_mean_field(y, dip_lambda, phi)),
         label='modified mean-field-theory')
plt.plot(alphas, magnetizations / N, 'o', label='simulation results')
plt.plot(y, initial_susceptibility * y,
         label=r'linear approximation at $\alpha = 0$')
plt.legend(fontsize=20)
plt.show()

"""
Explanation: We also want to plot the linear approximation at $\alpha = 0$ to see for which values of $\alpha$ this approximation holds. We use the initial susceptibility calculated in the first chapter of this part as the gradient. As we want the gradient of $M^*$ with respect to $\alpha$ which fulfills the relation
\begin{equation}
    \frac{\partial M^*}{\partial \alpha} = \frac{1}{M_\mathrm{sat}}\frac{\partial M}{\partial \left( \frac{\mu_0\mu}{k_\mathrm{B}T} H\right)} = \frac{k_\mathrm{B}T~V}{\mu_0\mu^2N}\frac{\partial M}{\partial H} = \frac{k_\mathrm{B}T~V}{\mu_0\mu^2N}~\chi
\end{equation}
we have to scale our calculated initial susceptibility $\chi_{init}$ by a factor to get it in our dimensionless units.
Now we plot the resulting curves together with our simulation results and the linear approximation
End of explanation
"""


