Repository: British-Oceanographic-Data-Centre/COAsT (language: python)

Issue:
606 circulation and fluxes in global model data jp
@jasontempestholt These are the fixes I did during the review of your Pull Request.
If you are happy with these, we (you) can update your branch with them. Then I will merge your branch into Develop.

Use "Files Changed" tab (5cm above) to see the differences:

0. updated to the latest version of Develop
1. changed a few names to make them more COAsT
2. abstracted and simplified your subsetting method in Gridded()
3. refactored the velocity_on_t_points so that the differences across indices are not a function of dimension ordering (I'm also hoping it might scale better with xarray backend magic)



Discussion / hints:

Fix patch:
diff --git a/coast/__init__.py b/coast/__init__.py
index 37279aa6..46a20af8 100644
--- a/coast/__init__.py
+++ b/coast/__init__.py
@@ -28,6 +28,7 @@
 from .data.copernicus import Copernicus, Product
 from ._utils.experiments_file_handling import experiments
 from ._utils.experiments_file_handling import nemo_filename_maker
+from .diagnostics.circulation import CurrentsOnT
 
 # Set default for logging level when coast is imported
 import logging
diff --git a/coast/data/gridded.py b/coast/data/gridded.py
index 46b3368a..786d74ec 100644
--- a/coast/data/gridded.py
+++ b/coast/data/gridded.py
@@ -63,18 +63,22 @@ def _setup_grid_obj(self, chunks, multiple, **kwargs):
 
         Args:
             chunks: This is a setting for xarray as to whether dask (parrell processing) should be on and how it works
-            multiple: falg to tell if we are loading one or more files
+            multiple: flag to tell if we are loading one or more files
             **kwargs: pass direct to loaded xarray dataset
+                lims = [x_dim index 1, x_dim_index 2, y_dim index 1, y_dim_index 2] - subset region defined from
+                                                                            lower left to upper right corners
+                calculate_bathymetry [boolean]: default-False
         """
         self.set_grid_vars()
         self.set_dimension_mapping()
         self.set_variable_mapping()
-
+        lims = kwargs.get("lims", [])
         if self.fn_data is not None:
             self.load(self.fn_data, chunks, multiple)
 
-        self.set_dimension_names(self.config.dataset.dimension_map)
-        self.set_variable_names(self.config.dataset.variable_map)
+            self.set_dimension_names(self.config.dataset.dimension_map)
+            self.set_variable_names(self.config.dataset.variable_map)
+            self.dataset = self.spatial_subset(self.dataset, lims)  # Trim data size if indices specified
 
         if self.fn_domain is None:
             self.filename_domain = ""  # empty store for domain fileanme
@@ -88,8 +92,9 @@ def _setup_grid_obj(self, chunks, multiple, **kwargs):
             for key, value in kwargs.items():
                 dataset_domain[key] = value
 
+            dataset_domain = self.spatial_subset(dataset_domain, lims)  # Trim domain size if indices specified
             if self.fn_data is not None:
-                dataset_domain = self.trim_domain_size(dataset_domain)
+                dataset_domain = self.trim_domain_size(dataset_domain)  # Trim domain size if self.data is smaller
             self.set_timezero_depths(
                 dataset_domain, **kwargs
             )  # THIS ADDS TO dataset_domain. Should it be 'return'ed (as in trim_domain_size) or is implicit OK?
@@ -189,7 +194,7 @@ def get_contour_complex(self, var, points_x, points_y, points_z, tolerance: int
         smaller = self.dataset[var].sel(z=points_z, x=points_x, y=points_y, method="nearest", tolerance=tolerance)
         return smaller
 
-    def set_timezero_depths(self, dataset_domain, calculate_bathymetry=False):
+    def set_timezero_depths(self, dataset_domain, **kwargs):
         """
         Calculates the depths at time zero (from the domain_cfg input file)
         for the appropriate grid.
@@ -204,6 +209,7 @@ def set_timezero_depths(self, dataset_domain, calculate_bathymetry=False):
         # keyword to allow calcution of bathymetry from scale factors
 
         # All bathymetry should now be mapped to bathy_metry
+        calculate_bathymetry = kwargs.get("calculate_bathymetry", False)
         try:
             if calculate_bathymetry:  # calculate bathymetry from scale factors
                 bathymetry, mask, time_mask = self.calc_bathymetry(dataset_domain)
@@ -674,6 +680,27 @@ def construct_density(
         except AttributeError as err:
             error(err)
 
+    def spatial_subset(self, dataset, lims):
+        """
+        Specify indices to subset the data. Subset region defined as a 2D box from lower left to upper right corners
+        lims = [x_dim index_1, x_dim_index_2, y_dim index_1, y_dim_index_2] -
+        Modifies self.dataset
+        """
+        if len(lims) == 4:  # if lims are provided take a subset
+            debug(f"Trimming by indices: {lims}")
+            # subsetting will wrap longitude across dateline if 1st longitude is larger than 2nd.
+            if "x_dim" in dataset.dims:
+                if lims[0] < lims[1]:  # usual case
+                    dataset = dataset.isel(y_dim=range(lims[2], lims[3]), x_dim=range(lims[0], lims[1]))
+                else:  # longitude  wrap around
+                    nx = dataset.dims["x_dim"]
+                    ds1 = dataset.isel(y_dim=range(lims[2], lims[3]), x_dim=range(lims[0], nx))
+                    ds2 = dataset.isel(y_dim=range(lims[2], lims[3]), x_dim=range(0, lims[1]))
+                    dataset = xr.concat([ds1, ds2], dim="x_dim")
+            else:
+                print("limits not used as only work with datasets having dimension x_dim")
+        return dataset
+
     def trim_domain_size(self, dataset_domain):
         """
         Trim the domain variables if the dataset object is a spatial subset
@@ -696,12 +723,22 @@ def trim_domain_size(self, dataset_domain):
             )
 
             # Find the corners of the cut out domain.
-            [j0, i0] = self.find_j_i_domain(
-                lat=self.dataset.nav_lat[0, 0], lon=self.dataset.nav_lon[0, 0], dataset_domain=dataset_domain
-            )
-            [j1, i1] = self.find_j_i_domain(
-                lat=self.dataset.nav_lat[-1, -1], lon=self.dataset.nav_lon[-1, -1], dataset_domain=dataset_domain
-            )
+            try:
+                [j0, i0] = self.find_j_i_domain(
+                    lat=self.dataset.latitude[0, 0], lon=self.dataset.longitude[0, 0], dataset_domain=dataset_domain
+                )
+                [j1, i1] = self.find_j_i_domain(
+                    lat=self.dataset.latitude[-1, -1], lon=self.dataset.longitude[-1, -1], dataset_domain=dataset_domain
+                )
+                debug(f"trim_domain_size(): USED dataset.longitude")
+            except:  # if called before variables are re-mapped. Not very pretty...
+                [j0, i0] = self.find_j_i_domain(
+                    lat=self.dataset.nav_lat[0, 0], lon=self.dataset.nav_lon[0, 0], dataset_domain=dataset_domain
+                )
+                [j1, i1] = self.find_j_i_domain(
+                    lat=self.dataset.nav_lat[-1, -1], lon=self.dataset.nav_lon[-1, -1], dataset_domain=dataset_domain
+                )
+                debug(f"trim_domain_size(): USED dataset.nav_lon")
 
             dataset_subdomain = dataset_domain.isel(y_dim=slice(j0, j1 + 1), x_dim=slice(i0, i1 + 1))
             return dataset_subdomain
diff --git a/coast/diagnostics/circulation.py b/coast/diagnostics/circulation.py
new file mode 100644
index 00000000..aca5690b
--- /dev/null
+++ b/coast/diagnostics/circulation.py
@@ -0,0 +1,148 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Dec 21 16:47:59 2022
+
+@author: jholt
+"""
+from ..data.gridded import Gridded
+import numpy as np
+import xarray as xr
+import matplotlib.pylab as plt
+from .._utils.logging_util import debug
+
+
+class CurrentsOnT(Gridded):
+    """
+    Methods for currents co-located on T-points
+
+    Use like coast.Gridded() with a domain_cfg file name to create gridded object
+
+    """
+
+    def __init(self, fn_domain=None, config=None, **kwargs):
+        gridded = Gridded(fn_domain=fn_domain, config=config, **kwargs)
+        self.dataset = gridded.dataset
+
+    def currents_on_t(self, gridded_u, gridded_v):
+        """
+        Adds co-located velocity components and speed to CurrentsOnT object
+
+        Avoids use of indices in the hope that it will be more parallelisable. However in index space the mappings
+        take this form:
+        u_on_t_points[:, :, :, 1:] = 0.5 * (
+            gridded_u.dataset.u_velocity.values[:, :, :, 1:] + gridded_u.dataset.u_velocity.values[:, :, :, :-1]
+            )
+
+        v_on_t_points[:, :, 1:, :] = 0.5 * (
+            gridded_v.dataset.v_velocity.values[:, :, 1:, :] + gridded_v.dataset.v_velocity.values[:, :, :-1, :]
+            )
+
+        For arrays of shape and order: (z_dim, t_dim, y_dim, x_dim)
+
+        Parameters
+        ----------
+        gridded_u : TYPE
+            DESCRIPTION.
+        gridded_v : TYPE
+            DESCRIPTION.
+
+        Returns
+        -------
+        None.
+
+        """
+        ds_u = self.dataset.copy(deep=True)
+        # U velocity on T-points
+        ds_u["ut_velocity"] = 0.5 * (gridded_u.dataset.u_velocity.shift(x_dim=1) + gridded_u.dataset.u_velocity)
+        # replace wrapped (1st) longitude coord with zero
+        _, _lon = xr.broadcast(gridded_u.dataset.u_velocity, gridded_u.dataset.longitude)
+        ds_u["ut_velocity"] = ds_u["ut_velocity"].where(
+            _lon != _lon.isel(x_dim=0), 0
+        )  # keep values except where lon(x_dim=0)
+        del _, _lon
+        ds_u.coords["latitude"] = self.dataset.latitude
+        ds_u.coords["longitude"] = self.dataset.longitude
+        ds_u.coords["depth_0"] = self.dataset.depth_0
+        try:
+            self.dataset["ut_velocity"] = ds_u.ut_velocity.drop("depthu")
+        except:
+            self.dataset["ut_velocity"] = ds_u.ut_velocity
+            debug("Did not find depthu variable to drop - to avoid conflicts in z_dim dimension")
+
+        # V velocity on T-points
+        ds_v = self.dataset.copy(deep=True)
+
+        ds_v["vt_velocity"] = 0.5 * (gridded_v.dataset.v_velocity.shift(y_dim=1) + gridded_v.dataset.v_velocity)
+        # replace wrapped (1st) latitude coord with zero
+        _, _lat = xr.broadcast(gridded_v.dataset.v_velocity, gridded_v.dataset.latitude)
+        ds_v["vt_velocity"] = ds_v["vt_velocity"].where(
+            _lat != _lat.isel(y_dim=0), 0
+        )  # keep values except where lat(y_dim=0)
+        del _, _lat
+        ds_v.coords["latitude"] = self.dataset.latitude
+        ds_v.coords["longitude"] = self.dataset.longitude
+        ds_v.coords["depth_0"] = self.dataset.depth_0
+        try:
+            self.dataset["vt_velocity"] = ds_v.vt_velocity.drop("depthv")
+        except:
+            self.dataset["vt_velocity"] = ds_v.vt_velocity
+            debug("Did not find depthv variable to drop - to avoid conflicts in z_dim dimension")
+
+        # Speed on T-points
+        self.dataset["speed_t"] = np.sqrt(
+            np.square(self.dataset["ut_velocity"]) + np.square(self.dataset["vt_velocity"])
+        )
+
+    def quick_plot(self, name, Vmax=0.16, Np=3, headwidth=4, scale=50, time_value=None, **kwargs):
+        """
+        plot surface circulation
+        direction: unit vector
+        magnitude: shaded
+        """
+        # %%
+        from matplotlib import cm
+        from matplotlib.colors import LinearSegmentedColormap
+
+        nx = self.dataset.x_dim.size
+        ny = self.dataset.y_dim.size
+
+        if time_value == None:
+            SP = np.squeeze(self.dataset.speed_t)
+            US = np.squeeze(self.dataset.ut_velocity / SP)
+            VS = np.squeeze(self.dataset.vt_velocity / SP)
+        else:
+            SP = np.squeeze(self.dataset.speed_t[time_value, :, :])
+            US = np.squeeze(self.dataset.ut_velocity[time_value, :, :] / SP)
+            VS = np.squeeze(self.dataset.vt_velocity[time_value, :, :] / SP)
+        mask = self.dataset.bottom_level != 0
+        p = np.ma.masked_where(mask == 0, SP)
+        u = np.ma.masked_where(mask[0::Np, 0::Np] == 0, US[0::Np, 0::Np])
+        v = np.ma.masked_where(mask[0::Np, 0::Np] == 0, VS[0::Np, 0::Np])
+        # x=nemo_t.dataset.longitude[0::Np,0::Np]
+        # y=nemo_t.dataset.latitude[0::Np,0::Np]
+        X = self.dataset.longitude
+        Y = self.dataset.latitude
+        # create a light colour map
+        N_colours = int(Vmax * 100)
+        n_c = 2
+        cmap0 = cm.get_cmap("BrBG_r", lut=N_colours + n_c * 2)
+        colors = cmap0(np.arange(cmap0.N))
+        colors1 = colors[n_c : cmap0.N - n_c]
+        cmap1 = LinearSegmentedColormap.from_list("cmap1", colors1, cmap0.N - n_c * 2)
+
+        cmap1.set_bad([0.75, 0.75, 0.75])
+
+        fig = plt.figure()
+        fig.clear()
+        plt.pcolormesh(p, cmap=cmap1)
+        # plt.pcolormesh(X,Y,p,cmap=cmap1)
+        x = np.arange(0, nx, Np)
+        y = np.arange(0, ny, Np)
+
+        plt.clim([0, Vmax])
+        plt.colorbar(orientation="vertical", cmap=cmap1)
+        # plt.quiver(x,y,u,v,color=[0.1,0.1,0.1],headwidth=4,scale=50)
+        plt.quiver(x, y, u, v, color=[0.1, 0.1, 0.1], headwidth=headwidth, scale=scale)  # ,**kwargs)
+        plt.title("Surface Currents " + name)
+        return fig
diff --git a/coast/diagnostics/gridded_stratification.py b/coast/diagnostics/gridded_stratification.py
index 5c670505..d41855da 100644
--- a/coast/diagnostics/gridded_stratification.py
+++ b/coast/diagnostics/gridded_stratification.py
@@ -338,5 +338,5 @@ def quick_plot(self, var: xr.DataArray = None):
             plt.ylabel("latitude")
             plt.clim([0, 50])
             plt.colorbar()
-            plt.show()
+            # plt.show()
         return fig, ax


Test patch:
diff --git a/unit_testing/test_diagnostic_methods.py b/unit_testing/test_diagnostic_methods.py
index f054d90f..b34db670 100644
--- a/unit_testing/test_diagnostic_methods.py
+++ b/unit_testing/test_diagnostic_methods.py
@@ -130,3 +130,43 @@ def test_construct_pycnocline_depth_and_thickness(self):
             fig.tight_layout()
             fig.savefig(files.dn_fig + "strat_1st_mom.png")
             plt.close("all")
+
+    def test_circulation(self):
+        # %%
+        lims = [150, 250, 100, 350]
+        nemo_u = coast.Gridded(
+            fn_data=files.fn_nemo_grid_u_dat, fn_domain=files.fn_nemo_dom, config=files.fn_config_u_grid, lims=lims
+        )
+        nemo_v = coast.Gridded(
+            fn_data=files.fn_nemo_grid_v_dat, fn_domain=files.fn_nemo_dom, config=files.fn_config_v_grid, lims=lims
+        )
+
+        with self.subTest("Map velocities to t-points"):
+            nemo_t = coast.CurrentsOnT(fn_domain=files.fn_nemo_dom, config=files.fn_config_t_grid, lims=lims)
+            nemo_t.currents_on_t(nemo_u, nemo_v)
+            nemo_t.subset(z_dim=[0], t_dim=[0])
+
+            u1 = 0.5 * (
+                nemo_u.dataset.u_velocity[0, 0, 150, 60].values + nemo_u.dataset.u_velocity[0, 0, 150, 59].values
+            )
+            u2 = nemo_t.dataset.ut_velocity[0, 0, 150, 60].values
+            # print(f"u vel on u-pts: {nemo_u.dataset.u_velocity.isel(x_dim=slice(59,61), y_dim=slice(149,151), t_dim=0, z_dim=0).values}")
+            # print(f"u vel on t-pts: {nemo_t.dataset.ut_velocity.isel(x_dim=slice(59,61), y_dim=slice(149,151), t_dim=0, z_dim=0).values}")
+            v1 = 0.5 * (
+                nemo_v.dataset.v_velocity[0, 0, 150, 60].values + nemo_v.dataset.v_velocity[0, 0, 149, 60].values
+            )
+            v2 = nemo_t.dataset.vt_velocity[0, 0, 150, 60].values
+            speed = nemo_t.dataset.speed_t[0, 0, 150, 60].values
+
+            check1 = np.isclose(u1, u2)
+            check2 = np.isclose(v1, v2)
+            check3 = np.isclose(speed, 0.04988122)
+            self.assertTrue(check1, msg="check 1")
+            self.assertTrue(check2, msg="check 2")
+            self.assertTrue(check3, msg="check 3")
+
+        with self.subTest("Plot circulation on t-points"):
+            fig = nemo_t.quick_plot("Test")
+            fig.tight_layout()
+            fig.savefig(files.dn_fig + "surface_circulation.png")
+            plt.close("all")
diff --git a/unit_testing/test_gridded_initialisation.py b/unit_testing/test_gridded_initialisation.py
index e7d27bca..0c4c57f8 100644
--- a/unit_testing/test_gridded_initialisation.py
+++ b/unit_testing/test_gridded_initialisation.py
@@ -150,3 +150,16 @@ def test_gridded_compute_e3_from_ssh(self):
         reference = np.array([8.337016e08, 8.333972e08, 8.344886e08, 8.330722e08, 8.265948e08])
         check1 = np.allclose(cksum, reference)
         self.assertTrue(check1, msg="check1")
+
+    def test_gridded_subset(self):
+        lims = [150, 250, 100, 350]
+        nemo_t_subset = coast.Gridded(
+            fn_data=files.fn_nemo_dat, fn_domain=files.fn_nemo_dom, config=files.fn_config_t_grid, lims=lims
+        )
+        nemo_t = coast.Gridded(fn_data=files.fn_nemo_dat, fn_domain=files.fn_nemo_dom, config=files.fn_config_t_grid)
+
+        T1 = nemo_t.dataset.temperature[0, 0, 105, 155].values
+        T2 = nemo_t_subset.dataset.temperature[0, 0, 5, 5].values
+        check1 = np.isclose(T1, T2)
+
+        self.assertTrue(check1, msg="check1")
diff --git a/unit_testing/unit_test_contents.txt b/unit_testing/unit_test_contents.txt
index 5d250880..4038fa3b 100755
--- a/unit_testing/unit_test_contents.txt
+++ b/unit_testing/unit_test_contents.txt
@@ -11,6 +11,7 @@
       h. gridded_load_of_data_only
       i. gridded_load_subregion_with_domain
       j. gridded_load_variables_correctly_renamed
+      k. gridded_subset
 
 2. test_general_utils
       a. bst_to_gmt
@@ -27,9 +28,10 @@
       a. basic_conversion_to_xesmf
 
 5. test_diagnostic_methods
-      a. compute_vertical_spatial_derivative
-      b. construct_density
-      c. construct_pycnocline_depth_and_thickness
+      a. circulation
+      b. compute_vertical_spatial_derivative
+      c. construct_density
+      d. construct_pycnocline_depth_and_thickness
 
 6. test_transect_methods
       a. calculate_transport_velocity_and_depth