From b7e5f7c002956f9570d126c60adf33965502db54 Mon Sep 17 00:00:00 2001
From: Sebastian Kloth <sebastian.kloth@pkm.tu-darmstadt.de>
Date: Mon, 19 Feb 2024 15:28:38 +0100
Subject: [PATCH] Finished data management example

---
 examples/data_management.py | 73 +++++++++++++++++++++++++++++++++----
 1 file changed, 66 insertions(+), 7 deletions(-)

diff --git a/examples/data_management.py b/examples/data_management.py
index 769db88..c29e3d3 100644
--- a/examples/data_management.py
+++ b/examples/data_management.py
@@ -1,15 +1,74 @@
-DATA_DIR = "/data/skloth/results/workshop_data"
-np.save(f"{DATA_DIR}/msd.npy", result_msd_nojump)
-print(np.load(f"{DATA_DIR}/msd.npy"))
+# %%
+import numpy as np
+import matplotlib.pyplot as plt
 
+import mdevaluate as md
+
+DATA_DIR = "/data/skloth/python_packages/mdevaluate_examples/data"
+path_to_sim = "/data/skloth/sim/IL_water/C4MIM_BF4_x50/T300_nvt"
+trajectory = md.open(path_to_sim, topology="run.tpr", trajectory="out/traj_full.xtc")
+
+oxygen_water = trajectory.subset(atom_name="OW", residue_name="SOL")
+cation_indices = trajectory.subset(residue_name="amim").atom_subset.indices
+cation_com = md.coordinates.center_of_masses(trajectory, atom_indices=cation_indices)
+anion_com = trajectory.subset(atom_name="B", residue_name="BF4")
+
+
+# %%
+time, msd_oxygen_water = md.correlation.shifted_correlation(
+    md.correlation.msd, oxygen_water.nojump, segments=100, skip=0.1, average=True
+)
+time, msd_cation_com = md.correlation.shifted_correlation(
+    md.correlation.msd, cation_com.nojump, segments=100, skip=0.1, average=True
+)
+time, msd_anion_com = md.correlation.shifted_correlation(
+    md.correlation.msd, anion_com.nojump, segments=100, skip=0.1, average=True
+)
+
+plt.figure()
+plt.plot(time, msd_oxygen_water, ".", label="Water")
+plt.plot(time, msd_cation_com, ".", label="Cation")
+plt.plot(time, msd_anion_com, ".", label="Anion")
+plt.legend()
+plt.xscale("log")
+plt.yscale("log")
+plt.xlabel(r"$t$ in ps", fontsize=16)
+plt.ylabel(r"$\langle r^2\rangle(t)$ in nm$^2$", fontsize=16)
+plt.show()
+plt.close()
+
+# %%
 import pandas as pd
-msd_df = pd.DataFrame({"t": result_msd[0], "msd_water": result_msd[1]})
+msd_df = pd.DataFrame(
+    {"t": time,
+     "msd_water": msd_oxygen_water,
+     "msd_cation": msd_cation_com,
+     "msd_anion": msd_anion_com
+     }
+)
 print(msd_df)
-msd_df["T"] = 300
 msd_df["cation"] = "C4mim"
 msd_df["anion"] = "DCA"
+msd_df["x_w"] = 50
+msd_df["T"] = 300
 print(msd_df)
+
+
+# %%
+def add_description(dataframe, cation, anion, x_w, T):
+    dataframe["cation"] = cation
+    dataframe["anion"] = anion
+    dataframe["x_w"] = x_w
+    dataframe["T"] = T
+    return dataframe
+
+
+msd_df = add_description(msd_df, "C4mim", "DCA", 50, 300)
+print(msd_df)
+
+#%%
 msd_df.to_hdf(f"{DATA_DIR}/result.h5", key="msd")
+
 msd_df = pd.read_hdf(f"{DATA_DIR}/result.h5", key="msd")
-print(msd_df.info())
-md.utils.cleanup_h5(f"{DATA_DIR}/result.h5")
\ No newline at end of file
+
+md.utils.cleanup_h5(f"{DATA_DIR}/result.h5")