Correct broadcasting of 4D array in apply_ufunc

[trondkr]

Hi all,

I would like to run the calculation of the squared buoyancy at grid points (lat,lon) with time-varying properties at depth. This. gives me 4D variables as input to my function (time, depth, lat, lon). I use the gsw-xarray wrapper package for ocean package gsw from where I use the function Nsquared. The Nsquared function is calculated as a function of depth and the wrapper package enables DataArray as input. I am having some issues getting this to work when following the example in xarray and I am hoping you could help me. Below is a standalone pytest that currently fails with error (requires the gsw-xarray package to run):

 def _broadcast_shape(*args):
        """Returns the shape of the arrays that would result from broadcasting the
        supplied arrays against each other.
        """
        # use the old-iterator because np.nditer does not handle size 0 arrays
        # consistently
>       b = np.broadcast(*args[:32])
E       ValueError: shape mismatch: objects cannot be broadcast to a single shape.  Mismatch is between arg 0 with shape (40, 5, 5, 2) and arg 2 with shape (40,).

The test:

import pytest
import gsw_xarray as gsw
import xarray as xr
import numpy as np

"""
Note test requires use of gsw-xarray package to run Nsquared
https://github.com/DocOtak/gsw-xarray
"""


@pytest.fixture(scope="function")
def setup(request):
    request.cls.lons = np.arange(0, 2, 1)
    request.cls.lats = np.arange(0, 5, 1)
    request.cls.depths = np.arange(0, 40, 1)
    request.cls.times = np.arange(0, 5, 1)


@pytest.mark.usefixtures("setup")
class Test_gsw:
    def get_4d_coords(self):
        return {
            "time": (["time"], self.times),
            "depth": (["depth"], self.depths),
            "lat": (["y"], self.lats),
            "lon": (["x"], self.lons),
        }

    def create_4d_src_dataset(
            self, source_4d_depth_levels, name="temperature"
    ) -> xr.Dataset:
        return xr.DataArray(
            name=name,
            data=source_4d_depth_levels,
            coords=self.get_4d_coords(),
            dims=["time", "depth", "y", "x"],
        ).to_dataset()

    def test_N2(self):
        """Test to use xarray to calculate the square of the buoyancy frequency.

        Returns: None

        """
        CT = np.ones(
            (len(self.times), len(self.depths), len(self.lats), len(self.lons))
        )
        SA = 35 * np.ones(
            (len(self.times), len(self.depths), len(self.lats), len(self.lons))
        )

        ds_temp = self.create_4d_src_dataset(CT, name="temp")
        ds_sa = self.create_4d_src_dataset(SA, name="so")
        ds_temp = ds_temp.transpose("depth", "time", "y", "x")
        ds_sa = ds_sa.transpose("depth", "time", "y", "x")
        p = ds_temp["depth"]

        N2, pmid = xr.apply_ufunc(gsw.Nsquared,
                                  ds_sa["so"],
                                  ds_temp["temp"],
                                  p,
                                  input_core_dims=[["depth", "time", "y", "x"], ["depth", "time", "y", "x"], ["depth"]],
                                  output_core_dims=[["depth_n", "time", "y_n", "x_n"],
                                                    ["depth_n", "time", "y_n", "x_n"]],
                                  exclude_dims=set(("depth", "time", "y", "x",)),
                                  vectorize=True)
        print(N2)

        assert isinstance(N2, xr.Dataset)
        assert N2 is not None

The header of the Nsquaredfunction in gswlooks like this:

def Nsquared(SA, CT, p, lat=None, axis=0):
    """
    Calculate the square of the buoyancy frequency.
    Parameters
    ----------
    SA : array-like
        Absolute Salinity, g/kg
    CT : array-like
        Conservative Temperature (ITS-90), degrees C
    p : array-like
        Sea pressure (absolute pressure minus 10.1325 dbar), dbar
    lat : array-like, 1-D, optional
        Latitude, degrees.
    axis : int, optional
        The dimension along which pressure increases.
    Returns
    -------
    N2 : array
        Buoyancy frequency-squared at pressure midpoints, 1/s^2.
        The shape along the pressure axis dimension is one
        less than that of the inputs.
        (Frequency N is in radians per second.)
    p_mid : array
        Pressure at midpoints of p, dbar.
        The array shape matches N2.
    """

I appreciate any help in resolving this issue. Thanks. Cheers, Trond

[trondkr]

I got this to work by modifying my code slightly and properly broadcasting depth and storing it as a new variable. Perhaps not the most efficient use of xarray but it does work. Suggestions for improvements to efficiency and approach are most appreciated.

import os
import pytest
import gsw_xarray as gsw
import xarray as xr
import numpy as np
import pandas as pd

"""
Note test requires use of gsw-xarray package to run Nsquared
https://github.com/DocOtak/gsw-xarray
"""


@pytest.fixture(scope="function")
def setup(request):
    request.cls.lons = np.arange(0, 2, 1)
    request.cls.lats = np.arange(0, 5, 1)
    request.cls.depths = np.arange(0, 40, 1)
    request.cls.times = np.arange(np.datetime64("2000-01-01"), np.datetime64("2010-01-01"), np.timedelta64(1, "7D"))


@pytest.mark.usefixtures("setup")
class Test_gsw:
    def get_4d_coords(self):
        return {"depth": (["depth"], self.depths),
                "time": (["time"], self.times),
                "lat": (["y"], self.lats),
                "lon": (["x"], self.lons),
                }

    def create_4d_src_dataset(
            self, source_4d_depth_levels, pmid=None, name="temperature"
    ) -> xr.Dataset:
        if pmid is not None:
            self.depths = pmid

        return xr.DataArray(
            name=name,
            data=source_4d_depth_levels,
            coords=self.get_4d_coords(),
            dims=["depth", "time", "y", "x"],
        ).to_dataset()

    def test_N2(self):
        """Test to use xarray to calculate the square of the buoyancy frequency.

        Returns: None

        """
        CT = 10 * np.random.rand(len(self.depths), len(self.times), len(self.lats), len(self.lons))
        SA = 35 * np.random.rand(len(self.depths), len(self.times), len(self.lats), len(self.lons))

        ds_temp = self.create_4d_src_dataset(CT, name="temp")
        ds_sa = self.create_4d_src_dataset(SA, name="so")
        ds_sa = ds_sa.assign(depth_n=ds_sa["depth"])

        (ds1, ds2) = xr.broadcast(ds_sa, ds_temp)

        N2, pmid = xr.apply_ufunc(gsw.Nsquared,
                                  ds1["so"].values,
                                  ds2["temp"].values,
                                  ds1["depth_n"].values,
                                  input_core_dims=[["depth", "time", "y", "x"],
                                                   ["depth", "time", "y", "x"],
                                                   ["depth"]],
                                  output_core_dims=[["depth", "time", "y", "x"],
                                                    ["depth", "time", "y", "x"]],
                                  exclude_dims=set(("depth", "time", "y", "x",)),
                                  vectorize=True)

        ds_n2 = self.create_4d_src_dataset(N2, pmid=np.squeeze(pmid[:, 0, 0, 0]), name="n2")