Squashed commit of the following:

dcherian · dcherian · commit 5be5a35e88f9 · 2022-11-26T20:13:11.000-07:00
commit 51fb6e9 Author: Deepak Cherian <dcherian@users.noreply.github.com> Date: Tue Nov 15 13:16:11 2022 -0700 Some typing (#190) * Some typing * fixes. commit 9b01c48 Author: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Sat Nov 5 09:02:07 2022 +0100 Add windows CI (#151) * Add windows CI * Update ci.yaml * Update ci.yaml * Make arg input the same as shown in pytest * Add dtype check * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * have expected and actual results on the same side * use np.intp for count expected * [revert] minimize test * specify dtypes * more fixers * more. * Fix groupby_reduce * [revert] only wiindows tests * more fixes? * more fixes. * more fix * Last fix? * Update .github/workflows/ci.yaml * revert * Better fix * Revert "revert" This reverts commit 3b79f6e. * better comment. * clean up test * Revert "Revert "revert"" This reverts commit 38438a2. * xfail labels dtype test * Revert "[revert] only wiindows tests" This reverts commit 232cf15. * Revert "[revert] minimize test" This reverts commit f993b31. * fix bad revert Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: dcherian <deepak@cherian.net> Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com> commit e3ea0e7 Author: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue Nov 1 08:26:30 2022 -0600 Bump mamba-org/provision-with-micromamba from 13 to 14 (#184) Bumps [mamba-org/provision-with-micromamba](https://github.com/mamba-org/provision-with-micromamba) from 13 to 14. - [Release notes](https://github.com/mamba-org/provision-with-micromamba/releases) - [Commits](mamba-org/provision-with-micromamba@v13...v14) --- updated-dependencies: - dependency-name: mamba-org/provision-with-micromamba dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] <support@github.com> Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> commit c440148 Author: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue Nov 1 08:26:17 2022 -0600 Bump xarray-contrib/ci-trigger from 1.1 to 1.2 (#183) Bumps [xarray-contrib/ci-trigger](https://github.com/xarray-contrib/ci-trigger) from 1.1 to 1.2. - [Release notes](https://github.com/xarray-contrib/ci-trigger/releases) - [Commits](xarray-contrib/ci-trigger@v1.1...v1.2) --- updated-dependencies: - dependency-name: xarray-contrib/ci-trigger dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> commit 471aa94 Author: Deepak Cherian <dcherian@users.noreply.github.com> Date: Wed Oct 26 19:39:44 2022 -0600 Use blockwise to extract final result (#182) * Use blockwise to extract final result for method="blockwise" * FOr all methods * bugfix * Try return_array from _finalize_results * Revert "Try return_array from _finalize_results" This reverts commit cb25e38. * Fixes. commit df0da40 Author: Deepak Cherian <dcherian@users.noreply.github.com> Date: Tue Oct 25 16:09:04 2022 -0600 Fix blockwise sort optimization (#181) commit c2c4e1d Author: Deepak Cherian <dcherian@users.noreply.github.com> Date: Tue Oct 25 13:41:40 2022 -0600 Support reindexing in simple_combine (#177) * Support reindexing in simple_combine For 1D combine, great improvement for cohorts-type reductions More memory but similar time for map-reduce. Note that the map-reduce intermediates are a worst case where there are no shared groups between the chunks being combined. This case is actually optimized in _group_combine where reindexing is skipped for reducing along a single axis. [ 68.75%] ··· =========== ========= ========= -- combine ----------- ------------------- kind grouped combine =========== ========= ========= cohorts 760M 631M mapreduce 981M 1.81G =========== ========= ========= [ 75.00%] ··· =========== ========== =========== -- combine ----------- ---------------------- kind grouped combine =========== ========== =========== cohorts 393±10ms 137±10ms mapreduce 652±10ms 611±400ms =========== ========== =========== Fix bug in unique * Fix bug with all NaN blocks commit 0db264a Author: Deepak Cherian <dcherian@users.noreply.github.com> Date: Mon Oct 24 12:27:12 2022 -0600 Update visualize for main changes (#179) * Update visualize for main changes * Update tests/__init__.py commit ccf578d Author: Deepak Cherian <dcherian@users.noreply.github.com> Date: Mon Oct 24 09:26:07 2022 -0600 Use npg sum_of_squares (#135) Bump to npg >= 0.9.19 Closes #107 commit c370b5d Author: dcherian <deepak@cherian.net> Date: Thu Oct 20 10:07:36 2022 -0600 Fix unique commit 336e2bb Author: dcherian <deepak@cherian.net> Date: Wed Oct 19 14:32:12 2022 -0600 Remove split-reduce to reduce some test time. commit 5431813 Author: dcherian <deepak@cherian.net> Date: Wed Oct 19 20:28:15 2022 -0600 Type more utility functions commit aa14656 Author: dcherian <deepak@cherian.net> Date: Wed Oct 19 20:05:56 2022 -0600 Simplify code for blockwise commit b1b43d9 Author: dcherian <deepak@cherian.net> Date: Wed Oct 19 20:01:47 2022 -0600 Test argmax early commit 47e0b38 Author: Deepak Cherian <dcherian@users.noreply.github.com> Date: Wed Oct 19 17:46:54 2022 -0600 Force reindex to be bool always (#176) * Force reindex to be bool always Closes #155 Turns out we weren't using the more efficient simple_combine with map_reduce in all cases because do_simple_combine was None when reindex was None. Now the default for map-reduce is reindex=True when (expected_groups is not None) or (expected_groups is None and by_is_dask is False)
diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
@@ -22,7 +22,7 @@ jobs:
           fetch-depth: 0
 
       - name: Set up conda environment
-        uses: mamba-org/provision-with-micromamba@v13
+        uses: mamba-org/provision-with-micromamba@v14
         with:
           environment-file: ci/environment.yml
           environment-name: flox-tests
diff --git a/.github/workflows/ci-additional.yaml b/.github/workflows/ci-additional.yaml
@@ -25,7 +25,7 @@ jobs:
       - uses: actions/checkout@v3
         with:
           fetch-depth: 2
-      - uses: xarray-contrib/ci-trigger@v1.1
+      - uses: xarray-contrib/ci-trigger@v1.2
         id: detect-trigger
         with:
           keyword: "[skip-ci]"
@@ -53,7 +53,7 @@ jobs:
           echo "TODAY=$(date  +'%Y-%m-%d')" >> $GITHUB_ENV
 
       - name: Setup micromamba
-        uses: mamba-org/provision-with-micromamba@34071ca7df4983ccd272ed0d3625818b27b70dcc
+        uses: mamba-org/provision-with-micromamba@v14
         with:
           environment-file: ${{env.CONDA_ENV_FILE}}
           environment-name: flox-tests
@@ -94,7 +94,7 @@ jobs:
         run: |
           echo "TODAY=$(date +'%Y-%m-%d')" >> $GITHUB_ENV
       - name: Setup micromamba
-        uses: mamba-org/provision-with-micromamba@34071ca7df4983ccd272ed0d3625818b27b70dcc
+        uses: mamba-org/provision-with-micromamba@v14
         with:
           environment-file: ${{env.CONDA_ENV_FILE}}
           environment-name: xarray-tests
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -24,7 +24,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: ["ubuntu-latest"]
+        os: ["ubuntu-latest", "windows-latest"]
         python-version: ["3.8", "3.10"]
     steps:
       - uses: actions/checkout@v3
@@ -34,7 +34,7 @@ jobs:
         run: |
           echo "PYTHON_VERSION=${{ matrix.python-version }}" >> $GITHUB_ENV
       - name: Set up conda environment
-        uses: mamba-org/provision-with-micromamba@v13
+        uses: mamba-org/provision-with-micromamba@v14
         with:
           environment-file: ci/environment.yml
           environment-name: flox-tests
@@ -43,8 +43,7 @@ jobs:
             python="${{ matrix.python-version }}"
       - name: Install flox
         run: |
-          python -m pip install -e .
-          conda list
+          python -m pip install --no-deps -e .
       - name: Run Tests
         run: |
           pytest -n auto --cov=./ --cov-report=xml
@@ -78,7 +77,7 @@ jobs:
         with:
           fetch-depth: 0 # Fetch all history for all branches and tags.
       - name: Set up conda environment
-        uses: mamba-org/provision-with-micromamba@v13
+        uses: mamba-org/provision-with-micromamba@v14
         with:
           environment-file: ci/${{ matrix.env }}.yml
           environment-name: flox-tests
@@ -101,7 +100,7 @@ jobs:
     steps:
       - uses: actions/checkout@v3
       - name: Set up conda environment
-        uses: mamba-org/provision-with-micromamba@v13
+        uses: mamba-org/provision-with-micromamba@v14
         with:
           environment-file: ci/upstream-dev-env.yml
           environment-name: flox-tests
@@ -123,7 +122,7 @@ jobs:
           repository: 'pydata/xarray'
           fetch-depth: 0 # Fetch all history for all branches and tags.
       - name: Set up conda environment
-        uses: mamba-org/provision-with-micromamba@v13
+        uses: mamba-org/provision-with-micromamba@v14
         with:
           environment-file: ci/requirements/environment.yml
           environment-name: xarray-tests
diff --git a/asv_bench/benchmarks/combine.py b/asv_bench/benchmarks/combine.py
@@ -1,3 +1,5 @@
+from functools import partial
+
 import numpy as np
 
 import flox
@@ -7,26 +9,31 @@
 N = 1000
 
 
+def _get_combine(combine):
+    if combine == "grouped":
+        return partial(flox.core._grouped_combine, engine="numpy")
+    else:
+        return partial(flox.core._simple_combine, reindex=False)
+
+
 class Combine:
     def setup(self, *args, **kwargs):
         raise NotImplementedError
 
-    @parameterized("kind", ("cohorts", "mapreduce"))
-    def time_combine(self, kind):
-        flox.core._grouped_combine(
+    @parameterized(("kind", "combine"), (("reindexed", "not_reindexed"), ("grouped", "simple")))
+    def time_combine(self, kind, combine):
+        _get_combine(combine)(
             getattr(self, f"x_chunk_{kind}"),
             **self.kwargs,
             keepdims=True,
-            engine="numpy",
         )
 
-    @parameterized("kind", ("cohorts", "mapreduce"))
-    def peakmem_combine(self, kind):
-        flox.core._grouped_combine(
+    @parameterized(("kind", "combine"), (("reindexed", "not_reindexed"), ("grouped", "simple")))
+    def peakmem_combine(self, kind, combine):
+        _get_combine(combine)(
             getattr(self, f"x_chunk_{kind}"),
             **self.kwargs,
             keepdims=True,
-            engine="numpy",
         )
 
 
@@ -47,7 +54,7 @@ def construct_member(groups):
             }
 
         # motivated by
-        self.x_chunk_mapreduce = [
+        self.x_chunk_not_reindexed = [
             construct_member(groups)
             for groups in [
                 np.array((1, 2, 3, 4)),
@@ -57,5 +64,7 @@ def construct_member(groups):
             * 2
         ]
 
-        self.x_chunk_cohorts = [construct_member(groups) for groups in [np.array((1, 2, 3, 4))] * 4]
+        self.x_chunk_reindexed = [
+            construct_member(groups) for groups in [np.array((1, 2, 3, 4))] * 4
+        ]
         self.kwargs = {"agg": flox.aggregations.mean, "axis": (3,)}
diff --git a/asv_bench/benchmarks/reduce.py b/asv_bench/benchmarks/reduce.py
@@ -6,7 +6,7 @@
 from . import parameterized
 
 N = 1000
-funcs = ["sum", "nansum", "mean", "nanmean", "max"]
+funcs = ["sum", "nansum", "mean", "nanmean", "max", "var", "nanvar"]
 engines = ["flox", "numpy"]
 
 
diff --git a/ci/docs.yml b/ci/docs.yml
@@ -7,7 +7,7 @@ dependencies:
   - xarray
   - numpy>=1.20
   - numpydoc
-  - numpy_groupies
+  - numpy_groupies>=0.9.19
   - toolz
   - matplotlib-base
   - myst-parser
diff --git a/ci/environment.yml b/ci/environment.yml
@@ -16,7 +16,7 @@ dependencies:
   - pytest-xdist
   - xarray
   - pre-commit
-  - numpy_groupies>=0.9.15
+  - numpy_groupies>=0.9.19
   - pooch
   - toolz
   - numba
diff --git a/ci/minimal-requirements.yml b/ci/minimal-requirements.yml
@@ -9,7 +9,7 @@ dependencies:
   - pytest-cov
   - pytest-xdist
   - numpy==1.20
-  - numpy_groupies==0.9.15
+  - numpy_groupies==0.9.19
   - pandas
   - pooch
   - toolz
diff --git a/ci/no-dask.yml b/ci/no-dask.yml
@@ -13,7 +13,7 @@ dependencies:
   - xarray
   - numpydoc
   - pre-commit
-  - numpy_groupies>=0.9.15
+  - numpy_groupies>=0.9.19
   - pooch
   - toolz
   - numba
diff --git a/ci/no-xarray.yml b/ci/no-xarray.yml
@@ -13,7 +13,7 @@ dependencies:
   - dask-core
   - numpydoc
   - pre-commit
-  - numpy_groupies>=0.9.15
+  - numpy_groupies>=0.9.19
   - pooch
   - toolz
   - numba
diff --git a/flox/aggregate_npg.py b/flox/aggregate_npg.py
@@ -9,14 +9,43 @@ def _get_aggregate(engine):
 
 
 def sum_of_squares(
-    group_idx, array, engine, *, axis=-1, func="sum", size=None, fill_value=None, dtype=None
+    group_idx,
+    array,
+    engine,
+    *,
+    axis=-1,
+    size=None,
+    fill_value=None,
+    dtype=None,
 ):
 
     return _get_aggregate(engine).aggregate(
         group_idx,
-        array**2,
+        array,
         axis=axis,
-        func=func,
+        func="sumofsquares",
+        size=size,
+        fill_value=fill_value,
+        dtype=dtype,
+    )
+
+
+def nansum_of_squares(
+    group_idx,
+    array,
+    engine,
+    *,
+    axis=-1,
+    size=None,
+    fill_value=None,
+    dtype=None,
+):
+
+    return _get_aggregate(engine).aggregate(
+        group_idx,
+        array,
+        axis=axis,
+        func="nansumofsquares",
         size=size,
         fill_value=fill_value,
         dtype=dtype,
@@ -55,19 +84,6 @@ def nanprod(group_idx, array, engine, *, axis=-1, size=None, fill_value=None, dt
     )
 
 
-def nansum_of_squares(group_idx, array, engine, *, axis=-1, size=None, fill_value=None, dtype=None):
-    return sum_of_squares(
-        group_idx,
-        array,
-        engine=engine,
-        func="nansum",
-        size=size,
-        fill_value=fill_value,
-        axis=axis,
-        dtype=dtype,
-    )
-
-
 def _len(group_idx, array, engine, *, func, axis=-1, size=None, fill_value=None, dtype=None):
     result = _get_aggregate(engine).aggregate(
         group_idx,
diff --git a/flox/aggregations.py b/flox/aggregations.py
@@ -39,7 +39,7 @@ def generic_aggregate(
             method_ = getattr(aggregate_npg, func)
             method = partial(method_, engine=engine)
         except AttributeError:
-            aggregate = npg.aggregate_np if engine == "numpy" else npg.aggregate_nb
+            aggregate = aggregate_npg._get_aggregate(engine).aggregate
             method = partial(aggregate, func=func)
     else:
         raise ValueError(
diff --git a/flox/core.py b/flox/core.py
diff --git a/flox/visualize.py b/flox/visualize.py
diff --git a/setup.cfg b/setup.cfg
diff --git a/tests/test_core.py b/tests/test_core.py
