From 03e65ef0613e5238631e8af3abd5cf091fe5b28a Mon Sep 17 00:00:00 2001
From: Michael Uleysky <uleysky@poi.dvo.ru>
Date: Tue, 11 Apr 2023 15:02:40 +1000
Subject: [PATCH] Added support for velocity variables

---
 include/layereddata.h |  2 ++
 src/layereddata.cpp   | 39 ++++++++++++++++++++++++---------------
 2 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/include/layereddata.h b/include/layereddata.h
index 2d03338..1f79e43 100644
--- a/include/layereddata.h
+++ b/include/layereddata.h
@@ -227,6 +227,8 @@ class LayeredData
   if(stname == "ocean_mixed_layer_thickness_defined_by_sigma_theta") return "mld";
   if(stname == "sea_surface_height_above_geoid") return "ssh";
   if(stname == "sea_surface_elevation") return "ssh";
+  if(stname == "eastward_sea_water_velocity") return "u";
+  if(stname == "northward_sea_water_velocity") return "v";
   if(stname == "upward_sea_water_velocity") return "w";
   return "";
  }
diff --git a/src/layereddata.cpp b/src/layereddata.cpp
index 14dd073..7828919 100644
--- a/src/layereddata.cpp
+++ b/src/layereddata.cpp
@@ -22,6 +22,9 @@ MString LayeredData::Info() const
  if(vars.contains("ptemp") && vars.contains("sal")) vars.emplace("temp");
  if(vars.contains("temp") && vars.contains("sal")) vars.emplace("ptemp");
 
+ if(vars.contains("u") && vars.contains("v")) vars.emplace("U");
+ if(vars.contains("u") && vars.contains("v")) vars.emplace("U2");
+
  MString svars;
  {
   bool first = true;
@@ -146,6 +149,7 @@ std::pair<const BaseParameters*, MString> LayeredData::Parameters(michlib_intern
   if(ppar->varname == "temp" && HaveVar("ptemp") && HaveVar("sal")) varexist = true;
   if(ppar->varname == "ptemp" && HaveVar("temp") && HaveVar("sal")) varexist = true;
   if(ppar->varname == "pdens" && (HaveVar("ptemp") || HaveVar("temp")) && HaveVar("sal")) varexist = true;
+  if((ppar->varname == "U" || ppar->varname == "U2") && HaveVar("u") && HaveVar("v")) varexist = true;
   if(!varexist) return {nullptr, "Variable " + ppar->varname + " not exists in this dataset"};
  }
  if(args.contains("layer")) ppar->layer = args.at("layer").ToInteger<size_t>();
@@ -217,11 +221,6 @@ LayeredData::Data LayeredData::Read(const BaseParameters* ip, size_t i) const
    auto out = temp;
    for(size_t ind = 0; ind < out.N(); ind++)
    {
-    if(out.Lon(ind) != temp.Lon(ind) || out.Lon(ind) != sal.Lon(ind) || out.Lat(ind) != temp.Lat(ind) || out.Lat(ind) != sal.Lat(ind))
-    {
-     michlib::errmessage("Internal error");
-     exit(1);
-    }
     if(temp.IsFill(ind) || sal.IsFill(ind))
      out.V(ind) = out.Fillval();
     else
@@ -241,11 +240,6 @@ LayeredData::Data LayeredData::Read(const BaseParameters* ip, size_t i) const
    auto out = temp;
    for(size_t ind = 0; ind < out.N(); ind++)
    {
-    if(out.Lon(ind) != temp.Lon(ind) || out.Lon(ind) != sal.Lon(ind) || out.Lat(ind) != temp.Lat(ind) || out.Lat(ind) != sal.Lat(ind))
-    {
-     michlib::errmessage("Internal error");
-     exit(1);
-    }
     if(temp.IsFill(ind) || sal.IsFill(ind))
      out.V(ind) = out.Fillval();
     else
@@ -266,11 +260,6 @@ LayeredData::Data LayeredData::Read(const BaseParameters* ip, size_t i) const
    auto out = temp;
    for(size_t ind = 0; ind < out.N(); ind++)
    {
-    if(out.Lon(ind) != temp.Lon(ind) || out.Lon(ind) != sal.Lon(ind) || out.Lat(ind) != temp.Lat(ind) || out.Lat(ind) != sal.Lat(ind))
-    {
-     michlib::errmessage("Internal error");
-     exit(1);
-    }
     if(temp.IsFill(ind) || sal.IsFill(ind))
      out.V(ind) = out.Fillval();
     else
@@ -279,6 +268,26 @@ LayeredData::Data LayeredData::Read(const BaseParameters* ip, size_t i) const
    }
    return out;
   }
+  // U and U2 from u and v
+  if(p->varname == "U" || p->varname == "U2")
+  {
+   bool              square = p->varname == "U2";
+   struct Parameters params = *p;
+   params.varname           = "u";
+   auto u                   = Read(&params, i);
+   params.varname           = "v";
+   auto v                   = Read(&params, i);
+   if(!(u && v)) return Data();
+   auto out = u;
+   for(size_t ind = 0; ind < out.N(); ind++)
+   {
+    if(u.IsFill(ind) || v.IsFill(ind))
+     out.V(ind) = out.Fillval();
+    else
+     out.V(ind) = square ? (u(ind) * u(ind) + v(ind) * v(ind)) : michlib::Hypot(u(ind), v(ind));
+   }
+   return out;
+  }
   return Data();
  }