#!/usr/bin/env python
# coding: utf-8

# # Analysing simulation results: Domain wall pair conversion
# 
# ## Simulation
# ### Problem description
# 
# We want to simulate a domain wall conversion in a two-dimensional thin film sample with:
# 
# - exchange energy constant $A = 15 \,\text{pJ}\,\text{m}^{-1}$,
# - Dzyaloshinskii-Moriya energy constant $D = 3 \,\text{mJ}\,\text{m}^{-2}$,
# - uniaxial anisotropy constant $K = 0.5 \,\text{MJ}\,\text{m}^{-3}$ with $\hat{\mathbf{u}} = (0, 0, 1)$ in the out of plane direction,
# - gyrotropic ratio $\gamma = 2.211 \times 10^{5} \,\text{m}\,\text{A}^{-1}\,\text{s}^{-1}$, and
# - Gilbert damping $\alpha=0.3$.
# 
# Please carry out the following steps:
# 
# 1. Create the following geometry with discretisation cell size $(2 \,\text{nm}, 2 \,\text{nm}, 2 \,\text{nm})$.
#     
#     <img src="figures/dw_pair_conversion_geometry.png" width="400">  
# 
# 2. Initialise the magnetisation so that when relaxes, a domain pair is present in the narrower part of the geometry.
# 
# 3. Relax the system. Is a domain wall pair contained in the constrained part?
# 
# 4. Apply the spin polarised current in the positive $x$ direction with velocity $\mathbf{u} = (400, 0, 0) \,\text{m}\,\text{s}^{-1}$, with $\beta=0.5$.
# 
# 5. Evolve the system over $0.2 \,\text{ns}$. What did you get? [1]
# 
# ### References
# 
# [1] Zhou, Y., & Ezawa, M. (2014). A reversible conversion between a skyrmion and a domain-wall pair in a junction geometry. *Nature Communications* **5**, 8. https://doi.org/10.1038/ncomms5652
# 
# 
# ### Solution

# In[1]:


import oommfc as oc
import discretisedfield as df
import micromagneticmodel as mm

Ms = 5.8e5  # saturation magnetisation (A/m)
A = 15e-12  # exchange energy constant (J/)
D = 3e-3  # Dzyaloshinkii-Moriya energy constant (J/m**2)
K = 0.5e6  # uniaxial anisotropy constant (J/m**3)
u = (0, 0, 1)  # easy axis
gamma0 = 2.211e5  # gyromagnetic ratio (m/As)
alpha = 0.3  # Gilbert damping

system = mm.System(name="dw_pair_conversion")
system.energy = (
    mm.Exchange(A=A)
    + mm.DMI(D=D, crystalclass="Cnv_z")
    + mm.UniaxialAnisotropy(K=K, u=u)
)
system.dynamics = mm.Precession(gamma0=2.211e5) + mm.Damping(alpha=alpha)

p1 = (0, 0, 0)
p2 = (150e-9, 50e-9, 2e-9)
cell = (2e-9, 2e-9, 2e-9)

region = df.Region(p1=p1, p2=p2)
mesh = df.Mesh(region=region, cell=cell)


def Ms_fun(pos):
    x, y, z = pos
    if x < 50e-9 and (y < 15e-9 or y > 35e-9):
        return 0
    else:
        return Ms


def m_init(pos):
    x, y, z = pos
    if 30e-9 < x < 40e-9:
        return (0.1, 0.1, -1)
    else:
        return (0.1, 0.1, 1)


system.m = df.Field(mesh, nvdim=3, value=m_init, norm=Ms_fun, valid="norm")

system.m.z.sel("z").mpl.scalar()


# In[2]:


md = oc.MinDriver()
md.drive(system)

system.m.z.sel("z").mpl.scalar()


# In[3]:


ux = 400  # velocity in x direction (m/s)
beta = 0.5  # non-adiabatic STT parameter

system.dynamics += mm.ZhangLi(u=ux, beta=beta)


# In[4]:


td = oc.TimeDriver()
td.drive(system, t=0.2e-9, n=200, verbose=2)

system.m.orientation.z.sel("z").mpl.scalar()


# As a result, we got a skyrmion formed in the wider region.
# 
# ## Data analysis
# 
# We can use the `micromagneticdata` package for post processing of simulation data.
# 
# We get access to:
# - initial magnetisation
# - magnetisation snapshots from the simulation, e.g. at the timesteps specified in `time_driver.drive(...)`
# - tabular data recorded during the simulation via `ubermagtable`, e.g. total energy as a function of time

# In[5]:


import micromagneticdata as mdata


# ### `Data`
# 
# We can access data from past simulation runs based on the system name used to run the simulation.

# In[6]:


data = mdata.Data("dw_pair_conversion")


# The `Data` object contains all simulation runs of the `System`, these are called drives. We can check how many drives does our data has:

# In[7]:


data.n


# We have used two drivers in our simulation, which has created two corresponding drives.
# 
# We can also extract the details using the `info` property.

# In[8]:


data.info


# ### `Drive`
# We can explicitly check the two drivers seperately by indexing the `Data` object.

# In[9]:


data[0].info


# A `drive` object gives acces to all data created from a single simulation run, e.g. call to `min_driver.drive()`
# 
# We can obtain the intial magnetisation with `m0`. We get a `discretisedfield.Field` and can use its plotting methods to visualise the initial configuration

# In[10]:


data[0].m0.orientation.sel("z").z.mpl.scalar()


# We can check the number of simulation snapshots saved for the drive (the initial magnetisation `m0` is always excluded):

# In[11]:


data[0].n


# For the first drive, the energy minimisation, we only have one single element, saved at the end of the energy minimisation.
# 
# To obtain the finial state we can access the first data element of the drive.

# In[12]:


data[0][0].orientation.sel("z").z.mpl.scalar()


# We can also use negative indices to access elements of the drive starting from the back.
# 
# In this case we only have a single element in the drive, so the first and the last element are identical:

# In[13]:


data[0][0] == data[0][-1]


# Let us now study the time drive (`data[1]`) in more detail. For convenience we assign it to a new name `time_drive` (and use negative indexing to select the last drive).

# In[14]:


time_drive = data[-1]
time_drive.info


# We can observe the final magnetisation state of the time driver by indexing the last drive, i.e., `time_drive[-1]`

# In[15]:


time_drive[-1].orientation.sel("z").z.mpl.scalar()


# We can create an interactive plot with a slider for the time to interactively investigate the conversion of the domain wall pair into a skyrmion.

# In[16]:


time_drive.hv(
    kdims=["x", "y"],
    scalar_kw={"clim": (-Ms, Ms), "cmap": "coolwarm"},
)


# ### `Table`
# 
# We also get access to all tabular data recorded during the drive. The data is made available as `pandas.Dataframe`.

# In[17]:


time_drive.table.data


# The `table` object provides a convenient plotting method. We can use it to visualise e.g. the total energy or the average magnetisation as a function of time. Ubermag, more precisely, `ubermagtable` automatically uses time for the x axis based on the metadata of the drive.

# In[18]:


time_drive.table.mpl(y=["E"])


# In[19]:


time_drive.table.mpl(y=["mx", "my", "mz"])


# ### Derived quantities using callbacks
# 
# We can also compute derived quantities for each time step and visualise them. As an example, we compute the topological charge density. Discretisedfield provides a function for this in `discretisedfield.tools`.
# 
# First, we compute the topolgical charge density for the final magnetisation state and plot it.

# In[20]:


final_magnetisation = time_drive[-1]
top_charge_last_step = df.tools.topological_charge_density(final_magnetisation.sel("z"))

top_charge_last_step.mpl.scalar(cmap="Blues")


# To compute a function for each time step, we use the `register_callback` method of `Drive`. We can pass in a function that will be called for each element of the drive. The function is only evaluated once we access an element of the drive. The `register_callback` method returns a new `Drive` object so that we still have access to the original data if required.
# 
# In this example we write a callback function that gets one magnetisation `m` and then selects a 2D slice normal to the `z` direction and computes the topological charge density.

# In[21]:


def top_charge_plane(m):
    return df.tools.topological_charge_density(m.sel("z"))


top_charge = time_drive.register_callback(top_charge_plane)


# We can now plot the data of the new `top_charge` drive object. We manually compute suitable colour limits based on the topological charge in the final state to avoid changes of the colour range when moving the slider. (The plotting function always has only access to the current frame and would re-compute a suitable but changing colour range for each frame, which makes it difficult to judge the changes in topological charge.)

# In[22]:


c_min = top_charge_last_step.array.min()
c_max = top_charge_last_step.array.max()

top_charge.hv.scalar(kdims=["x", "y"], clim=(c_min, c_max))


# ### Conversion to `xarray.DataArray`
# 
# The drive object also provides a method to convert all magnetisation data into a single `DataArray` suitable for further post-processing outside Ubermag.

# In[23]:


time_drive.to_xarray()