pydata / pydata/xarray

add join argument to xr.broadcast?

Open
#6,304 1 comment 4 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
4.2k
Forks
1.4k
Avg merge
2d 15h
Merged PRs (30d)
14

Description

Is your feature request related to a problem?

xr.broadcast always does an outer join:

https://github.com/pydata/xarray/blob/de965f342e1c9c5de92ab135fbc4062e21e72453/xarray/core/alignment.py#L702

https://github.com/pydata/xarray/blob/de965f342e1c9c5de92ab135fbc4062e21e72453/xarray/core/alignment.py#L768

This is not how the (default) broadcasting (arithmetic join) works, e.g. the following first does an inner join and then broadcasts:

import xarray as xr

da1 = xr.DataArray([[0, 1, 2]], dims=("y", "x"), coords={"x": [0, 1, 2]})
da2 = xr.DataArray([0, 1, 2, 3, 4], dims="x", coords={"x": [0, 1, 2, 3, 4]})
da1 + da2
<xarray.DataArray (y: 1, x: 3)>
array([[0, 2, 4]])
Coordinates:
  * x        (x) int64 0 1 2
Dimensions without coordinates: y
Describe the solution you'd like

Add a join argument to xr.broadcast. I would propose to leave the default as is

def broadcast(*args, exclude=None, join="outer"):
    args = align(*args, join=join, copy=False, exclude=exclude) 
Describe alternatives you've considered
  • We could make broadcast respect options -> arithmetic_join but that would be a breaking change and I am not sure how the deprecation should/ would be handled...
  • We could leave it as is.
Additional context
  • xr.broadcast should not be used often because this is should happen automatically in most cases
  • in #6059 I use broadcast because I couldn't get it to work otherwise (maybe there is a better way?). However, the "outer elements" are immediately discarded again - so it's kind of pointless to do an outer join.
import numpy as np
import xarray as xr

da = xr.DataArray(np.arange(6).reshape(3, 2), coords={"dim_0": [0, 1, 2]})
w = xr.DataArray([1, 1, 1, 1, 1, 1], coords={"dim_0": [0, 1, 2, 4, 5, 6]})
da.weighted(w).quantile(0.5)

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in xarray/core/alignment.py at the linked xr.broadcast implementation and review how it calls align, including the existing exclude handling. Check the requested join behavior against the arithmetic example and the weighted quantile context. Done means xr.broadcast accepts the proposed join argument while preserving the current outer-join default.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.