107 lines
2.3 KiB
Python
107 lines
2.3 KiB
Python
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import pandas as pd
|
|
import argparse
|
|
|
|
|
|
def plot_2d(df: pd.DataFrame):
|
|
df = pd.read_csv("temp.csv")
|
|
|
|
fig = plt.figure()
|
|
plt.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95)
|
|
ax = fig.add_subplot()
|
|
|
|
# Analytically computed solution
|
|
ts = df['tb']
|
|
x2s = -ts / 2
|
|
x1s = -ts * x2s
|
|
|
|
ax.scatter(x1s, x2s)
|
|
|
|
ax.set_xlabel("x1")
|
|
ax.set_ylabel("x2")
|
|
# ax.set_zlabel("t")
|
|
|
|
df = df.reset_index()
|
|
|
|
for _, row in df.iterrows():
|
|
x1b = row['x1b']
|
|
x2b = row['x2b']
|
|
# x1p = row['x1p']
|
|
# x2p = row['x2p']
|
|
x1e = row['x1e']
|
|
x2e = row['x2e']
|
|
x1n = row['x1n']
|
|
x2n = row['x2n']
|
|
|
|
ax.plot([x1b, x1e], [x2b, x2e], 'b')
|
|
ax.plot([x1e, x1n], [x2e, x2n], 'r')
|
|
|
|
plt.show()
|
|
|
|
|
|
def plot_3d(df: pd.DataFrame):
|
|
df = pd.read_csv("temp.csv")
|
|
|
|
fig = plt.figure()
|
|
plt.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95)
|
|
ax = fig.add_subplot(projection='3d')
|
|
|
|
# Analytically computed solution
|
|
ts = df['tb']
|
|
x2s = -ts / 2
|
|
x1s = -ts * x2s
|
|
|
|
ax.scatter(x1s, x2s)
|
|
ax.scatter(x1s, x2s, ts)
|
|
|
|
ax.set_xlabel("x1")
|
|
ax.set_ylabel("x2")
|
|
ax.set_zlabel("t")
|
|
|
|
df = df.reset_index()
|
|
|
|
for _, row in df.iterrows():
|
|
x1b = row['x1b']
|
|
x2b = row['x2b']
|
|
tb = row['tb']
|
|
# x1p = row['x1p']
|
|
# x2p = row['x2p']
|
|
# tp = row['tp']
|
|
x1e = row['x1e']
|
|
x2e = row['x2e']
|
|
te = row['te']
|
|
x1n = row['x1n']
|
|
x2n = row['x2n']
|
|
tn = row['tn']
|
|
|
|
ax.plot([x1b, x1e], [x2b, x2e], 'b', zs=[tb, te])
|
|
ax.plot([x1e, x1n], [x2e, x2n], 'r', zs=[te, tn])
|
|
|
|
plt.show()
|
|
|
|
|
|
def main():
|
|
# Parse command line arguments
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument("--input", "-i", type=str, required=True,
|
|
default="temp.csv", help="Filename of the csv file")
|
|
parser.add_argument("--projection", "-p", type=str, required=False,
|
|
default="2d", help="2d or 3d plot")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Plot data
|
|
df = pd.read_csv(args.input)
|
|
|
|
if args.projection == "2d":
|
|
plot_2d(df)
|
|
elif args.projection == "3d":
|
|
plot_3d(df)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|