From 285b9476bcd16d7be3cc8b17cb86760566a19646 Mon Sep 17 00:00:00 2001
From: "Edgar P. Burkhart" <git@edgarpierre.fr>
Date: Mon, 28 Mar 2022 10:59:35 +0200
Subject: [PATCH] Added velocity to sws_ola

---
 olaflow/processing/olaflow.py | 15 +++++++++++++++
 olaflow/processing/sws_ola.py | 30 +++++++++++++++++++++++++++---
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/olaflow/processing/olaflow.py b/olaflow/processing/olaflow.py
index 520289b..48d9165 100644
--- a/olaflow/processing/olaflow.py
+++ b/olaflow/processing/olaflow.py
@@ -25,6 +25,21 @@ class OFModel:
                 )
             )
 
+    def write_vector_field(self, field, values):
+        with open(self._root.joinpath("0", field), "r") as aw_file:
+            aw_raw = aw_file.read()
+
+        with open(self._root.joinpath("0", field), "w") as aw_file:
+            aw_file.write(
+                re.sub(
+                    r"(?<=\(\n).*?(?=\n\))",
+                    "\n".join(map(lambda x: f"({' '.join(x.astype('str'))})", values)),
+                    aw_raw,
+                    count=1,
+                    flags=re.S,
+                )
+            )
+
     @property
     def x(self):
         return self._x
diff --git a/olaflow/processing/sws_ola.py b/olaflow/processing/sws_ola.py
index 93a0f11..c10673c 100644
--- a/olaflow/processing/sws_ola.py
+++ b/olaflow/processing/sws_ola.py
@@ -5,6 +5,7 @@ import pathlib
 
 import numpy as np
 from scipy import interpolate
+import matplotlib.pyplot as plt
 
 from .olaflow import OFModel
 
@@ -22,9 +23,16 @@ config = configparser.ConfigParser()
 config.read("config.ini")
 
 sws_out = pathlib.Path(config.get("swash", "np_out"))
-x = np.load(sws_out.joinpath("xp.npy"))
-t = np.load(sws_out.joinpath("tsec.npy"))
-watl = np.load(sws_out.joinpath("watl.npy"))
+def data(var):
+    return np.load(sws_out.joinpath(f"{var}.npy"))
+
+x = data("xp")
+t = data("tsec")
+
+watl = data("watl")
+zk = data("zk")
+velk = data("velk")
+vz = data("vz")
 
 olaflow_root = pathlib.Path(config.get("olaflow", "root"))
 model = OFModel(olaflow_root)
@@ -33,4 +41,20 @@ model.read_mesh()
 watl_t = interpolate.interp1d(x, watl[680])
 alpha_water = np.where(model.z < watl_t(model.x), 1, 0)
 
+zk_t = interpolate.interp1d(x, zk[680])
+velk_t = interpolate.interp1d(x, velk[680,:,0,:])(model.x)
+vz_t = interpolate.interp1d(x, vz[680])(model.x)
+zk_tl = zk_t(model.x)
+
+print(velk.shape)
+ux = np.zeros(model.x.shape)
+uy = np.zeros(model.x.shape)
+uz = np.zeros(model.x.shape)
+print(zk_tl.shape, velk.shape, vz.shape)
+for zk_l, velk_l, vz_l in zip(zk_tl, velk_t, vz_t):
+    ux = np.where(model.z < zk_l, velk_l, ux)
+    uz = np.where(model.z < zk_l, vz_l, uz)
+
+print(np.stack((ux,uy,uz)).T)
 model.write_field("alpha.water", alpha_water)
+model.write_vector_field("U", np.stack((ux, uy, uz)).T)