Support basic xarray-functionality

.github/workflows/ReceivePR.yml

@@ -27,6 +27,7 @@ jobs:
       - name: Test
         run: |
           pip install .[dev]
+          pip install xarray
           pre-commit run --all-files
           python -m unittest
 

.github/workflows/benchmark_main.yml

@@ -0,0 +1,46 @@
+name: Benchmark main and save
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  benchmark-main:
+    name: Benchmark main and save
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+      - name: Setup MPI
+        uses: mpi4py/setup-mpi@v1
+      - name: Use Python 3.10
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.10.11 # Perun only supports 3.8 and ahead
+          architecture: x64
+      - name: Test
+        run: |
+            pip install torch==1.12.1+cpu torchvision==0.13.1+cpu torchaudio==0.12.1 -f https://download.pytorch.org/whl/torch_stable.html
+            pip install xarray
+            pip install .[cb]
+            PERUN_RUN_ID=N4 mpirun -n 4 python benchmarks/cb/main.py
+            jq -s flatten bench_data/*.json > bench_data/all_benchmarks.json
+      - name: Save benchmark result and update gh-pages-chart
+        if: ${{github.ref == 'refs/heads/main'}}
+        uses: benchmark-action/github-action-benchmark@v1
+        with:
+          github-token: ${{secrets.GITHUB_TOKEN}}
+          # Benchmark action input and output
+          tool: 'customSmallerIsBetter'
+          output-file-path: bench_data/all_benchmarks.json
+          # external-data-json-path: ./cache/benchmark-data.json
+          # Alert configuration
+          fail-on-alert: false  # Don't fail on main branch
+          comment-on-alert: true
+          # Save benchmarks from the main branch
+          save-data-file: true
+          # Pages configuration
+          auto-push: true
+          gh-pages-branch: gh-pages
+          benchmark-data-dir-path: dev/bench
+      # Upload the updated cache file for the next job by actions/cache

.github/workflows/benchmark_pr.yml

@@ -0,0 +1,46 @@
+name: Benchmark PR
+on:
+  pull_request:
+    types: [opened, synchronize, reopened, labeled]
+    branches: [main]
+
+jobs:
+  benchmark-pr:
+    name: Benchmark PR
+    if: contains(github.event.pull_request.labels.*.name, 'benchmark PR')
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+      - name: Setup MPI
+        uses: mpi4py/setup-mpi@v1
+      - name: Use Python 3.10
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.10.11 # Perun only supports 3.8 and ahead
+          architecture: x64
+      - name: Test
+        run: |
+            pip install torch==1.12.1+cpu torchvision==0.13.1+cpu torchaudio==0.12.1 -f https://download.pytorch.org/whl/torch_stable.html
+            pip install xarray
+            pip install .[cb]
+            PERUN_RUN_ID=N4 mpirun -n 4 python benchmarks/cb/main.py
+            jq -s flatten bench_data/*.json > bench_data/all_benchmarks.json
+      - name: Compare benchmark result
+        if: ${{github.ref != 'refs/heads/main'}}
+        uses: benchmark-action/github-action-benchmark@v1
+        with:
+          github-token: ${{secrets.GITHUB_TOKEN}}
+          # Benchmark action input and output
+          tool: 'customSmallerIsBetter'
+          output-file-path: bench_data/all_benchmarks.json
+          # external-data-json-path: ./cache/benchmark-data.json
+          # Alert configuration
+          fail-on-alert: true
+          comment-on-alert: true
+          # Ignore results from non main branches.
+          save-data-file: false
+          # Pages configuration
+          auto-push: false
+          gh-pages-branch: gh-pages
+          benchmark-data-dir-path: dev/bench

.github/workflows/ci.yaml

@@ -52,5 +52,6 @@ jobs:
           pip install pytest
           pip install ${{ matrix.pytorch-version }} --extra-index-url https://download.pytorch.org/whl/cpu
           pip install ${{ matrix.install-options }}
+          pip install xarray
           mpirun -n 3 pytest heat/
           mpirun -n 4 pytest heat/

heat/core/manipulations.py

@@ -9,7 +9,7 @@
 
 from typing import Iterable, Type, List, Callable, Union, Tuple, Sequence, Optional
 
-from .communication import MPI
+from .communication import MPI, sanitize_comm, Communication
 from .dndarray import DNDarray
 
 from . import arithmetics
@@ -21,6 +21,7 @@
 from . import tiling
 from . import types
 from . import _operations
+from . import devices
 
 __all__ = [
     "balance",
@@ -37,6 +38,7 @@
     "flip",
     "fliplr",
     "flipud",
+    "from_numpy",
     "hsplit",
     "hstack",
     "moveaxis",
@@ -1141,6 +1143,34 @@ def flipud(a: DNDarray) -> DNDarray:
     return flip(a, 0)
 
 
+def from_numpy(
+    x: np.ndarray,
+    split: Optional[int] = None,
+    device: Optional[Union[str, devices.Device]] = None,
+    comm: Optional[Communication] = None,
+) -> DNDarray:
+    """
+    Creates DNDarray from given NumPy Array. The data type is determined by the data type of the Numpy Array.
+    Split-dimension, device and communicator can be prescribed as usual.
+    Inverse of :meth:`DNDarray.numpy()`.
+    """
+    dtype = types.canonical_heat_type(x.dtype)
+    device = devices.sanitize_device(device)
+    comm = sanitize_comm(comm)
+    xht = DNDarray(
+        torch.from_numpy(x).to(device.torch_device),
+        x.shape,
+        dtype=dtype,
+        split=None,
+        device=device,
+        comm=comm,
+        balanced=True,
+    )
+    if split is not None:
+        xht.resplit_(split)
+    return xht
+
+
 def hsplit(x: DNDarray, indices_or_sections: Iterable) -> List[DNDarray, ...]:
     """
     Split array into multiple sub-DNDarrays along the 2nd axis (horizontally/column-wise).

heat/dxarray/__init__.py

@@ -0,0 +1,7 @@
+"""
+import into heat.dxarray namespace
+"""
+
+from .dxarray import *
+from .dxarray_sanitation import *
+from .dxarray_manipulations import *