added a hacky cell2edges function for scalars

pyicon/pyicon_sparse.py

@@ -132,7 +132,7 @@ def sp_convert_tgrid_data(ds_tgrid, matrix_dtype="float32"):
     Parameters
     ----------
     ds_tgrid : xr.Dataset
-        raw, unprocessed tgrid
+        raw, unprocessed tgrid *or* converted tgrid (ds_IcD) dataset
 
     matrix_dtype : float type
         the data type of the sparse matrices. Default is "float32". Should
@@ -154,7 +154,10 @@ def sp_convert_tgrid_data(ds_tgrid, matrix_dtype="float32"):
     Then convert by:
     >>> ds_SpD = pyic.sp_convert_tgrid_data(ds_tgrid)
     """
-    ds_SpD = convert_tgrid_data(ds_tgrid)
+    if not hasattr(ds_tgrid, "converted_tgrid"):
+        ds_SpD = convert_tgrid_data(ds_tgrid)
+    else:
+        ds_SpD = ds_tgrid.copy(deep=True)
 
     div_matrix = _construct_div_matrix(ds_SpD, dtype=matrix_dtype)
     ds_SpD["div_matrix"] = xr.DataArray(div_matrix, dims=["cell", "edge"])
@@ -374,3 +377,103 @@ def _construct_grad_matrix(ds_IcD, dtype="float32"):
         shape=(ds_IcD.sizes["edge"], ds_IcD.sizes["cell"]),
     )
     return grad_matrix
+
+
+def _construct_scalar_cell2edges_matrix(ds_IcD, dtype="float32"):
+    rows = []
+    cols = []
+    data = []
+    cell_separation = ds_IcD["edge_cell_distance"].sum("nc_e")
+    for nc_e in range(2):
+        rows += [ds_IcD["edge"].values]
+        cols += [ds_IcD["adjacent_cell_of_edge"].isel(nc_e=nc_e).values]
+        data += [ds_IcD["edge_cell_distance"].isel(nc_e=nc_e) / cell_separation]
+        
+    row = np.hstack(rows)
+    col = np.hstack(cols)
+    data = np.hstack(data)
+
+    col[col == -1] = ds_IcD.sizes["cell"] - 1
+
+    coordinates = np.vstack([row, col])
+    
+    scalar_cell2edges_matrix = sparse.COO(
+        coordinates,
+        data.astype(dtype),
+        shape=(ds_IcD.sizes["edge"], ds_IcD.sizes["cell"]),
+    )
+    return scalar_cell2edges_matrix
+
+
+def _construct_scalar_cell2edges_matrix(ds_IcD, dtype="float32"):
+    rows = []
+    cols = []
+    data = []
+    cell_separation = ds_IcD["edge_cell_distance"].sum("nc_e")
+    for nc_e in range(2):
+        rows += [ds_IcD["edge"].values]
+        cols += [ds_IcD["adjacent_cell_of_edge"].isel(nc_e=nc_e).values]
+        data += [1 - ds_IcD["edge_cell_distance"].isel(nc_e=nc_e) / cell_separation]
+        
+    row = np.hstack(rows)
+    col = np.hstack(cols)
+    data = np.hstack(data)
+
+    col[col == -1] = ds_IcD.sizes["cell"] - 1
+
+    coordinates = np.vstack([row, col])
+    
+    scalar_cell2edges_matrix = sparse.COO(
+        coordinates,
+        data.astype(dtype),
+        shape=(ds_IcD.sizes["edge"], ds_IcD.sizes["cell"]),
+    )
+    return scalar_cell2edges_matrix
+
+
+def sp_scalar_cell2edges(ds_SpD, scalar):
+    """Interpolates a scalar from a cell center to an edge
+
+    Parameters
+    ----------
+    ds_SpD : xr.Dataset
+        pyicon_sparse dataset containing coordinate info
+
+    scalar : xr.DataArray
+        scalar at cell center
+
+
+    Returns
+    -------
+    scalar_on_edges : xr.DataArray
+        the scalar interpolated to the edges of the grid
+        
+    Notes
+    -----
+    The function performs an distance weighted interpolation from the center of
+    a cell to the edges. The mapping isn't strictly defined in the ICON model;
+    however, it is often useful.
+    """
+    scalar = _check_chunk(scalar, "cell")
+    
+    if "scalar_cell2edges_matrix" not in ds_SpD.data_vars:
+        scalar_cell2edges_matrix = _construct_scalar_cell2edges_matrix(
+            ds_SpD
+        )
+        ds_SpD["scalar_cell2edges_matrix"] = xr.DataArray(scalar_cell2edges_matrix, dims=["edge", "cell"])
+        
+    
+    scalar_on_edges = xr.apply_ufunc(
+        _multiply_sparse,
+        scalar,
+        kwargs={
+            "scipy_sparse_matrix": ds_SpD["scalar_cell2edges_matrix"].data.to_scipy_sparse(),
+        },
+        dask="parallelized",
+        input_core_dims=[["cell"]],
+        output_core_dims=[["edge"]],
+        dask_gufunc_kwargs={"output_sizes": {"edge": ds_SpD.sizes["edge"]}},
+        vectorize=True,
+        output_dtypes=[scalar.dtype],
+    )
+    return scalar_on_edges