How can I use dimensions and data in xarray to create a new array of data?

Question

I'm running into issues with DataSets in Xarray . I'd like to apply a function to data and coordinates already in the dataset to compute a new set of data (essentially, w = f ( x, y, z ), where x , y , and z are floats).

There are a series of functions that need to be applied, and the dimensions of the resulting array need to expand when a new variable is pulled in. So f(x) should be a 1D array, g(f(x), y) should be 2D, and h(g(f(x), y), z) should be a 3D array created from the previous 2D array. I'm able to use np.outer to force the 1D -> 2D transformation, but I'm having trouble with the 2D -> 3D conversion. Given the amount of trouble I'm having, I suspect I'm coming at the problem from the wrong direction.

Below is a snippet of some of the code. Can someone help me understand the correct way to approach this problem from a NumPy/Xarray standpoint (I'm really trying avoid dropping to a loop for this work…)

import numpy as np
import xarray as xr

data = xr.Dataset(
    data_vars={
        "abs": (("x"), np.ones((200)))
    },
    coords={
        "x": np.linspace(1.5,1.55,200),
        "y": np.arange(5e5),
        "z": np.arange(125)
    }
)

def intensity_decay(da):
    "Return decay curve (exponential decay)."
    init_power = 1e5
    decay = init_power * np.exp(
        -1 * np.outer(da.abs, da.y * 10 ** (-7))
        )
    decay[decay < 1e-3] = 0
    return decay

data["depth_decay"] = (["x", "y"], intensity_decay(data))

def radial_decay(da):
    "Return radial decay curve."
    return np.outer(da.depth_decay, np.exp(-da.z))

radial_decay(data.isel(x=[1, 4, 10])).shape

As you can see, the radial_decay function doesn't broadcast into the z dimension. It seems like Xarray should support this type of operation, but I don't know how to approach the problem, or even where to start in the documentation. I currently have two datasets to which I'd like to apply these function, so I could do it by hand, but I want to make a framework I can use on future data as well.

Answer 1

After playing around a bit and using the function structure found in this answer , I was able to come up with an output of the right shape.

By defining the radial_decay function as follows, I can get an output array with three dimensions.

def radial_decay(da):
    "Return radial decay curve."
    da["radial_decay"] = (["x", "y", "z"], da.depth_decay * np.exp(-da.z))
    return da

With this structure, the output of radial_decay(data.isel(x=[1, 4, 10])).radial_decay.shape is (3, 500000, 125) , rather than (1500000,125) using the function defined in the original question.

While this seems to work (I still need to verify that the values are correct), I'd like to know if there are any other (better?) methods to achieve this goal, particularly if there are better ways with just NumPy to do this (in case I don't have access to Xarray for some reason in the future).