pydata / pydata/xarray

optimize align for scalars at least

Open
#8,350 5 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

What happened?

Here's a simple rescaling calculation:

import numpy as np
import xarray as xr

ds = xr.Dataset(
    {"a": (("x", "y"), np.ones((300, 400))), "b": (("x", "y"), np.ones((300, 400)))}
)
mean = ds.mean() # scalar
std = ds.std() # scalar
rescaled = (ds - mean) / std

The profile for the last line shows 30% (!!!) time spent in align (really reindex_like) except there's nothing to reindex when only scalars are involved!

image

This is a small example inspired by a ML pipeline where this normalization is happening very many times in a tight loop.

cc @benbovy

What did you expect to happen?

A fast path for when no reindexing needs to happen.

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 with the reported (ds - mean) / std calculation and inspect the align/reindex_like path used when the operands are scalar. Profile the operation to confirm the overhead, then verify that scalar-only operations avoid unnecessary reindexing while preserving the calculation's results and behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
data, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.