Rewrite pairwise to remove concatenation from blockwise

[aktech]

This rewrites pairwise distance to fix the concatenation in blockwise. The implementation is motivated from dask's matmul PR.

Fixes #375

I ran this on MalariaGEN data, it took nearly the same time as blocks implementation on 50 Workers, 16 GB each. The time comparions can be seen in the notebook here

The dask reports can be seen here:

Blocks implementation:

• dask report: 20 min 28 sec
• task stream

Blockwise implementation:

• dask report:
• task stream: 21 min 31 sec

Important note: These times are only rough estimates for comparison, not for quoting. When I ran the blocks implementation on the same configuration of coiled cloud couple of months ago it took ~ 10 min, now its taking more, since workers are dying due to: OSError: Timed out during handshake while connecting to tls:.... I am not sure why is that happening since no code change has bee made in the blocks implementation. I am communicating with coiled team to find the root cause at the moment.

cc @ravwojdyla

Update

A quick performance for correlation on only 10% of malariaGEN data:

• Notebook
• Blocks dask report: 56.25 sec
• Blockwise dask report: 58.94 sec

TODO:

• Remove notebooks and performance reports before merging.

[ravwojdyla]

Thanks for the update @aktech! Scanned this PR, looks good, some small comments below, but nothing critical. As discussed in the chat, it would be nice to see the correlation metric performance report between blocks and blockwise.

[aktech]

@ravwojdyla I think I have addressed all the comments. I have added a quick performance test for correlation metric as well (See PR description above). I was initially seeing huge difference in performance between blockwise and blocks implementation, as in the blockwise was incredibly slower.

After the analysis of the dask performance reports, I found out the compute time was same and the issue as due to serialisation/deserialisation as I had included the np.empty(..) repeatedly and lambdas in the implementation which increased the serialisation/deserialisation significantly. Now I have fixed those issues and blockwise is as performant as blocks implementation.

[aktech]

pre-commit seems to fail for unrelated changes:

sgkit/io/utils.py:186: error: Incompatible types in assignment (expression has type "number[Any]", variable has type "int")  [assignment]
sgkit/window.py:118: error: Item "ndarray" of "Union[ndarray, Any]" has no attribute "chunks"  [union-attr]
sgkit/stats/regenie.py:657: error: Incompatible types in assignment (expression has type "ndarray", variable has type "Optional[Sequence[float]]")  [assignment]
sgkit/stats/regenie.py:699: error: Argument "alphas" to "_stage_1" has incompatible type "Optional[Sequence[float]]"; expected "Optional[ndarray]"  [arg-type]
sgkit/stats/regenie.py:704: error: Argument "alphas" to "_stage_2" has incompatible type "Optional[Sequence[float]]"; expected "Optional[ndarray]"  [arg-type]
sgkit/stats/hwe.py:51: error: unused 'type: ignore' comment
sgkit/stats/hwe.py:94: error: unused 'type: ignore' comment
sgkit/stats/pca.py:133: error: Argument 1 to "__call__" of "ufunc" has incompatible type "DataArray"; expected "Union[Union[int, float, complex, str, bytes, generic], Sequence[Union[int, float, complex, str, bytes, generic]], Sequence[Sequence[Any]], _SupportsArray]"  [arg-type]
sgkit/io/vcfzarr_reader.py:199: error: Item "ndarray" of "Union[ndarray, Any]" has no attribute "values"  [union-attr]
sgkit/tests/test_pca.py:32: error: Returning Any from function declared to return "ndarray"  [no-any-return]
sgkit/io/bgen/bgen_reader.py:148: error: Returning Any from function declared to return "ndarray"  [no-any-return]
sgkit/io/bgen/bgen_reader.py:365: error: Non-overlapping container check (element type: "dtype[Any]", container item type: "Type[unsignedinteger[Any]]")  [comparison-overlap]
sgkit/tests/test_preprocessing.py:88: error: List item 0 has incompatible type "float"; expected "int"  [list-item]
sgkit/tests/test_preprocessing.py:124: error: Module has no attribute "bool"; maybe "bool_" or "bool8"?  [attr-defined]
sgkit/tests/test_preprocessing.py:162: error: Module has no attribute "bool"; maybe "bool_" or "bool8"?  [attr-defined]
sgkit/io/plink/plink_reader.py:87: error: Returning Any from function declared to return "ndarray"  [no-any-return]
sgkit/io/plink/plink_reader.py:268: error: Item "ndarray" of "Union[ndarray, Any]" has no attribute "compute"  [union-attr]
Found 17 errors in 10 files (checked 87 source files)

[tomwhite]

@aktech I filed #451 for these (unrelated) failures

[mergify]

This PR has conflicts, @aktech please rebase and push updated version 🙏

[aktech]

The tests failures are unrelated, comments have been addressed. @ravwojdyla

[codecov-io]

Codecov Report

Merging #447 (e754773) into master (9150392) will not change coverage.
The diff coverage is 100.00%.

@@            Coverage Diff            @@
##            master      #447   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           34        34           
  Lines         2361      2384   +23     
=========================================
+ Hits          2361      2384   +23     

Impacted Files	Coverage Δ
sgkit/distance/api.py	100.00% <100.00%> (ø)
sgkit/distance/metrics.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 9150392...e754773. Read the comment docs.

[ravwojdyla]

Good work @aktech! (please remove the notebooks before the merge)

[aktech]

@ravwojdyla Thanks for guiding me throughout, I have removed the notebooks.