starting to look at divergence at dual edges

pyicon/pyicon_calc_xr.py

+import warnings
 print('numpy')
 import numpy as np
 print('xarray')
@@ -418,18 +419,92 @@ def xr_edges2edges_via_cell(ds_IcD, vn_e, scalar, dze='const'):
 
 ## Mapping between cells and vertices
 
-def xr_calc_cell2vertex_coeff(ds_IcD):
-    raise(NotImplementedError)
-    dist_vector = ds_IcD.vert_cart_vec - ds_IcD.cell_cart_vec.isel(cell=ds_IcD.cells_of_vertex)
-    orientation = (dist_vector*ds_IcD.edge_prim_norm).sum(dim='cart')
-    dist_vector *= np.sign(orientation)
-    edge2cell_coeff_cc = (  ds_IcD.edge_prim_norm*ds_IcD.grid_sphere_radius
-                            * np.sqrt((dist_vector**2).sum(dim='cart'))
-                            / ds_IcD.dual_edge_length )
-    edge2cell_coeff_cc = edge2cell_coeff_cc.transpose('edge', 'nc', 'cart')
-    return edge2cell_coeff_cc
+def xr_calc_cell2vertex_coeff(ds_IcD, method=None):
+    """ Calculates cell2vertex coefficients for scalars
+
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+        A converted tgrid dataset
+
+    method : ["equal", "sector_area", "distance", "area_dist"]
+        Method used for weighting
+
+    Returns
+    -------
+    cell2vertex_coeff : xr.DataArray
+        Array of coefficients for mapping cells to vertices, with dimensions
+        ("vertex", "ne"). Note that technically "ne" should be replaced by "nc"
+
+    Notes
+    -----
+    Hexagonal barycentric interpolants are not uniquely defined. This function
+    offers several different options, and a user may even choose their own.
+    The sole constraint on the interpolants is that they sum to one; however,
+    it is advisable to use physical arguments when determining the
+    coefficients. For instance, we would expect points to be more influenced by
+    points nearer to them or representing a greater fraction of the hexagon our
+    point is in. If our hexagon has symmetry, we may want weights which reflect
+    these symmetries, be this inversion, rotation or reflection.
+    """
+    if method is None:
+        warnings.warn("No cell2vertex method specified, falling back on 'equal'")
+        method = "equal"
+    
+    cells_of_vertex = ds_IcD["cells_of_vertex"]
+    cell_cart_vec_of_vertex = ds_IcD["cell_cart_vec"].isel(cell=cells_of_vertex)
+    
+    if method == "equal":
+        cell2vertex_coeff = 1 / 6 * xr.ones_like(ds_IcD["cells_of_vertex"])
+    
+    elif method == "sector_area":
+        # This method is almost conservative.
+        vecAP = -cell_cart_vec_of_vertex + ds_IcD["vert_cart_vec"]
+        
+        vecAFp = 0.5 * (-cell_cart_vec_of_vertex + cell_cart_vec_of_vertex.roll(ne=-1))
+        vec_areaFpAP = xr.cross(vecAP, vecAFp, dim="cart")
+        areaFpAP = 0.5 * np.sqrt(xr.dot(vec_areaFpAP, vec_areaFpAP, dims="cart"))
+
+        vecABp = 0.5 * (-cell_cart_vec_of_vertex + cell_cart_vec_of_vertex.roll(ne=1))
+        vec_areaABpP = xr.cross(vecAP, vecABp, dim="cart")
+        areaABpP = 0.5 * np.sqrt(xr.dot(vec_areaABpP, vec_areaABpP, dims="cart"))
+
+        areaFpABpP = areaFpAP + areaABpP
+        cell2vertex_coeff = areaFpABpP / areaFpABpP.sum(dim="ne")
+    
+    elif method == "distance":
+        vecPA = (ds_IcD["vert_cart_vec"] - cell_cart_vec_of_vertex)
+        distPA = np.sqrt(xr.dot(vecPA, vecPA, dims="cart"))
+        cell2vertex_coeff = (1 / distPA) / (1 / distPA).sum(dim="ne")
+    
+    elif method == "area_dist":
+        # First do the sector area
+        vecAP = -cell_cart_vec_of_vertex + ds_IcD["vert_cart_vec"]
+        
+        vecAFp = 0.5 * (-cell_cart_vec_of_vertex + cell_cart_vec_of_vertex.roll(ne=-1))
+        vec_areaFpAP = xr.cross(vecAP, vecAFp, dim="cart")
+        areaFpAP = 0.5 * np.sqrt(xr.dot(vec_areaFpAP, vec_areaFpAP, dims="cart"))
+
+        vecABp = 0.5 * (-cell_cart_vec_of_vertex + cell_cart_vec_of_vertex.roll(ne=1))
+        vec_areaABpP = xr.cross(vecAP, vecABp, dim="cart")
+        areaABpP = 0.5 * np.sqrt(xr.dot(vec_areaABpP, vec_areaABpP, dims="cart"))
+
+        areaFpABpP = areaFpAP + areaABpP        
+    
+        # Now do the distance
+        vecPA = (ds_IcD["vert_cart_vec"] - cell_cart_vec_of_vertex)
+        distPA = np.sqrt(xr.dot(vecPA, vecPA, dims="cart"))
 
-def xr_cell2vertex(ds_IcD, da_b, cell2vertex_coeff_cc=None, fixed_vol_norm=None):
+        # Now calculate the coefficients
+        cell2vertex_coeff = (areaFpABpP / distPA)
+        cell2vertex_coeff = cell2vertex_coeff / cell2vertex_coeff.sum(dim="ne")
+
+    else:
+        raise ValueError("Selected method cannot be found")
+    
+    return cell2vertex_coeff
+
+def xr_cell2vertex(ds_IcD, da_b, cell2vertex_coeff=None, fixed_vol_norm=None):
     """ Interpolate cell variables on to vertices
     Parameters
     ----------
@@ -439,29 +514,22 @@ def xr_cell2vertex(ds_IcD, da_b, cell2vertex_coeff_cc=None, fixed_vol_norm=None)
     da_b : xr.DataArray
         The data to be mapped onto cell vertices. This should be a scalar.
 
+    cell2vertex_coeff : xr.DataArray or string
+        This gives the weightings to be used during interpolation. If an array,
+        it should have the dimensions ("vertex", "ne") <<NEED TO FIX SO NE IS NC>>
+        If a string it should take a valid method argument for xr_calc_cell2vertex_coeff
+        
     Returns
     -------
-    ds_icd : xr.Dataset
-        A tgrid dataset compatible with pyicon functions
+    da_b_vert : xr.DataArray
+        The array da_b interpolated onto vertices
 
     Notes
     -----
-    This function uses hexagonal barycentric interpolation. Each vertex is
-    (normally) surrounded by 6 cell centres. The value at the vertex should
-    depend only upon these vertices. We express the position of the cell centre
-    as a sum of contributions from each point of the hexagon. We then use these
-    weights to interpolate. Hexagonal barycentric interpolant coefficients
-    aren't uniquely defined. Here, for the hexagon ABCDEF and point P, we
-    choose coefficients
-        x_A = [ABC]/[PBC]
-        y_B = [BCD]/[PCD]
-        z_C = [CDE]/[PDE]
-        u_D = [DEF]/[PEF]
-        v_E = [EFA]/[PFA]
-        w_F = [FAB]/[PAB]
-    where [xyz] represents the area of triangle xyz. We can then interpolate
-    a function f onto the point P using
-    f(P) = x_A f(A) + y_B f(B) + z_C f(C) + u_D f(D) + v_E f(E) + w_F f(F)
+    This function is useful for calculating potential vorticity using it's
+    C-grid divergence form (see for instance Morel et al. 2019: 
+    https://doi.org/10.1016/j.ocemod.2019.04.004). In this form we require
+    cell centre scalars (such as density) on vertices (vorticity points).
 
     What isn't immediately clear to me is what happens at edge points where
     there may be less valid centres surrounding the vertex? Perhaps we need to
@@ -469,23 +537,12 @@ def xr_cell2vertex(ds_IcD, da_b, cell2vertex_coeff_cc=None, fixed_vol_norm=None)
 
     Furthermore, do we need to apply any volume normalisation? 
     """
-    raise(NotImplementedError)
-    if cell2vertex_coeff_cc is None:
-        cell2vertex_coeff_cc = xr_calc_cell2vertex_coeff(ds_IcD)
-    ic0 = ds_IcD.adjacent_cell_of_edge.isel(nc=0).data
-    ic1 = ds_IcD.adjacent_cell_of_edge.isel(nc=1).data
-    ptp_vn = (
-        (
-            p_vn_c.isel(cell=ic0).rename({'cell': 'edge'})#.chunk(dict(edge=ic0.size))
-            * edge2cell_coeff_cc_t.isel(nc=0)
-        ).sum(dim='cart')                                            
-        +
-        (
-            p_vn_c.isel(cell=ic1).rename({'cell': 'edge'})#.chunk(dict(edge=ic0.size))
-            * edge2cell_coeff_cc_t.isel(nc=1)
-        ).sum(dim='cart') 
-    )
-    return ptp_vn
+    if type(cell2vertex_coeff) != xr.DataArray:
+        # Calculate the coefficients
+        cell2vertex_coeff = xr_calc_cell2vertex_coeff(ds_IcD, method=cell2vertex_coeff)
+    
+    da_b_vert = (da_b * cell2vertex_coeff).sum(dim="ne")
+    return da_b_vert
 
 ## Divergence
 
@@ -504,6 +561,38 @@ def xr_calc_div(ds_IcD, vector, div_coeff=None):
     ).sum(dim='nv')
     return div_of_vector
 
+
+def xr_calc_div_dual_coeff():
+    pass
+
+
+def xr_calc_div_dual(ds_IcD, vector, div_dual_coeff=None):
+    """ Calculates the divergence of vectors on dual edges
+
+    Parameters
+    ----------
+    ds_IcD : xr.Dataset
+
+    vector : xr.DataArray
+        A vector defined on dual edge points (i.e. pointing towards vertices)
+
+    div_dual_coeff : None or xr.DataArray
+
+    Returns
+    -------
+    div_of_vector : xr.DataArray
+        The divergence of the dual-edge vector, defined on vertices
+    """
+    if div_dual_coeff is None:
+        div_dual_coeff = xr.calc_div_dual_coeff()
+    
+    div_of_vector = (
+        vector.isel(edge=ds_IcD["edges_of_vertex"])
+        * div_dual_coeff
+        ).sum(dim="nv")  # Strictly speaking this should be ne...
+    
+    return div_of_vector
+
 ## Gradient
 
 def xr_calc_grad_coeff(ds_IcD):

pyicon/tests/test_pyicon_calc_xr.py

@@ -95,3 +95,12 @@ def test_nabla_funcs(processed_tgrid):
     assert np.allclose(curl_of_grad, 0)
 
     # Should include other tests in the future if any refactoring is done
+
+
+@pytest.mark.parametrize("method", [None, "equal", "sector_area", "distance", "area_dist"])
+def test_cell2vertex(processed_tgrid, method):
+    # Check the cell2vertex coefficients sum to one or nan.
+    weights = pyic.xr_calc_cell2vertex_coeff(processed_tgrid, method=method)
+    sums_to_1 = np.isclose(weights.sum(dim="ne", skipna=False), 1)
+    sums_to_nan = np.isnan(weights.sum(dim="ne", skipna=False))
+    assert np.all(sums_to_nan | sums_to_1)
\ No newline at end of file