Resolve "add support for cuboid pywasp structure in xrsite from_pywasp_pwc"

py_wake/site/xrsite.py

@@ -93,7 +93,6 @@ class XRSite(Site):
             return x_i * 0
 
     def interp(self, var, coords, deg=False):
-
         # Interpolate via EqDistRegGridInterpolator (equidistance regular grid interpolator) which is much faster
         # than xarray.interp.
         # This function is comprehensive because var can contain any combinations of coordinates (i or (xy,h)) and wd,ws
@@ -321,35 +320,54 @@ class XRSite(Site):
             "A", "k", and "wdfreq".
 
         """
-        renames = {
+        pwc = pwc.copy()
+
+        # Drop coordinates that are not needed
+        for coord in ["sector_floor", "sector_ceil", "crs"]:
+            if coord in pwc.coords:
+                pwc = pwc.drop_vars(coord)
+
+        # Get the spatial dims
+        if "point" in pwc.dims:
+            xyz_dims = ("point",)
+            xy_dims = ("point",)
+        elif all(d in pwc.dims for d in ["west_east", "south_north"]):
+            xyz_dims = ("west_east", "south_north", "height")
+            xy_dims = ("west_east", "south_north")
+        else:  # pragma: no cover
+            raise ValueError(f"No spatial dimensions found on dataset!")
+
+        # Make the dimensin order as needed
+        pwc = pwc.transpose(*xyz_dims, "sector", ...)
+
+        ws_mean = xr.apply_ufunc(
+            weibull.mean, pwc["A"], pwc["k"], dask="allowed"
+        )
+
+        pwc["Speedup"] = ws_mean / ws_mean.max(dim=xy_dims)
+
+        # Add TI if not already present
+        for var in ["turbulence_intensity"]:
+            if var not in pwc.data_vars:
+                pwc[var] = pwc["A"] * 0.0
+
+        new_names = {
             "wdfreq": "Sector_frequency",
             "A": "Weibull_A",
             "k": "Weibull_k",
             "turbulence_intensity": "TI",
             "sector": "wd",
             "point": "i",
+            "stacked_point": "i",
             "west_east": "x",
             "south_north": "y",
             "height": "h",
         }
 
-        pwc_renamed = pwc.rename({k: v for k, v in renames.items() if k in pwc})
-        pwc_renamed = pwc_renamed.transpose("i", "wd", ...)
-        for coord in ["sector_floor", "sector_ceil"]:
-            if coord in pwc_renamed.coords:
-                pwc_renamed = pwc_renamed.drop_vars(coord)
-
-        ws_mean = xr.apply_ufunc(
-            weibull.mean, pwc_renamed["Weibull_A"], pwc_renamed["Weibull_k"]
-        )
-        speedup = ws_mean / ws_mean.max(dim="i")
-        if "Speedup" not in pwc_renamed.data_vars:
-            pwc_renamed["Speedup"] = speedup
-
-        # Add TI and P if not already present
-        for var in ["P", "TI"]:
-            if var not in pwc_renamed.data_vars:
-                pwc_renamed[var] = pwc_renamed["Weibull_A"] * 0.0
+        pwc_renamed = pwc.rename({
+            old_name: new_name for old_name, new_name in new_names.items()
+            if old_name in pwc or old_name in pwc.dims
+        })
 
         return cls(pwc_renamed, **kwargs)
 

py_wake/tests/test_sites/test_xrsite.py

@@ -93,11 +93,17 @@ def complex_ws_site():
 
 
 @pytest.fixture
-def pywasp_pwc():
+def pywasp_pwc_point():
     pwc_file = tfp + "pwc_parquo_fictio_small.nc"
     return xr.open_dataset(pwc_file)
 
 
+@pytest.fixture
+def pywasp_pwc_cuboid():
+    pwc_file = tfp + "pwc_cuboid_small.nc"
+    return xr.open_dataset(pwc_file)
+
+
 def test_uniform_local_wind(uniform_site):
     site = uniform_site
     x_i = y_i = np.arange(5)
@@ -520,9 +526,9 @@ def test_interp_special_cases():
     npt.assert_array_almost_equal(ip1.data, ip2.data)
 
 
-def test_from_pywasp_pwc(pywasp_pwc):
+def test_from_pywasp_pwc_point(pywasp_pwc_point):
 
-    site = XRSite.from_pywasp_pwc(pywasp_pwc)
+    site = XRSite.from_pywasp_pwc(pywasp_pwc_point)
 
     A = np.array([
         [5.757, 6.088, 5.369],
@@ -590,6 +596,8 @@ def test_from_pywasp_pwc(pywasp_pwc):
     wd = np.linspace(0.0, 360.0, 13)
     i = np.arange(3)
 
+    assert "i" in site.ds.dims
+
     npt.assert_array_equal(site.ds.Weibull_A.values[..., :12], A.T)
     npt.assert_array_equal(site.ds.Weibull_k.values[..., :12], k.T)
     npt.assert_array_equal(site.ds.Sector_frequency.values[..., :12], wdfreq.T)
@@ -599,3 +607,14 @@ def test_from_pywasp_pwc(pywasp_pwc):
     npt.assert_array_equal(site.ds.i.values, i)
     npt.assert_array_equal(site.ds.wd.values, wd)
     npt.assert_allclose(site.ds.Speedup.values, speedups, rtol=1e-8)
+
+
+def test_from_pywasp_pwc_cuboid(pywasp_pwc_cuboid):
+
+    site = XRSite.from_pywasp_pwc(pywasp_pwc_cuboid)
+
+    assert all(d in site.ds.dims for d in ["x", "y", "h"])
+    assert site.ds.wd.size == 13
+    npt.assert_almost_equal(site.ds.Weibull_A.values[0, 0, 0, 0], 6.208063)
+    npt.assert_almost_equal(site.ds.Weibull_k.values[1, 1, 0, 1], 2.4550781)
+    npt.assert_almost_equal(site.ds.Sector_frequency.values[0, 0, 0, 0], 0.05975334)