Add 3D sp graph

notebooks/lib.py

@@ -68,6 +68,16 @@ def read_load(load_dir: str, reader=read_xy_alt):
     return pd.concat([reader(path) for path in paths])
 
 
+def augment_read_with_sp(inner_reader):
+    def hoc(path: str):
+        probability = float(Path(path).parent.name)
+        df = inner_reader(path)
+        df['probability'] = probability
+        return df
+
+    return hoc
+
+
 def read_sp_xy(specific_probability_dir: str):
     probability = float(Path(specific_probability_dir).name)
     df = read_load(specific_probability_dir)
@@ -76,11 +86,16 @@ def read_sp_xy(specific_probability_dir: str):
     return df
 
 
-def read_sp(sp_dir: str):
+def read_sp(sp_dir: str, inner_reader=read_xy_alt):
     if not Path(sp_dir).exists():
         raise Exception("Root does not exist")
 
-    return pd.concat([read_sp_xy(specific_probability_dir) for specific_probability_dir in glob(f'{sp_dir}/*')])
+    reader = augment_read_with_sp(inner_reader)
+
+    return pd.concat([
+        read_load(specific_probability_dir, reader)
+        for specific_probability_dir in glob(f'{sp_dir}/*')
+    ])
 
 
 def convergent_tail_index(series, tol):
@@ -126,3 +141,20 @@ def mean_across(df):
         .replace([np.inf, -np.inf], np.nan)
 
     return data
+
+
+def aggregate_sp_fd(df):
+    by_run = df.groupby(['probability', 'N'])
+
+    by_probability = by_run.agg(
+        overall_fd=('fd', lambda fd: np.mean(fd[-100:])),
+        overall_fd_std=('fd', 'std')
+    ).reset_index().groupby('probability')
+
+    data = by_probability.agg(
+        fd=('overall_fd', 'mean'),
+        # TODO Check stats
+        fd_std=('overall_fd_std', lambda std: np.sqrt(np.mean(np.square(std))))
+    )
+
+    return data

notebooks/sp-fd-3d.py

@@ -0,0 +1,18 @@
+from matplotlib import pyplot as plt
+
+from notebooks.lib import read_sp, read_xyz_alt, aggregate_sp_fd
+
+data_3d_sp = read_sp("../data/rust-3d-offaxis-sp", read_xyz_alt)
+sp_fd_data = aggregate_sp_fd(data_3d_sp)
+
+# %%
+plt.fill_between(sp_fd_data.index, sp_fd_data.fd - sp_fd_data.fd_std, sp_fd_data.fd + sp_fd_data.fd_std, alpha=0.2, label=f"Standard error band")
+plt.plot(sp_fd_data.index, sp_fd_data.fd, color='tab:blue', label='fd mean, seeds = 100')
+
+plt.xlabel("$p_{stick}$")
+plt.ylabel("$fd$")
+plt.legend()
+
+plt.savefig('../figures/sp-fd-3d.svg')
+plt.savefig('../figures/sp-fd-3d.png')
+plt.show()

notebooks/sp-fd.py

@@ -1,25 +1,13 @@
-import numpy as np
 from matplotlib import pyplot as plt
 
-from notebooks.lib import read_sp
+from notebooks.lib import read_sp, aggregate_sp_fd
 
-c_sp = read_sp("../data/stick-probability")
-by_run = c_sp.groupby(['probability', 'N'])
-
-by_probability = by_run.agg(
-    overall_fd=('fd', lambda fd: np.mean(fd[-100:])),
-    overall_fd_std=('fd', 'std')
-).reset_index().groupby('probability')
-
-ggg = by_probability.agg(
-    fd=('overall_fd', 'mean'),
-    # TODO Check stats
-    fd_std=('overall_fd_std', lambda std: np.sqrt(np.mean(np.square(std))))
-)
+data_2d_sp = read_sp("../data/stick-probability")
+sp_fd_data = aggregate_sp_fd(data_2d_sp)
 
 # %%
-plt.fill_between(ggg.index, ggg.fd - ggg.fd_std, ggg.fd + ggg.fd_std, alpha=0.2, label=f"Standard error band")
-plt.plot(ggg.index, ggg.fd, color='tab:blue', label='fd mean, seeds = 100')
+plt.fill_between(sp_fd_data.index, sp_fd_data.fd - sp_fd_data.fd_std, sp_fd_data.fd + sp_fd_data.fd_std, alpha=0.2, label=f"Standard error band")
+plt.plot(sp_fd_data.index, sp_fd_data.fd, color='tab:blue', label='fd mean, seeds = 100')
 
 plt.xlabel("$p_{stick}$")
 plt.ylabel("$fd$")