Homec4science
Diffusion DD_1D-SRA (main) Browse

DD_1D-SRA/CODE/DATASET/plot_dataset.py
f855cb73c1c5main

Last Change
Raw File

plot_dataset.py
View Options

# Plot the dataset from a pickle file
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
# Location of the dataset
path = r'C:/Travail/EPFL/Cours/M2/Projet CSE/LAMMPS/Dataset/'
def multiline(xs, ys, c, ax=None, **kwargs):
"""
Plot lines with different colorings
Parameters
----------
xs : iterable container of x coordinates
ys : iterable container of y coordinates
c : iterable container of numbers mapped to colormap
ax (optional): Axes to plot on.
kwargs (optional): passed to LineCollection
Notes:
len(xs) == len(ys) == len(c) is the number of line segments
len(xs[i]) == len(ys[i]) is the number of points for each line (indexed by i)
Returns
-------
lc : LineCollection instance.
"""
# find axes
ax = plt.gca() if ax is None else ax
# create LineCollection
segments = [np.column_stack([x, y]) for x, y in zip(xs, ys)]
lc = LineCollection(segments, **kwargs)
# set coloring of line segments
# Note: I get an error if I pass c as a list here... not sure why.
lc.set_array(np.asarray(c))
# add lines to axes and rescale
# Note: adding a collection doesn't autoscalee xlim/ylim
ax.add_collection(lc)
ax.autoscale()
return lc
df = pd.read_pickle(path + 'dataset')
col = df.columns
dataset = {}
for i in range(int(len(col)/2)):
Pressure = int(col[2*i].split()[1][2:-2])
Strain = df[col[2*i]].dropna()
Stress = df[col[2*i+1]].dropna()
dataset[Pressure] = np.array([[e,s] for e,s in zip(Strain,Stress)])
pressures = [p for p in dataset.keys()]
pressures.sort()
pressures_kPa = [int(p/1e3) for p in dataset.keys()]
pressures_kPa.sort()
n_lines = len(pressures)
xs = np.array([dataset[p][:,0] for p in pressures],dtype=object)
ys = np.array([dataset[p][:,1]/1e6 for p in pressures],dtype=object)
colors = np.arange(n_lines)
fig, ax = plt.subplots()
lc = multiline(xs, ys, pressures_kPa, cmap='bwr', lw=2)
axcb = fig.colorbar(lc)
axcb.set_label('Confinement pressure [kPa]')
ax.set_title('Constitutive law')
ax.set_xlabel('Strain')
ax.set_ylabel('Stress [MPa]')
plt.savefig(path + 'constitutive.pdf')
plt.show()
c4science · Help