diff --git a/cgp/genome.py b/cgp/genome.py
index 3c18a87f..4e8f035d 100644
--- a/cgp/genome.py
+++ b/cgp/genome.py
@@ -107,7 +107,7 @@ def dna(self) -> List[int]:
     def dna(self, value: List[int]) -> None:
         self._validate_dna(value)
         self._dna = value
-        self._initialize_unkown_parameters()
+        self._initialize_unknown_parameters()

     @property
     def _n_hidden(self) -> int:
@@ -332,6 +332,9 @@ def _is_output_region(self, region_idx: int) -> bool:
     def _is_function_gene(self, gene_idx: int) -> bool:
         return (gene_idx % self._length_per_region) == 0

+    def _is_hidden_input_gene(self, gene_idx: int, region_idx: int) -> bool:
+        return self._is_hidden_region(region_idx) & ((gene_idx % self._length_per_region) != 0)
+
     def _is_active_input_gene(self, gene_idx: int) -> bool:
         input_index = gene_idx % self._length_per_region
         assert input_index > 0
@@ -346,87 +349,113 @@ def _is_active_input_gene(self, gene_idx: int) -> bool:
         else:
             assert False  # should never be reached

+    def _select_gene_indices_for_mutation(
+        self, mutation_rate: float, rng: np.random.RandomState
+    ) -> List[int]:
+        """Selects the gene indices for mutations
+
+        Parameters
+        ----------
+        mutation_rate : float
+            Probability of a gene to be mutated, between 0 (excluded) and 1 (included).
+        rng : numpy.random.RandomState
+            Random number generator instance to use for selecting the indices.
+
+        Returns
+        ----------
+        selected_gene_indices: np.ndarray
+            indices of the genes selected for mutation.
+        """
+
+        selected_gene_indices = np.nonzero(rng.rand(len(self.dna)) < mutation_rate)[0]
+        return selected_gene_indices
+
     def mutate(self, mutation_rate: float, rng: np.random.RandomState):
         """Mutate the genome.

         Parameters
         ----------
         mutation_rate : float
-            Proportion of genes to be mutated, between 0 and 1.
+            Probability of a gene to be mutated, between 0 (excluded) and 1 (included).
         rng : numpy.random.RandomState
             Random number generator instance to use for mutating.

+
         Returns
         ----------
         bool
             True if only inactive regions of the genome were mutated, False otherwise.
         """

-        def count_dna_differences(dna0: List[int], dna1: List[int]) -> int:
-            return len([1 for gene0, gene1 in zip(dna0, dna1) if gene0 != gene1])
-
-        n_mutations = int(mutation_rate * len(self.dna))
-        assert n_mutations > 0
-
         graph = CartesianGraph(self)
         active_regions = graph.determine_active_regions()
-
         dna = list(self._dna)
-
         only_silent_mutations = True
-        while count_dna_differences(self.dna, dna) < n_mutations:

-            gene_idx = rng.randint(0, self._n_genes)
-            region_idx = gene_idx // self._length_per_region
+        selected_gene_indices = self._select_gene_indices_for_mutation(mutation_rate, rng)

-            if self._is_input_region(region_idx):
-                continue  # nothing to do here
+        for (gene_idx, allele) in zip(selected_gene_indices, np.array(dna)[selected_gene_indices]):

-            elif self._is_output_region(region_idx):
-                silent = self._mutate_output_region(dna, gene_idx, rng)
+            region_idx = gene_idx // self._length_per_region

-            elif self._is_hidden_region(region_idx):
-                silent = self._mutate_hidden_region(dna, gene_idx, active_regions, rng)
+            permissible_values = self._determine_alternative_permissible_values(
+                gene_idx, allele, region_idx
+            )
+            if len(permissible_values) > 0:

-            else:
-                assert False  # should never be reached
+                dna[gene_idx] = rng.choice(permissible_values)
+                silent = region_idx not in active_regions

-            only_silent_mutations = only_silent_mutations and silent
+                only_silent_mutations = only_silent_mutations and silent

         self.dna = dna
         return only_silent_mutations

-    def _mutate_output_region(
-        self, dna: List[int], gene_idx: int, rng: np.random.RandomState
-    ) -> bool:
-        assert self._is_gene_in_output_region(gene_idx)
+    def _determine_alternative_permissible_values(
+        self, gene_idx: int, allele: int, region_idx: int
+    ) -> List[int]:

-        if not self._is_function_gene(gene_idx) and self._is_active_input_gene(gene_idx):
-            permissible_inputs = self._permissible_inputs_for_output_region()
-            dna[gene_idx] = rng.choice(permissible_inputs)
-            return False
-        else:
-            return True
+        if self._is_input_region(region_idx):
+            return []  # genes in input regions have no alternative permissible values

-    def _mutate_hidden_region(
-        self, dna: List[int], gene_idx: int, active_regions: List[int], rng: np.random.RandomState
-    ) -> bool:
+        elif self._is_hidden_region(region_idx):
+            return self._determine_alternative_permissible_values_hidden_gene(
+                gene_idx, allele, region_idx
+            )

-        assert self._is_gene_in_hidden_region(gene_idx)
+        elif self._is_output_region(region_idx):
+            return self._determine_alternative_permissible_values_output_gene(gene_idx, allele)

-        region_idx = gene_idx // self._length_per_region
-        silent_mutation = region_idx not in active_regions
+        else:
+            assert False  # should never be reached

+    def _determine_alternative_permissible_values_hidden_gene(
+        self, gene_idx: int, allele: int, region_idx: int
+    ) -> List[int]:
         if self._is_function_gene(gene_idx):
-            dna[gene_idx] = self._primitives.sample_allele(rng)
-            return silent_mutation
+            permissible_values = list(np.arange(len(self._primitives._primitives)))

+        elif self._is_hidden_input_gene(gene_idx, region_idx):
+            permissible_values = self._permissible_inputs(region_idx)
+
+        else:
+            assert False
+        permissible_values.remove(allele)
+        return permissible_values
+
+    def _determine_alternative_permissible_values_output_gene(
+        self, gene_idx: int, allele: int
+    ) -> List[int]:
+        input_index = (
+            gene_idx % self._length_per_region
+        )  # assumes that the second gene in all output regions is the index of the input
+        if input_index == 1:
+            permissible_values = self._permissible_inputs_for_output_region()
+            permissible_values.remove(allele)
         else:
-            permissible_inputs = self._permissible_inputs(region_idx)
-            dna[gene_idx] = rng.choice(permissible_inputs)
+            permissible_values = []

-            silent_mutation = silent_mutation or (not self._is_active_input_gene(gene_idx))
-            return silent_mutation
+        return permissible_values

     @property
     def primitives(self) -> Primitives:
@@ -484,7 +513,7 @@ def update_parameters_from_torch_class(self, torch_cls: "torch.nn.Module") -> bo

         return any_parameter_updated

-    def _initialize_unkown_parameters(self) -> None:
+    def _initialize_unknown_parameters(self) -> None:
         for region_idx, region in self.iter_hidden_regions():
             node_id = region[0]
             node_type = self._primitives[node_id]
diff --git a/cgp/population.py b/cgp/population.py
index f2ccce23..114d70f9 100755
--- a/cgp/population.py
+++ b/cgp/population.py
@@ -25,8 +25,7 @@ def __init__(
         n_parents : int
             Number of parent individuals.
         mutation_rate : float
-            Rate of mutations determining the number of genes to be
-            mutated for offspring creation, between 0 and 1.
+            Probability of a gene to be mutated, between 0 (excluded) and 1 (included).
         seed : int
             Seed for internal random number generator.
         genome_params : dict
@@ -34,8 +33,8 @@ def __init__(
         """
         self.n_parents = n_parents  # number of individuals in parent population

-        if not (0.0 < mutation_rate and mutation_rate < 1.0):
-            raise ValueError("mutation rate needs to be in (0, 1)")
+        if not (0.0 < mutation_rate and mutation_rate <= 1.0):
+            raise ValueError("mutation rate needs to be in (0, 1]")
         self._mutation_rate = mutation_rate  # probability of mutation per gene

         self.seed = seed

diff --git a/test/test_ea_mu_plus_lambda.py b/test/test_ea_mu_plus_lambda.py
index cf0e1b2b..8f0521ea 100644
--- a/test/test_ea_mu_plus_lambda.py
+++ b/test/test_ea_mu_plus_lambda.py
@@ -196,6 +196,8 @@ def objective(individual):
         individual.fitness = float(individual.idx)
         return individual

+    population_params["mutation_rate"] = 1.0  # ensures every offspring has mutations
+
     pop = cgp.Population(**population_params, genome_params=genome_params)
     ea = cgp.ea.MuPlusLambda(**ea_params)

diff --git a/test/test_genome.py b/test/test_genome.py
index 55d1a892..17106b86 100644
--- a/test/test_genome.py
+++ b/test/test_genome.py
@@ -3,6 +3,7 @@

 import cgp
 from cgp.genome import ID_INPUT_NODE, ID_OUTPUT_NODE, ID_NON_CODING_GENE
+from cgp.cartesian_graph import CartesianGraph

 def test_check_dna_consistency():
@@ -346,53 +347,91 @@ def test_is_gene_in_output_region(rng_seed):
     assert not genome._is_gene_in_output_region(11)

-def test_mutate_hidden_region(rng_seed):
+def test_mutation_rate(rng_seed, mutation_rate):
+    n_inputs = 1
+    n_outputs = 1
+    n_columns = 4
+    n_rows = 3
+    genome = cgp.Genome(n_inputs, n_outputs, n_columns, n_rows, None, (cgp.Add, cgp.Sub))
     rng = np.random.RandomState(rng_seed)
-    genome = cgp.Genome(1, 1, 3, 1, None, (cgp.Add, cgp.ConstantFloat))
-    dna = [
+    genome.randomize(rng)
+
+    def count_n_immutable_genes(n_inputs, n_outputs, n_rows):
+        length_per_region = genome.primitives.max_arity + 1  # function gene + input gene addresses
+        n_immutable_genes = n_inputs * length_per_region  # none of the input genes are mutable
+        n_immutable_genes += n_outputs * (
+            length_per_region - 1
+        )  # only one gene per output can be mutated
+        if n_inputs == 1:
+            n_immutable_genes += (
+                n_rows * genome.primitives.max_arity
+            )  # input gene addresses in the first (hidden) column can't be mutated
+            # if only one input node exists
+        return n_immutable_genes
+
+    def count_mutations(dna0, dna1):
+        n_differences = 0
+        for (allele0, allele1) in zip(dna0, dna1):
+            if allele0 != allele1:
+                n_differences += 1
+        return n_differences
+
+    n_immutable_genes = count_n_immutable_genes(n_inputs, n_outputs, n_rows)
+    n_mutations_mean_expected = mutation_rate * (len(genome.dna) - n_immutable_genes)
+    n_mutations_std_expected = np.sqrt(
+        (len(genome.dna) - n_immutable_genes) * mutation_rate * (1 - mutation_rate)
+    )
+
+    n = 10000
+    n_mutations = []
+    for _ in range(n):
+        dna_old = genome.dna
+        genome.mutate(mutation_rate, rng)
+        n_mutations.append(count_mutations(dna_old, genome.dna))
+
+    assert np.mean(n_mutations) == pytest.approx(n_mutations_mean_expected, rel=0.04)
+    assert np.std(n_mutations) == pytest.approx(n_mutations_std_expected, rel=0.04)
+
+
+def test_only_silent_mutations(genome_params, mutation_rate, rng_seed):
+    genome = cgp.Genome(**genome_params)
+    rng = np.random.RandomState(rng_seed)
+    genome.randomize(rng)
+
+    only_silent_mutations = genome.mutate(mutation_rate=0, rng=rng)
+    assert only_silent_mutations is True
+
+    only_silent_mutations = genome.mutate(mutation_rate=1, rng=rng)
+    assert not only_silent_mutations
+
+    dna_fixed = [
+        ID_INPUT_NODE,
+        ID_NON_CODING_GENE,
+        ID_NON_CODING_GENE,
         ID_INPUT_NODE,
         ID_NON_CODING_GENE,
         ID_NON_CODING_GENE,
+        2,
         1,
         0,
         0,
-        1,
         0,
         0,
         0,
         0,
-        2,
-        ID_OUTPUT_NODE,
-        3,
-        ID_NON_CODING_GENE,
-    ]
-    genome.dna = list(dna)
-    active_regions = cgp.CartesianGraph(genome).determine_active_regions()
-
-    # mutating any gene in inactive region returns True
-    assert genome._mutate_hidden_region(list(dna), 3, active_regions, rng) is True
-    assert genome._mutate_hidden_region(list(dna), 4, active_regions, rng) is True
-    assert genome._mutate_hidden_region(list(dna), 5, active_regions, rng) is True
-
-    # mutating function gene in active region returns False
-    assert genome._mutate_hidden_region(list(dna), 6, active_regions, rng) is False
-    # mutating inactive genes in active region returns True
-    assert genome._mutate_hidden_region(list(dna), 7, active_regions, rng) is True
-    assert genome._mutate_hidden_region(list(dna), 8, active_regions, rng) is True
-
-    # mutating any gene in active region without silent genes returns False
-    assert genome._mutate_hidden_region(list(dna), 9, active_regions, rng) is False
-    assert genome._mutate_hidden_region(list(dna), 10, active_regions, rng) is False
-    assert genome._mutate_hidden_region(list(dna), 11, active_regions, rng) is False
-
-
-def test_mutate_output_region(rng_seed):
-    rng = np.random.RandomState(rng_seed)
-    genome = cgp.Genome(1, 1, 2, 1, None, (cgp.Add,))
-    dna = [
-        ID_INPUT_NODE,
-        ID_NON_CODING_GENE,
-        ID_NON_CODING_GENE,
+        0,
+        0,
+        0,
+        0,
+        0,
+        0,
+        0,
+        0,
+        0,
+        0,
+        0,
+        0,
+        0,
         0,
         0,
         0,
@@ -403,32 +442,110 @@ def test_mutate_output_region(rng_seed):
         2,
         ID_NON_CODING_GENE,
     ]
+    genome.dna = dna_fixed
+    graph = CartesianGraph(genome)
+    active_regions = graph.determine_active_regions()
+    length_per_region = genome.primitives.max_arity + 1  # function gene + input gene addresses
+    gene_to_be_mutated_non_active = (
+        3 * length_per_region
+    )  # operator gene of 2nd hidden node, a silent node in this dna configuration
+
+    def select_gene_indices_silent(mutation_rate, rng):
+        selected_gene_indices = [gene_to_be_mutated_non_active]
+        return selected_gene_indices
+
+    genome._select_gene_indices_for_mutation = (
+        select_gene_indices_silent  # monkey patch the selection of indices to select a silent gene
+    )
+    genome.dna = dna_fixed
+    only_silent_mutations = genome.mutate(mutation_rate, rng)
+    assert only_silent_mutations is True
+
+    gene_to_be_mutated_active = (
+        active_regions[-1] * length_per_region
+    )  # operator gene of the 1st active hidden node,
+    # should always be mutable and result in a non-silent-mutation
+
+    def select_gene_indices_non_silent(mutation_rate, rng):
+        selected_gene_indices = [gene_to_be_mutated_active]
+        return selected_gene_indices
+
+    genome._select_gene_indices_for_mutation = select_gene_indices_non_silent
+    # monkey patch the selection of indices to select a non-silent gene
+    only_silent_mutations = genome.mutate(mutation_rate, rng)
+    assert not only_silent_mutations
+
+
+def test_permissible_values_hidden(rng_seed):
+    genome_params = {
+        "n_inputs": 2,
+        "n_outputs": 1,
+        "n_columns": 3,
+        "n_rows": 3,
+        "levels_back": 2,
+        "primitives": (cgp.Add, cgp.Sub, cgp.ConstantFloat),
+    }
+
+    genome = cgp.Genome(**genome_params)
+
+    # test function gene
+    gene_idx = 6  # function gene of first hidden region
+    allele = 0  #
+    region_idx = None  # can be any number, since not touched for function gene
+    permissible_func_gene_values = genome._determine_alternative_permissible_values_hidden_gene(
+        gene_idx, allele, region_idx
+    )
+    assert permissible_func_gene_values == [
+        1,
+        2,
+    ]  # function idx 1,2 (cgp.Sub, cgp.ConstantFloat)
+
+    # test input gene
+    gene_idx = 16  # first input gene in second column of hidden region
+    allele = 0
+    region_idx = gene_idx // (
+        genome.primitives.max_arity + 1
+    )  # function gene + input gene addresses
+    permissible_input_gene_values = genome._determine_alternative_permissible_values_hidden_gene(
+        gene_idx, allele, region_idx
+    )

-    assert genome._mutate_output_region(list(dna), 9, rng) is True
-    assert genome._mutate_output_region(list(dna), 10, rng) is False
-    assert genome._mutate_output_region(list(dna), 11, rng) is True
+    assert permissible_input_gene_values == [
+        1,
+        2,
+        3,
+        4,
+    ]  # gene indices of input, 1st hidden layer

-@pytest.mark.parametrize("mutation_rate", [0.02, 0.05, 0.2])
-def test_correct_number_of_mutations(mutation_rate, rng_seed):
+def test_permissible_values_output(rng_seed):
+    genome_params = {
+        "n_inputs": 2,
+        "n_outputs": 1,
+        "n_columns": 3,
+        "n_rows": 3,
+        "levels_back": 2,
+        "primitives": (cgp.Add, cgp.Sub, cgp.ConstantFloat),
+    }
+    genome = cgp.Genome(**genome_params)

-    n_inputs = 2
-    n_outputs = 1
-    n_columns = 10
-    n_rows = 2
-    levels_back = 5
-    primitives = (cgp.Add, cgp.Sub, cgp.Mul, cgp.Div, cgp.ConstantFloat)
+    gene_idx_function = 33
+    gene_idx_input0 = 34
+    gene_idx_input1 = 35

-    rng = np.random.RandomState(rng_seed)
-    genome = cgp.Genome(n_inputs, n_outputs, n_columns, n_rows, levels_back, primitives)
-    genome.randomize(rng)
+    allele = 5  # set current value for input0 -> should be excluded from permissible values

-    n_mutations = 0
-    genome_new = genome.clone()
-    genome_new.mutate(mutation_rate, rng)
-    for (gene_0, gene_1) in zip(genome.dna, genome_new.dna):
-        if gene_0 != gene_1:
-            n_mutations += 1
+    permissible_values = genome._determine_alternative_permissible_values_output_gene(
+        gene_idx_function, allele
+    )
+    assert permissible_values == []

-    n_mutations_expected = int(mutation_rate * len(genome.dna))
-    assert n_mutations == n_mutations_expected
+    permissible_values = genome._determine_alternative_permissible_values_output_gene(
+        gene_idx_input0, allele
+    )
+    assert permissible_values == [0, 1, 2, 3, 4, 6, 7, 8, 9, 10]
+
+    permissible_values = genome._determine_alternative_permissible_values_output_gene(
+        gene_idx_input1, allele
+    )
+    assert permissible_values == []
diff --git a/test/test_hl_api.py b/test/test_hl_api.py
index 9ea76dd1..50837686 100644
--- a/test/test_hl_api.py
+++ b/test/test_hl_api.py
@@ -65,6 +65,7 @@ def test_parallel_population(population_params, genome_params, ea_params):
         fitness_per_n_processes.append(
             _test_population(population_params, genome_params, ea_params)
         )
+
     assert fitness_per_n_processes[0] == pytest.approx(fitness_per_n_processes[1])
     assert fitness_per_n_processes[0] == pytest.approx(fitness_per_n_processes[2])

Issue warning when `mutation_rate` is high and `mutate` potentially takes long
#157 made sure that exactly the correct number of mutations occurs in a genome by checking that the mutated genome differs from the original one in `n_mutations` position. when using a high mutation rate and small number of primitives, this can lead to a long runtime for `mutate` as it struggles to find positions in the genome to mutate and alternative genes to insert. we should figure out exactly under which conditions this tends to occur and issue a user warning when we detect such a situation.

@HenrikMettler, since you're already familiar with that part of the code, maybe you want to look into this?

btw, an alternative solution to this problem would be to have a maximum number iterations in the while loop (https://github.com/Happy-Algorithms-League/hal-cgp/blob/master/cgp/genome.py#L377) and return even if the desired number of differences was not reached while issuing a warning to the user
What about keeping track of the genes that were previously mutated and only allowing `gene_idx` to take different values?
Perhaps I am missing something here, but won't this problem go away once we finally fix the current implementation to not apply a fixed number of mutations but rather mutate each gene with a fixed probability, as the variable names suggest, i.e. once #124 is closed?
Actually #124 (sorry I will do a PR soonish) only calculates a different `n_mutations` for every individual in every generation (from a binomial distribution). It does not apply a mutation with a probability to every single gene (this would be computationally more expensive I think). So the `while` loop in `Genome.mutate` still looks the same, just with different values for `n_mutations`
I doubt that making a coin flip for each gene in each individual is slower.
For instance, running this code:

```python
import numpy as np
import time

n_individuals = 10
n_genes = 15

p = 0.2

t0 = time.time()
for i in range(1000):
    np.random.binomial(n_genes, p, size=n_individuals)

print(f"Binomial took {time.time() - t0} seconds.")

t0 = time.time()
for i in range(1000):
    indx = np.random.rand(n_individuals, n_genes) < p

print(f"Bernoulli took {time.time() - t0} seconds.")
```

gives me:
```
Binomial took 0.0036504268646240234 seconds.
Bernoulli took 0.002948284149169922 seconds.
```

yes, it seems like you are right @mschmidt87, my intuition would've been with @HenrikMettler, oh well this is why one should always benchmark :flushed:

however, i don't think checking every gene individually would remove the fundamental problem that we would need to detect if a user wants to modify more genes than could actually be changed. two things can happen otherwise: i) we have an infinite loop ii) the effective mutation rate is smaller than the one requested by the user without them noticing.

one suggestion to resolve this (just discussed with @HenrikMettler) would be to determine the set of all "mutable genes", make sure that set is larger than `n_mutations`, and then pick `n_mutations` genes from that set to actually be mutated. a PR will follow
If I understand correctly, there are several types of genes that cannot be mutated:
- input region: no gene can be mutated
- output region: function gene and is_active_input gene cannot be mutated

Thus, the number of mutable genes in a genome is `len(genome) - n_immutable_genes`. This leads to a problem if the user understand `mutation_rate` as determining the number of mutated genes by doing `n_mutations = int(mutation_rate * len(self.dna))`. This is what we're doing now and how it is documented. If we do it like that, then we need to add additional checks etc.

However, if `mutation_rate` is understood as the probability for each mutable gene to be mutated, then these problems just go away. We can simply do:
```python
## is_mutable = array of length len(dna) with True/False for each gene
mutated_idx = np.where(np.logical_and(np.random.rand(len(dna)) < p, is_mutable))
dna[mutated_idx] = rng.choice(permissible_inputs, size=len(mutated_idx))
```

This is just a question of the exact definition/understanding of the mutation rate, and I would opt for the interpretation that is simpler and makes more sense, in my opinion.
I agree with your analysis and (of course short ;)) implementation. i would like to add that for some parameter configurations there might be only one option for a gene, which makes it questionable whether replacing this gene with itself should be counted as a mutation or not.

that being said, I think the two approaches (calculating the total number of mutations once as `n_mutations = int(mutation_rate * len(self.dna))` or mutating each gene with a certain probability) are identical as long as we count trying to mutate an gene that can not be mutated (as in your definition above) as a successful mutation. in both cases, the number of genes that are different between parent and offspring will be smaller than expected based on the total number of genes and the mutation rate.

after this discussion and considering the implementation i'm starting to like this version a bit better, however the issue remains that adopting this version offsprings can be generated _which are genetically identical to their parents_. as we rely on mutations to drive search this null operation seems wasteful to me, but from quickly scanning other implementations it seems quite common. this might be related to, as you said, this interpretation of the mutation rate being easier to explain/understand.

@HenrikMettler do you have any strong feelings either way? i mean besides having already working code for one ;)
I don't, given that computationally there seems to be no significant difference. (I vaguely remember reading somewhere in the literature that there is one, so I will quickly check again - but probably it is just my memory fooling me)

I would just argue that for the `permissible_inputs` we do:
`permissible_inputs = self._permissible_inputs()`
`permissible_inputs.remove(dna[idx])`

Since a gene for which this list is empty (eg. a hidden gene, that can only have one possible input, in a specific graph structure) it should be that: `is_mutable = false`  I would do this before defining which elements should be mutated


diff --git a/sgkit/distance/api.py b/sgkit/distance/api.py
index 70141b598..7df70eb16 100644
--- a/sgkit/distance/api.py
+++ b/sgkit/distance/api.py
@@ -1,14 +1,20 @@
+import typing
+
 import dask.array as da
 import numpy as np
+from typing_extensions import Literal

 from sgkit.distance import metrics
 from sgkit.typing import ArrayLike

+MetricTypes = Literal["euclidean", "correlation"]
+

 def pairwise_distance(
     x: ArrayLike,
-    metric: str = "euclidean",
-) -> np.ndarray:
+    metric: MetricTypes = "euclidean",
+    split_every: typing.Optional[int] = None,
+) -> da.array:
     """Calculates the pairwise distance between all pairs of row vectors in the
     given two dimensional array x.

@@ -47,6 +53,13 @@ def pairwise_distance(
     metric
         The distance metric to use. The distance function can be
         'euclidean' or 'correlation'.
+    split_every
+        Determines the depth of the recursive aggregation in the reduction
+        step. This argument is directly passed to the call to``dask.reduction``
+        function in the reduce step of the map reduce.
+
+        Omit to let dask heuristically decide a good default. A default can
+        also be set globally with the split_every key in dask.config.

     Returns
     -------
@@ -63,41 +76,91 @@ def pairwise_distance(
     >>> from sgkit.distance.api import pairwise_distance
     >>> import dask.array as da
     >>> x = da.array([[6, 4, 1,], [4, 5, 2], [9, 7, 3]]).rechunk(2, 2)
-    >>> pairwise_distance(x, metric='euclidean')
+    >>> pairwise_distance(x, metric='euclidean').compute()
     array([[0.        , 2.44948974, 4.69041576],
            [2.44948974, 0.        , 5.47722558],
            [4.69041576, 5.47722558, 0.        ]])

     >>> import numpy as np
     >>> x = np.array([[6, 4, 1,], [4, 5, 2], [9, 7, 3]])
-    >>> pairwise_distance(x, metric='euclidean')
+    >>> pairwise_distance(x, metric='euclidean').compute()
     array([[0.        , 2.44948974, 4.69041576],
            [2.44948974, 0.        , 5.47722558],
            [4.69041576, 5.47722558, 0.        ]])

     >>> x = np.array([[6, 4, 1,], [4, 5, 2], [9, 7, 3]])
-    >>> pairwise_distance(x, metric='correlation')
-    array([[1.11022302e-16, 2.62956526e-01, 2.82353505e-03],
-           [2.62956526e-01, 0.00000000e+00, 2.14285714e-01],
-           [2.82353505e-03, 2.14285714e-01, 0.00000000e+00]])
+    >>> pairwise_distance(x, metric='correlation').compute()
+    array([[-4.44089210e-16,  2.62956526e-01,  2.82353505e-03],
+           [ 2.62956526e-01,  0.00000000e+00,  2.14285714e-01],
+           [ 2.82353505e-03,  2.14285714e-01,  0.00000000e+00]])
     """
-
     try:
-        metric_ufunc = getattr(metrics, metric)
+        metric_map_func = getattr(metrics, f"{metric}_map")
+        metric_reduce_func = getattr(metrics, f"{metric}_reduce")
+        n_map_param = metrics.N_MAP_PARAM[metric]
     except AttributeError:
         raise NotImplementedError(f"Given metric: {metric} is not implemented.")

     x = da.asarray(x)
-    x_distance = da.blockwise(
-        # Lambda wraps reshape for broadcast
-        lambda _x, _y: metric_ufunc(_x[:, None, :], _y),
-        "jk",
+    if x.ndim != 2:
+        raise ValueError(f"2-dimensional array expected, got '{x.ndim}'")
+
+    # setting this variable outside of _pairwise to avoid it's recreation
+    # in every iteration, which eventually leads to increase in dask
+    # graph serialisation/deserialisation time significantly
+    metric_param = np.empty(n_map_param, dtype=x.dtype)
+
+    def _pairwise(f: ArrayLike, g: ArrayLike) -> ArrayLike:
+        result: ArrayLike = metric_map_func(f[:, None, :], g, metric_param)
+        # Adding a new axis to help combine chunks along this axis in the
+        # reduction step (see the _aggregate and _combine functions below).
+        return result[..., np.newaxis]
+
+    # concatenate in blockwise leads to high memory footprints, so instead
+    # we perform blockwise without contraction followed by reduction.
+    # More about this issue: https://github.com/dask/dask/issues/6874
+    out = da.blockwise(
+        _pairwise,
+        "ijk",
         x,
-        "ji",
+        "ik",
         x,
-        "ki",
-        dtype="float64",
-        concatenate=True,
+        "jk",
+        dtype=x.dtype,
+        concatenate=False,
+    )
+
+    def _aggregate(x_chunk: ArrayLike, **_: typing.Any) -> ArrayLike:
+        """Last function to be executed when resolving the dask graph,
+        producing the final output. It is always invoked, even when the reduced
+        Array counts a single chunk along the reduced axes."""
+        x_chunk = x_chunk.reshape(x_chunk.shape[:-2] + (-1, n_map_param))
+        result: ArrayLike = metric_reduce_func(x_chunk)
+        return result
+
+    def _chunk(x_chunk: ArrayLike, **_: typing.Any) -> ArrayLike:
+        return x_chunk
+
+    def _combine(x_chunk: ArrayLike, **_: typing.Any) -> ArrayLike:
+        """Function used for intermediate recursive aggregation (see
+        split_every argument to ``da.reduction below``).  If the
+        reduction can be performed in less than 3 steps, it will
+        not be invoked at all."""
+        # reduce chunks by summing along the -2 axis
+        x_chunk_reshaped = x_chunk.reshape(x_chunk.shape[:-2] + (-1, n_map_param))
+        return x_chunk_reshaped.sum(axis=-2)[..., np.newaxis]
+
+    r = da.reduction(
+        out,
+        chunk=_chunk,
+        combine=_combine,
+        aggregate=_aggregate,
+        axis=-1,
+        dtype=x.dtype,
+        meta=np.ndarray((0, 0), dtype=x.dtype),
+        split_every=split_every,
+        name="pairwise",
     )
-    x_distance = da.triu(x_distance, 1) + da.triu(x_distance).T
-    return x_distance.compute()  # type: ignore[no-any-return]
+
+    t = da.triu(r)
+    return t + t.T
diff --git a/sgkit/distance/metrics.py b/sgkit/distance/metrics.py
index 9466671dc..1c0b3433b 100644
--- a/sgkit/distance/metrics.py
+++ b/sgkit/distance/metrics.py
@@ -1,52 +1,112 @@
-"""This module implements various distance metrics."""
+"""
+This module implements various distance metrics. To implement a new distance
+metric, two methods needs to be written, one of them suffixed by 'map' and other
+suffixed by 'reduce'. An entry for the same should be added in the N_MAP_PARAM
+dictionary below.
+"""

 import numpy as np
 from numba import guvectorize

 from sgkit.typing import ArrayLike

+# The number of parameters for the map step of the respective distance metric
+N_MAP_PARAM = {
+    "correlation": 6,
+    "euclidean": 1,
+}
+
+
+@guvectorize(  # type: ignore
+    [
+        "void(float32[:], float32[:], float32[:], float32[:])",
+        "void(float64[:], float64[:], float64[:], float64[:])",
+        "void(int8[:], int8[:], int8[:], float64[:])",
+    ],
+    "(n),(n),(p)->(p)",
+    nopython=True,
+    cache=True,
+)
+def euclidean_map(
+    x: ArrayLike, y: ArrayLike, _: ArrayLike, out: ArrayLike
+) -> None:  # pragma: no cover
+    """Euclidean distance "map" function for partial vector pairs.
+
+    Parameters
+    ----------
+    x
+        An array chunk, a partial vector
+    y
+        Another array chunk, a partial vector
+    _
+        A dummy variable to map the size of output
+    out
+        The output array, which has the squared sum of partial vector pairs.
+
+    Returns
+    -------
+    An ndarray, which contains the output of the calculation of the application
+    of euclidean distance on the given pair of chunks, without the aggregation.
+    """
+    square_sum = 0.0
+    m = x.shape[0]
+    # Ignore missing values
+    for i in range(m):
+        if x[i] >= 0 and y[i] >= 0:
+            square_sum += (x[i] - y[i]) ** 2
+    out[:] = square_sum
+
+
+def euclidean_reduce(v: ArrayLike) -> ArrayLike:  # pragma: no cover
+    """Corresponding "reduce" function for euclidean distance.
+
+    Parameters
+    ----------
+    v
+        The euclidean array on which map step of euclidean distance has been
+        applied.
+
+    Returns
+    -------
+    An ndarray, which contains square root of the sum of the squared sums obtained from
+    the map step of `euclidean_map`.
+    """
+    out: ArrayLike = np.sqrt(np.einsum("ijkl -> ij", v))
+    return out
+

 @guvectorize(  # type: ignore
     [
-        "void(float32[:], float32[:], float32[:])",
-        "void(float64[:], float64[:], float64[:])",
-        "void(int8[:], int8[:], float64[:])",
+        "void(float32[:], float32[:], float32[:], float32[:])",
+        "void(float64[:], float64[:], float64[:], float64[:])",
+        "void(int8[:], int8[:], int8[:], float64[:])",
     ],
-    "(n),(n)->()",
+    "(n),(n),(p)->(p)",
     nopython=True,
     cache=True,
 )
-def correlation(x: ArrayLike, y: ArrayLike, out: ArrayLike) -> None:  # pragma: no cover
-    """Calculates the correlation between two vectors.
+def correlation_map(
+    x: ArrayLike, y: ArrayLike, _: ArrayLike, out: ArrayLike
+) -> None:  # pragma: no cover
+    """Pearson correlation "map" function for partial vector pairs.

     Parameters
     ----------
     x
-        [array-like, shape: (M,)]
-        A vector
+        An array chunk, a partial vector
     y
-        [array-like, shape: (M,)]
-        Another vector
+        Another array chunk, a partial vector
+    _
+        A dummy variable to map the size of output
     out
         The output array, which has the output of pearson correlation.

     Returns
     -------
-    A scalar representing the pearson correlation coefficient between two vectors x and y.
-
-    Examples
-    --------
-    >>> from sgkit.distance.metrics import correlation
-    >>> import dask.array as da
-    >>> import numpy as np
-    >>> x = da.array([4, 3, 2, 3], dtype='i1')
-    >>> y = da.array([5, 6, 7, 0], dtype='i1')
-    >>> correlation(x, y).compute()
-    1.2626128
-
-    >>> correlation(x, x).compute()
-    -1.1920929e-07
+    An ndarray, which contains the output of the calculation of the application
+    of pearson correlation on the given pair of chunks, without the aggregation.
     """
+
     m = x.shape[0]
     valid_indices = np.zeros(m, dtype=np.float64)

@@ -66,61 +126,49 @@ def correlation(x: ArrayLike, y: ArrayLike, out: ArrayLike) -> None:  # pragma:
             _y[valid_idx] = y[i]
             valid_idx += 1

-    cov = ((_x - _x.mean()) * (_y - _y.mean())).sum()
-    denom = (_x.std() * _y.std()) / _x.shape[0]
-
-    value = np.nan
-    if denom > 0:
-        value = 1.0 - (cov / (_x.std() * _y.std()) / _x.shape[0])
-    out[0] = value
+    out[:] = np.array(
+        [
+            np.sum(_x),
+            np.sum(_y),
+            np.sum(_x * _x),
+            np.sum(_y * _y),
+            np.sum(_x * _y),
+            len(_x),
+        ]
+    )

 @guvectorize(  # type: ignore
     [
-        "void(float32[:], float32[:], float32[:])",
-        "void(float64[:], float64[:], float64[:])",
-        "void(int8[:], int8[:], float64[:])",
+        "void(float32[:, :], float32[:])",
+        "void(float64[:, :], float64[:])",
     ],
-    "(n),(n)->()",
+    "(p, m)->()",
     nopython=True,
     cache=True,
 )
-def euclidean(x: ArrayLike, y: ArrayLike, out: ArrayLike) -> None:  # pragma: no cover
-    """Calculates the euclidean distance between two vectors.
-
+def correlation_reduce(v: ArrayLike, out: ArrayLike) -> None:  # pragma: no cover
+    """Corresponding "reduce" function for pearson correlation
     Parameters
     ----------
-    x
-        [array-like, shape: (M,)]
-        A vector
-    y
-        [array-like, shape: (M,)]
-        Another vector
+    v
+        The correlation array on which pearson corrections has been
+        applied on chunks
     out
-        The output scalar, which has the output of euclidean between two vectors.
+        An ndarray, which is a symmetric matrix of pearson correlation

     Returns
     -------
-    A scalar representing the euclidean distance between two vectors x and y.
-
-    Examples
-    --------
-    >>> from sgkit.distance.metrics import euclidean
-    >>> import dask.array as da
-    >>> import numpy as np
-    >>> x = da.array([4, 3, 2, 3], dtype='i1')
-    >>> y = da.array([5, 6, 7, 0], dtype='i1')
-    >>> euclidean(x, y).compute()
-    6.6332495807108
-
-    >>> euclidean(x, x).compute()
-    0.0
-
+    An ndarray, which contains the result of the calculation of the application
+    of euclidean distance on all the chunks.
     """
-    square_sum = 0.0
-    m = x.shape[0]
-    # Ignore missing values
-    for i in range(m):
-        if x[i] >= 0 and y[i] >= 0:
-            square_sum += (x[i] - y[i]) ** 2
-    out[0] = np.sqrt(square_sum)
+    v = v.sum(axis=0)
+    n = v[5]
+    num = n * v[4] - v[0] * v[1]
+    denom1 = np.sqrt(n * v[2] - v[0] ** 2)
+    denom2 = np.sqrt(n * v[3] - v[1] ** 2)
+    denom = denom1 * denom2
+    value = np.nan
+    if denom > 0:
+        value = 1 - (num / denom)
+    out[0] = value

diff --git a/sgkit/tests/test_distance.py b/sgkit/tests/test_distance.py
index 3432b6421..f5bd44861 100644
--- a/sgkit/tests/test_distance.py
+++ b/sgkit/tests/test_distance.py
@@ -10,7 +10,7 @@
     squareform,
 )

-from sgkit.distance.api import pairwise_distance
+from sgkit.distance.api import MetricTypes, pairwise_distance
 from sgkit.typing import ArrayLike

@@ -98,6 +98,27 @@ def test_distance_euclidean(
     np.testing.assert_almost_equal(distance_matrix, expected_matrix)

+@pytest.mark.parametrize(
+    "size, chunk, split_every, metric",
+    [
+        ((100, 100), (25, 10), 5, "euclidean"),
+        ((100, 100), (20, 25), 3, "euclidean"),
+        ((100, 100), (25, 10), 5, "correlation"),
+        ((100, 100), (20, 25), 3, "correlation"),
+    ],
+)
+def test_pairwise_split_every(
+    size: typing.Tuple[int, int],
+    chunk: typing.Tuple[int, int],
+    split_every: int,
+    metric: MetricTypes,
+) -> None:
+    x = get_vectors(size=size, chunk=chunk)
+    distance_matrix = pairwise_distance(x, metric=metric, split_every=split_every)
+    expected_matrix = squareform(pdist(x, metric=metric))
+    np.testing.assert_almost_equal(distance_matrix, expected_matrix)
+
+
 def test_distance_ndarray() -> None:
     x = get_vectors(array_type="np")
     distance_matrix = pairwise_distance(x, metric="euclidean")
@@ -115,7 +136,7 @@ def test_distance_ndarray() -> None:
     ],
 )
 def test_missing_values(
-    metric: str,
+    metric: MetricTypes,
     metric_func: typing.Callable[[ArrayLike, ArrayLike], np.float64],
     dtype: str,
 ) -> None:
@@ -136,20 +157,31 @@ def test_missing_values(

 @pytest.mark.parametrize(
-    "dtype, expected",
+    "metric, dtype, expected",
     [
-        ("i8", "float64"),
-        ("f4", "float32"),
-        ("f8", "float64"),
+        ("euclidean", "i8", "float64"),
+        ("euclidean", "f4", "float32"),
+        ("euclidean", "f8", "float64"),
+        ("correlation", "i8", "float64"),
+        ("correlation", "f4", "float32"),
+        ("correlation", "f8", "float64"),
     ],
 )
-def test_data_types(dtype, expected):
+def test_data_types(metric: MetricTypes, dtype: str, expected: str) -> None:
     x = get_vectors(dtype=dtype)
-    distance_matrix = pairwise_distance(x)
+    distance_matrix = pairwise_distance(x, metric=metric).compute()
     assert distance_matrix.dtype.name == expected

 def test_undefined_metric() -> None:
     x = get_vectors(array_type="np")
     with pytest.raises(NotImplementedError):
-        pairwise_distance(x, metric="not-implemented-metric")
+        pairwise_distance(x, metric="not-implemented-metric")  # type: ignore[arg-type]
+
+
+def test_wrong_dimension_array() -> None:
+    with pytest.raises(ValueError):
+        pairwise_distance(da.arange(6).reshape(1, 2, 3))
+
+    with pytest.raises(ValueError):
+        pairwise_distance(da.arange(10))

Pairwise distance scalability
Raising this issue to revisit the scalability of our pairwise distance calculation and whether it's worth returning to a map-reduce style implementation that would allow chunking along both dimensions.

In the work that @aktech is doing on early scalability demonstrations (#345) there are some memory usage difficulties emerging. @aktech is I believe trying to run a pairwise distance computation over data from Ag1000G phase 2, using all SNPs and samples from a single chromosome arm. This is of the order ~20 million variants and ~1000 samples. With the current implementation, it is hard to get this to run on systems with average memory/CPU ratios, below 12G/CPU. My understanding is that, essentially, this is because the pairwise distance implementation currently does not support chunking along the variants dimension, and so to reduce memory footprint you need short chunks along the samples dimension. Depending on how the input data have been chunked natively, this may be suboptimal, i.e., you may need to run the computation with sample chunks that are (much) shorter than the native storage.

If this is correct, then it raises two questions for discussion.

First, should we revisit the map-reduce implementation of pairwise distance? This would allow chunking on both samples and variants dimensions, and so could naturally make use of whatever the underlying chunking of the data in storage, without large memory footprint.

Secondly, do we ever really need to run pairwise distance on arrays that are large in the variants dimension? I.e., do we care about scaling this up to large numbers of variants? xref https://github.com/pystatgen/sgkit/pull/306#issuecomment-714654217

> Secondly, do we ever really need to run pairwise distance on arrays that are large in the variants dimension? I.e., do we care about scaling this up to large numbers of variants? xref #306 (comment)

I think you mean in the *samples* dimension? (If so, I  agree - ~10K is as many samples as we could hope to support for O(n^2) like this)
> Secondly, do we ever really need to run pairwise distance on arrays that are large in the variants dimension? I.e., do we care about scaling this up to large numbers of variants? xref [#306 (comment)](https://github.com/pystatgen/sgkit/pull/306#issuecomment-714654217)

On this, it is true that for many analyses it's entirely reasonable to work around memory limitations by downsampling data along the variants dimensions. E.g., often when doing a PCA we randomly downsample variants, and this gives a very reasonable approximation. I.e., typically information about population structure is very clear from just a fraction of the variants.

In other words, the scalability test that @aktech is trying, running pairwise distance over all 20 million SNPs from a chromosome arm, is perhaps somewhat artificial, as you could reasonably downsample to 100,000 SNPs and get the information you need.

On the other side, we do have use cases where we want to compute pairwise distance using all SNPs. E.g., during QC, we check for samples that are technically identical, and check these match our expectations for replicate samples. To do this we compute pairwise distance at all SNPs. (To avoid unnecessary computation, we do an initial pass with a downsampled set of SNPs to exclude pairs of individuals that are obviously not identical, then rerun with all SNPs on remaining samples. But that's a detail.)

> > Secondly, do we ever really need to run pairwise distance on arrays that are large in the variants dimension? I.e., do we care about scaling this up to large numbers of variants? xref #306 (comment)
>
> I think you mean in the _samples_ dimension? (If so, I agree - ~10K is as many samples as we could hope to support for O(n^2) like this)

No, I did mean the variants dimension.

Re the samples dimension, yes agreed, probably ~10k samples is the most we could hope to support due to computational complexity (although if you have access to GPUs then perhaps that allows you to reach a little further).

My main point was about the variants dimension. Computation scales linearly on the variants dimension, there's no computational reason to expect any limit there. We currently have a practical limit due to memory requirements, which could be removed if we supported chunking along the variants dimension via the original map/reduce implementation. Should we revisit that?
I see, thanks. Well, from my perspective I'd find it surprising as a user to be limited on the number of variants we could process by memory restrictions - since it's a sample x sample comparison.

How have you used this in the past? Do you subset down to a representative set of sites, or just throw compute at it and let it run over the whole lot?
> I see, thanks. Well, from my perspective I'd find it surprising as a user to be limited on the number of variants we could process by memory restrictions - since it's a sample x sample comparison.

Yes, exactly.

> How have you used this in the past? Do you subset down to a representative set of sites, or just throw compute at it and let it run over the whole lot?

We've done both, more details [above](https://github.com/pystatgen/sgkit/issues/375#issuecomment-722983199).
Just to add, memory restrictions will be even tighter when moving this to GPU.
> Re the samples dimension, yes agreed, probably ~10k samples is the most we could hope to support due to computational complexity (although if you have access to GPUs then perhaps that allows you to reach a little further).

@alimanfoo I'm curious, I understand that the O(n^2) complexity is vicious, but I'm curious how exactly have you arrived at ~10k (including potential horizontal scaling via dask)?

Also a quick point regarding our current implementation, our current [blockwise](https://github.com/pystatgen/sgkit/blob/master/sgkit/distance/api.py#L91-L101):

```python
    x_distance = da.blockwise(
        # Lambda wraps reshape for broadcast
        lambda _x, _y: metric_ufunc(_x[:, None, :], _y),
        "jk",
        x,
        "ji",
        x,
        "ki",
        dtype="float64",
        concatenate=True,
    )
```

we specify `concatenate=True`, from [dask](https://github.com/dask/dask/blob/master/dask/array/blockwise.py#L78-L83):

> Any index, like ``i`` missing from the output index is interpreted as a contraction (note that this differs from Einstein convention; repeated indices do not imply contraction.)  In the case of a contraction the passed function should expect an iterable of blocks on any array that holds that index.  To receive arrays concatenated along contracted dimensions instead pass ``concatenate=True``.

So afaiu in #345 we are concatenating on the variants blocks, and thus potentially the memory issue(?). @aktech what's the reasoning behind `concatenate=True`? I might have missed it, but have we evaluated an implementation that takes `concatenate=None` (default)?
> @alimanfoo I'm curious, I understand that the O(n^2) complexity is vicious, but I'm curious how exactly have you arrived at ~10k (including potential horizontal scaling via dask)?

I pulled this one out of thin air first. It was just an arbitrary number from experience - in practise, somewhere between 1K and 100K.
> So afaiu in #345 we are concatenating on the variants blocks, and thus potentially the memory issue(?). @aktech what's the reasoning behind concatenate=True? I might have missed it, but have we evaluated an implementation that takes concatenate=None (default)?

@ravwojdyla The `concatenate=True` as you noted is concatenation along variant blocks is made to pass full vectors to the distance metrics functions for distance calculation. If it is `None` the `x` and `y` will be passed to that lamba function as a list of chunks and you cannot calculate distance metric with non-full vectors i.e. chunks, unless you do it via Map-Reduce approach, which we did evaluated, here is the implementation in this [commit](https://github.com/pystatgen/sgkit/blob/e5f0388bf28cf7ade9afc8fcfcfbb785ee228fa3/sgkit/distance/core.py). This was slower than the current implementation, but perhaps seems like more scalable.

To summarise what I think of the current state is it's worth evaluating the map reduce implementation for scalability.
Thanks @aktech, I was curious about the use of `blockwise` with `concatenate=False`, and specifically the memory footprint. The commit you reference doesn't use `blockwise` but instead works directly on blocks. I understand that both might look similar. I was also curious if the `iterable of blocks on any array that holds that index` was lazy or not(?) I can't find a documentation on this, and when I try it on a dummy data, I actually get a `np.ndarray` (which would suggest it isn't lazy), in which case the memory consumption is likely the same for `concatenate` equal `True/False`, which would be rather disappointing (it would still be interesting to double check that tho, plus whether it duplicates the memory required for chunks), I've also noticed @eric-czech [comment](https://github.com/pystatgen/sgkit/pull/306#issuecomment-704294220):

> I'm also thinking that there is a way to express the loop we have now that references `.blocks` as a `blockwise` call like this, though I haven't been able to get it to work yet.

which might suggest we have not evaluated it.

+1 to trying the "map-reduce" implementation on the #345. Going forward iff the performance is important right now, we could also have both implementations that switch based on the shape of the data (and/or parameter), unfortunately that would have a real maintenance cost (so that's something we need to be cautious about).
> Thanks @aktech, I was curious about the use of blockwise with concatenate=False, and specifically the memory footprint. The commit you reference doesn't use blockwise but instead works directly on blocks. I understand that both might look similar. I was also curious if the iterable of blocks on any array that holds that index was lazy or not(?) I can't find a documentation on this, and when I try it on a dummy data, I actually get a np.ndarray (which would suggest it isn't lazy), in which case the memory consumption is likely the same for concatenate equal True/False, which would be rather disappointing (it would still be interesting to double check that tho, plus whether it duplicates the memory required for chunks), I've also noticed @eric-czech comment:

@ravwojdyla Yes, that's very true, I was also expecting the iterable of blocks to be lazy, which would make the memory consumption very low compared to the current implementation.

I only realised this after I wrote a map reduce implementation with `blockwise` yesterday. Here is the map reduce implementation with `blockwise`: https://github.com/aktech/sgkit/blob/pairwise-mr/sgkit/distance/api.py#L53. It might not be the most efficient though (could probably be improved), but gives a rough idea of how it might look like. I tried it yesterday and the memory consumption was not any better (if not worse). So far the the previous map reduce implementation I shared in the commit in my previous comment was better than these two.

Also, one of the reasons that the iterable of blocks are not lazy is probably due to the fact, that the `func` (`blockwise`'s first argument) is not called unless you call `.compute` on it. You would rather expect it would call it before you do `.compute` on it and make an efficient graph. I don't know if that's a feature or a bug, but in most of dask's internal codebase `blockwise` is called with `concatenate=True`, which gives me the impression that they are not scalable in both dimenstions.
@aktech so in general it "seems" like for `blockwise` iff you have a contraction for specific index, all blocks along that index must fit into memory regardless of the `concatenate` option (which might actually double the memory required for those blocks, first to keep the chunks, then to build up the input for the `blockwise` function (?), tho this part likely depends on dask). I wish this was in the documentation, plus the `iterable of blocks` is a bit misleading. I wonder if we should open an issue to add that information, wdyt?
> @aktech so in general it "seems" like for blockwise iff you have a contraction for specific index, all block along that index must fit into memory regardless of the concatenate option

Yes, true.

> (which might actually double the memory required for those blocks, first to keep the chunks, then to build up the input for the blockwise function (?), tho this part likely depends on dask).

What do you mean by "keep the chunks", those chunks are computed once, from the original array to make the input for `blockwise`'s `func`  I believe. The memory consumption would be twice atleast for sure from the chunks from `_x` and `_y`, which are passed to the `func`.

> I wish this was in the documentation, plus the iterable of blocks is a bit misleading. I wonder if we should open an issue to add that information, wdyt?

Yes, I think so. That would be worth mentioning.
> What do you mean by "keep the chunks", those chunks are computed once, from the original array to make the input for blockwise's func I believe. The memory consumption would be twice atleast for sure from the chunks from _x and _y, which are passed to the func.

@aktech yea, per process (threads within process will reuse memory if needed), I just wonder at which stage exactly will Dask free up the memory required for chunks on each process and how exactly will it build up the input for `blockwise`. ~On a dummy data it looks like the chunks are coming in a `ndarray` (not py `list`)~. Also if the cluster has multiple processes and they require the same chunks they will likely fetch that from a process that computed them. So essentially it might be that we would need memory for chunks, \_and\_ memory for concatenated array or an "array of chunks" (depending on the `concatenate` parameter). ~The documentation says that `blockwise` for `concatenate=None` will pass `list` or `iterable` (depending where you look) of blocks, but locally I'm actually seeing a numpy array of blocks being pass into the function given to `blockwise`, which afaiu might requires more memory to build up the array than to pass a list/iterable of materialized/numpy chunks (depending on how soon can dask free up memory of the chunk). I wonder if this is some kind of local optimization or a bug.~ (clarification https://github.com/pystatgen/sgkit/issues/375#issuecomment-724204043)

Edit: but just to be clear, I agree that the `blockwise` implementations do not look at all promising for any computation that has a large number of elements in the contraction index.
>  On a dummy data it looks like the chunks are coming in a ndarray (not py list).

Ah that's interesting, for me that is only in the first iteration on the func you'll see an empty numpy array, if you check on the next iteration it will be a Python list. That execution of the function happens even before you call `.compute`, with empty arrays. Do you see arrays in all iterations?
@aktech I think I found the issue, if there is a contraction and `concatenate=None` the input will be lists (as expected), if `concatenate=True` the input will be numpy array, but also if there is no contraction, the input is numpy arrays (regardless of the `concatenate`).
>  but also if there is no contraction, the input is numpy arrays (regardless of the concatenate).

Interesting, good to know that.

By the way, I discovered something interesting about utilising the symmetry of matrix in dask. We had assumed so far that we can drop duplicate calculation via the following line, this is basically making the lower triangular matrix 0 as the distance matrix is a symmetric matrix.

https://github.com/pystatgen/sgkit/blob/682d58fbb39f91fcf7e441e3d4a937b91123fcb4/sgkit/distance/api.py#L102

I was trying to prove this today, but from what is looks like, our assumption was not true, unless I am missing something. [Here is a notebook](https://nbviewer.jupyter.org/gist/aktech/36e9b4fd04cba99682b2b6179039474a) of that analysis showing the the above line doesn't makes dask utilise the symmetry of a matrix to drop duplicate calculations, in-fact that takes more time.

At the moment I am trying to evaluate the [map-reduce approach](https://github.com/pystatgen/sgkit/blob/e5f0388bf28cf7ade9afc8fcfcfbb785ee228fa3/sgkit/distance/core.py), which does not use the `blockwise`.
FYI: https://stackoverflow.com/questions/64774771/does-blockwise-allow-iteration-over-out-of-core-arrays

cc: @mrocklin, if you get a moment we could use your help understanding the scalability limits of `blockwise`.  A more specific question is framed in the SO post.
Answered.  I also hypothesize in my answer that you're running into
problems with tensordot style applications, which are heavily dependent on
chunking structure.  See https://github.com/dask/dask/issues/2225 .  If
folks here were interested in implementing some of the automatic rechunking
heuristics mentioned in that issue that would be very welcome.

On Wed, Nov 11, 2020 at 8:52 AM Eric Czech <notifications@github.com> wrote:

> FYI:
> https://stackoverflow.com/questions/64774771/does-blockwise-allow-iteration-over-out-of-core-arrays
>
> cc: @mrocklin <https://github.com/mrocklin>, if you get a moment we could
> use your help understanding the scalability limits of blockwise. A more
> specific question is framed in the SO post.
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/pystatgen/sgkit/issues/375#issuecomment-725535583>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTHV4IZIDCSTXUBCNV3SPK6OPANCNFSM4TMOJR4Q>
> .
>

> If folks here were interested in implementing some of the automatic rechunking heuristics mentioned in that issue that would be very welcome.

Thanks @mrocklin.  Would a simplistic example of this be to decrease chunking along non-contracted axes as the contracted axes get larger so that the concatenated chunks still fit in memory?  As a simpler example, say we want to do matrix multiplication via blockwise on 100 x 100 arrays with 10 x 10 chunks.  Now say we want to do the same on 1M x 1M arrays with 10 x 10 chunks -- we now need to fit 1M / 10 chunks in memory instead of 100 / 10 chunks so we could decrease the chunks in the first axes to try to compensate for this.  This is what those heuristics would accomplish correct?
If I recall correctly the heuristic was pretty unintuitive.  I think that
one wanted intermediate result chunks that were as large as possible
because this cut down significantly on how much communication and
duplication of the input arrays that you had to do.  I don't fully recall
though.

My recollection is that tensordot computations end up blowing up memory
more often because of the all-to-all nature of the computation.  The input
dataset ends up being copied onto distributed memory many times when one
doesn't do things carefully.

On Wed, Nov 11, 2020 at 10:36 AM Eric Czech <notifications@github.com>
wrote:

> If folks here were interested in implementing some of the automatic
> rechunking heuristics mentioned in that issue that would be very welcome.
>
> Thanks @mrocklin <https://github.com/mrocklin>. Would a simplistic
> example of this be to decrease chunking along non-contracted axes as the
> contracted axes get larger so that the concatenated chunks still fit in
> memory? As a simpler example, say we want to do matrix multiplication via
> blockwise on 100 x 100 arrays with 10 x 10 chunks. Now say we want to do
> the same on 1M x 1M arrays with 10 x 10 chunks -- we now need to fit 1M /
> 10 chunks in memory instead of 100 / 10 chunks so we could decrease the
> chunks in the first axes to try to compensate for this. This is what those
> heuristics would accomplish correct?
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/pystatgen/sgkit/issues/375#issuecomment-725591045>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTDR5T3JKTSNXNADYWLSPLKUPANCNFSM4TMOJR4Q>
> .
>

I ran the map-reduce pairwise implementation on the MalariaGEN dataset on [Coiled Cloud](https://coiled.io/) dask cluster with the following configuration:

- Workers: 25
- Memory: 16 GB each
- CPU: 4
- Nthreads: 1

### Code:
- Implementation in [branch:pairwise-mr-blocks](https://github.com/aktech/sgkit/blob/pairwise-mr-blocks/sgkit/distance/api.py) (this includes missing value handling, similar to current implementation)
- Notebook used is [here](https://nbviewer.jupyter.org/gist/aktech/7e9011e4194a9cb38cb4ecbe4b3f7232)
- Metric used: `euclidean`

---
Here are the results:

### 1. Chunk size: 100 MB

- Time taken: 32 min 59 sec (1978.87 s)
- [Dask Report](https://rawcdn.githack.com/aktech/aktech.github.io/3a98bdb0f511b773f6a642a5dd5c6b314c87eaaa/sgkit/pairwise-analysis/dask-report-25W-16G-4cpu-1thread.html)
- [Task Stream](https://rawcdn.githack.com/aktech/aktech.github.io/3a98bdb0f511b773f6a642a5dd5c6b314c87eaaa/sgkit/pairwise-analysis/task-stream-25W-16G-4cpu-1threads.html)

Takeaway: Everything is quick except the `euclidean_map` task, which is the bottleneck as you can see from the report above, which makes sense because that is something non-parallelizable from a dask point of view. The chunk size is 100 MB, which means a 100 MB chunk is passed to guvectorized `euclidean_map`, which runs serially.

### 2. (a) Chunk Size 100 MB, using Numba's `target="parallel"`
- Time Taken: 11 min 20 sec (680.34 s)
- [Dask Report](https://rawcdn.githack.com/aktech/aktech.github.io/3a98bdb0f511b773f6a642a5dd5c6b314c87eaaa/sgkit/pairwise-analysis/dask-report-25W-16G-4cpu-1threads-parallel.html)
- [Task Stream](https://rawcdn.githack.com/aktech/aktech.github.io/3a98bdb0f511b773f6a642a5dd5c6b314c87eaaa/sgkit/pairwise-analysis/task-stream-25W-16G-4cpu-1threads-parallel.html)

Takeaway: Since, the chunk is non parallelizable from dask point of view as it is dispatched to `guvectorize`, one way to speed things is to use Numba’s `target="parallel"`, which would parallelize the individual chunk.

One thing to note about this approach is ideally dask’s nthreads should be 1 in this case, because Dask and Numba don’t share their threadpool, which might cause contention, see [this issue](https://github.com/dask/dask/issues/5229#issuecomment-518822534)

### 2. (b) Chunk Size 50 MB, using Numba's `target="parallel"`

- Time Taken: 12 min 14 sec (733.56 s)
- [Dask Report](https://rawcdn.githack.com/aktech/aktech.github.io/3a98bdb0f511b773f6a642a5dd5c6b314c87eaaa/sgkit/pairwise-analysis/dask-report-25W-16G-4cpu-1threads-parallel-50mb-chunks-parallel.html)
[Task Stream](https://rawcdn.githack.com/aktech/aktech.github.io/3a98bdb0f511b773f6a642a5dd5c6b314c87eaaa/sgkit/pairwise-analysis/task-stream-25W-16G-4cpu-1threads-parallel-50mb-chunks-parallel.html)

Since the bottleneck in all these cases is the map function, any kind of parallelisation can speed things up, another approach could be parallelising the chunk with Dask instead of Numba, by getting rid of broadcasting and parallelising calls to guvectorize with two 1D arrays via Dask.

---
As mentioned earlier, I was not able to run the `blockwise` implementation due to memory issues, and the above map-reduce implementation seems to scale well.

## Update:

@eric-czech @alimanfoo

We are able to achieve similar speeds of numba `target="parallel"` by increasing number of threads in above dask workers, as suggested by Eric:

### 3. Same configuration as "1" above with `nthreads: 4`

- Time taken:  10 min 48 sec (648.08 s)
- [Dask Report](https://rawcdn.githack.com/aktech/aktech.github.io/0ab56162139c9614c2b14e4f110c3fe81f729c31/sgkit/pairwise-analysis/dask-report-16G-4cpu-4threads.html)
- [Task Stream](https://rawcdn.githack.com/aktech/aktech.github.io/0ab56162139c9614c2b14e4f110c3fe81f729c31/sgkit/pairwise-analysis/task-stream-16G-4cpu-4threads.html)
@eric-czech @alimanfoo

I was comparing it with scikit-allel as suggested by Eric in the dev call today. It doesn't seems to work with even a small array of the MalariaGEN data. Here is the [notebook for the same](https://nbviewer.jupyter.org/gist/aktech/d6c952599d0f3b1cb7bfbc7ff49f63a0).
Hi @aktech

> I was comparing it with scikit-allel as suggested by Eric in the dev call today. It doesn't seems to work with even a small array of the MalariaGEN data. Here is the [notebook for the same](https://nbviewer.jupyter.org/gist/aktech/d6c952599d0f3b1cb7bfbc7ff49f63a0).

Looks like you need a petabyte of memory :-)

Perhaps Eric was suggesting to compare with the [scikit-allel v2 prototype](https://github.com/scikit-allel/skallel-stats)? The pairwise distance function in scikit-allel v1 is just scipy pdist, not dask enabled.

> Perhaps Eric was suggesting to compare with the scikit-allel v2 prototype? The pairwise distance function in scikit-allel v1 is just scipy pdist, not dask enabled.

Yes, [that's the one](https://github.com/scikit-allel/skallel-stats/blob/master/src/skallel_stats/distance/api.py#L5) I tried. Here is the [environment](https://gist.github.com/aktech/dec74d199feacb18d12e8501ef028923) (Note that, I had to remove pinning from libraries to resolve conflicts) I used for workers. I also faced [this issue](https://github.com/scikit-allel/skallel-stats/issues/18), which was fixed by upgrading `llvmlite`.

Although, I just realised that code was calling numpy backend for some reason, even though a dask array was passed to it. I will play with it again to see why.
> Yes, [that's the one](https://github.com/scikit-allel/skallel-stats/blob/master/src/skallel_stats/distance/api.py#L5) I tried.

Sorry, I should've looked more carefully.

I think the problem is transposing the input array. The skallel-stats pairwise distance function computes pairwise distance between columns, not rows. This is because the genotype data we have is typically arranged this way (columns are samples).
Just to add I have run the skallel pairwise distance on a simulated dataset approaching the size of Ag1000G: https://github.com/scikit-allel/skallel-stats/blob/master/examples/pairwise_distance.ipynb

Note that I created the simulated zarr dataset with chunking best for pairwise distance, i.e., no chunks along the samples dimension.

I just tried running on Ag1000G on malariagen datalab and it's a little awkward because the chunks are quite tall (~500,000) and skallel rechunks the data to remove any chunks in the samples dimension, so memory is a challenge. I can get it to run if I use dask rechunk on the input data to make chunks shorter, which is then suboptimal for I/O (because the same chunk is being read multiple times) but does work. But it doesn't seem to be running particularly well.

Perhaps try a simulated dataset for now?
> I just tried running on Ag1000G on malariagen datalab and it's a little awkward because the chunks are quite tall (~500,000) and skallel rechunks the data to remove any chunks in the samples dimension, so memory is a challenge. I can get it to run if I use dask rechunk on the input data to make chunks shorter, which is then suboptimal for I/O (because the same chunk is being read multiple times) but does work. But it doesn't seem to be running particularly well.

I also had hard time trying to run it on the MalariaGEN data, couldn't get it working.

> Perhaps try a simulated dataset for now?

I ran it on a simulated dataset, similar to the size of MalariaGEN. It seems to run very well, here is the [notebook](https://nbviewer.jupyter.org/gist/aktech/6621da206b436beea8baf707137efff3) for the same. It took about **5min 32s**, on 25 workers (same configuration as above).

> I ran it on a simulated dataset, similar to the size of MalariaGEN. It seems to run very well, here is the notebook for the same. It took about 5min 32s, on 25 workers (same configuration as above).

Also, to compare with the map-reduce implementation the same takes 18 minutes roughly, whereas it takes ~10 minutes on the MalariaGEN data. Also, note that the MalariaGEN data and simulated data are of nearly the same size (~28.5 GB)

~It seems generating random numbers is more expensive than reading from GCS in the map-reduce implementation.~
Update: (This sentence is incorrect, the slow speed was due to the chunking shape not due to random number generation)

@alimanfoo @eric-czech any thoughts on this? Is that fast enough?
So, just to clarify, we're looking at about a 4X slow-down for the map-reduce implementation? This seems OK, if it allows us to scale.
Hi @aktech, echoing @jeromekelleher, I would say that speed per se is less important than scalability. I.e., whether it takes 5 minutes or 18 minutes is less important than being able to run this robustly within reasonable memory limits, and demonstrating that increasing the cluster size makes the whole computation complete faster (up to some natural limit such as the number of chunks in the data).

> So, just to clarify, we're looking at about a 4X slow-down for the map-reduce implementation? This seems OK, if it allows us to scale.

That is true for simulated data with random values, although on the actual data its ~10 mins, which is 2X slow, assuming the skallel implementation is able to run on the MalariaGEN data with same performance as it runs on the simulated data.

> Also, to compare with the map-reduce implementation the same takes 18 minutes roughly.

So, this is reduced to 10 minutes with better chunking.

@alimanfoo I think it is able to scale well and memory hasn't been a problem from what I have seen from various experiments, I did on Map-reduce implementation (if each worker has 8GB memory allowance, which seems reasonable to me).

The chunk size was 100MB in all of these cases.

### 1. Macbook Pro 16 GB (2 workers, 4threads)

I am able to run the simulated data of same size as MalariaGEN on my laptop as well, here is the [notebook](https://nbviewer.jupyter.org/github/aktech/aktech.github.io/blob/master/work/sgkit/notebooks/sgkit_local_full.ipynb). Following are the links to reports:

- [Dask Performance report](https://iamit.in/work/sgkit/pairwise-analysis/dask-report-16G-4cpu-4threads-local-full.html)
- [Dask Task Stream](https://iamit.in/work/sgkit/pairwise-analysis/task-stream-16G-4cpu-4threads-local-full.html)

It took about **8 hours 19 minutes** overnight, given the fact that there would be some background tasks running on my laptop too and I also suspect, it might have slept and paused the workflow at night, but the point is I was able to run it on my laptop.

I chose two workers to give 8GB memory to each worker. I didn't had any memory issues. 8GB per worker seems like a good number for memory allowance.

### 2. Coiled cluster: 24 Worker (16 GB, 4 Cores each): MalariaGEN data

- [Dask Performance report](https://iamit.in/work/sgkit/pairwise-analysis/dask-report-16G-4cpu-4threads.html)
- [Dask Task Stream](https://iamit.in/work/sgkit/pairwise-analysis/task-stream-16G-4cpu-4threads.html)

Time Taken: **10 min 48 sec**

### 3. Coiled cluster: 50 Worker (16 GB, 4 Cores each): MalariaGEN data

- [Dask Performance report](https://iamit.in/work/sgkit/pairwise-analysis/dask-report-16G-4cpu-4threads-50W.html)
- [Dask Task Stream](https://iamit.in/work/sgkit/pairwise-analysis/task-stream-16G-4cpu-4threads-50W.html)

Time Taken: **5 min 54 sec**

Let me know if this sounds reasonable.

---

Also, thanks to @mrocklin's @coiled for generously increasing my cores limit to 200.

> Also, thanks to @mrocklin <https://github.com/mrocklin>'s

I don't remember doing this, but I'm glad that it happened regardless.  If
you all want an team account with more capacity then do let me know.

On Thu, Nov 19, 2020 at 1:44 AM Amit Kumar <notifications@github.com> wrote:

> Also, to compare with the map-reduce implementation the same takes 18
> minutes roughly.
>
> So, this is reduced to 10 minutes with better chunking.
>
> @alimanfoo <https://github.com/alimanfoo> I think it is able to scale
> well and memory hasn't been a problem from what I have seen from various
> experiments, I did on Map-reduce implementation (if each worker has 8GB
> memory allowance, which seems reasonable to me).
>
> The chunk size was 100MB in all of these cases.
> 1. Macbook Pro 16 GB (2 workers, 4threads)
>
> I am able to run the simulated data of same size as MalariaGEN on my
> laptop as well, here is the notebook
> <https://nbviewer.jupyter.org/github/aktech/aktech.github.io/blob/master/work/sgkit/notebooks/sgkit_local_full.ipynb>.
> Following are the links to reports:
>
>    - Dask Performance report
>    <https://iamit.in/work/sgkit/pairwise-analysis/dask-report-16G-4cpu-4threads-local-full.html>
>    - Dask Task Stream
>    <https://iamit.in/work/sgkit/pairwise-analysis/task-stream-16G-4cpu-4threads-local-full.html>
>
> It took about *8 hours 19 minutes* overnight, given the fact that there
> would be some background tasks running on my laptop too and I also suspect,
> it might have slept and paused the workflow at night, but the point is I
> was able to run it on my laptop.
>
> I chose two workers to give 8GB memory to each worker. I didn't had any
> memory issues. 8GB per worker seems like a good number for memory allowance.
> 2. Coiled cluster: 24 Worker (16 GB, 4 Cores each): MalariaGEN data
>
>    - Dask Performance report
>    <https://iamit.in/work/sgkit/pairwise-analysis/dask-report-16G-4cpu-4threads.html>
>    - Dask Task Stream
>    <https://iamit.in/work/sgkit/pairwise-analysis/task-stream-16G-4cpu-4threads.html>
>
> Time Taken: *10 min 48 sec*
> 3. Coiled cluster: 50 Worker (16 GB, 4 Cores each): MalariaGEN data
>
>    - Dask Performance report
>    <https://iamit.in/work/sgkit/pairwise-analysis/dask-report-16G-4cpu-4threads-50W.html>
>    - Dask Task Stream
>    <https://iamit.in/work/sgkit/pairwise-analysis/task-stream-16G-4cpu-4threads-50W.html>
>
> Time Taken: *5 min 54 sec*
>
> Let me know if this sounds reasonable.
> ------------------------------
>
> Also, thanks to @mrocklin <https://github.com/mrocklin>'s @coiled
> <https://github.com/coiled> for generously increasing my cores limit to
> 200.
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/pystatgen/sgkit/issues/375#issuecomment-730253208>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTCB54ENY4M4CMDR4B3SQTSGFANCNFSM4TMOJR4Q>
> .
>

Looking at the performance reports it seems like you're definitely not dask-communication bound, so I would not expect decreasing chunking to improve performance.  This aligns with @aktech 's intutition from the meeting.
> > Also, thanks to @mrocklin

> I don't remember doing this, but I'm glad that it happened regardless. If you all want an team account with more capacity then do let me know.

That was me :smile: but yes, if more capacity is needed either of us is happy to help!
Hey @mrocklin looking at [matmul](https://github.com/dask/dask/blob/ff3ea6c74e8736a5823d648defcf9164f083d266/dask/array/routines.py#L313), it looks like it must have contracted axes .. right?  It would be great to know if it's possible to implement matrix multiplication using reductions and no contractions instead since we could easily extend that to our more general pairwise metrics.
I am surprised to see that matmul does not just reuse the tensordot
implementation, or at least use the same trick used in tensordot.  That
trick is useful and definitely results in smoother computations.

On Thu, Nov 19, 2020 at 2:41 PM Eric Czech <notifications@github.com> wrote:

> Hey @mrocklin <https://github.com/mrocklin> looking at matmul
> <https://github.com/dask/dask/blob/ff3ea6c74e8736a5823d648defcf9164f083d266/dask/array/routines.py#L313>,
> it looks like it must have contracted axes .. right? It would be great to
> know if it's possible to implement matrix multiplication using reductions
> and no contractions instead since we could easily extend that to our more
> general pairwise metrics.
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/pystatgen/sgkit/issues/375#issuecomment-730682432>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTDTOAKWTCMCDDXIYZ3SQWNLHANCNFSM4TMOJR4Q>
> .
>

> It would be great to know if it's possible to implement matrix multiplication using reductions and no contractions instead since we could easily extend that to our more general pairwise metrics.

Could we try replacing calls to `x @ y` (in the GWAS code) with something like `da.tensordot(x, y, (1, 0))` to see if that has any effect?
Good idea @tomwhite, I tried swapping that in the notebook mentioned at https://github.com/pystatgen/sgkit/issues/390#issuecomment-730449672 and saw this on small simulated arrays:

<img width="973" alt="Screen Shot 2020-11-20 at 6 38 12 AM" src="https://user-images.githubusercontent.com/6130352/99796384-65eb2800-2afb-11eb-91c7-524865ffd579.png">

Do you think it is worth starting a conversation as a dask issue about changing the matmul implementation @mrocklin?

Btw in the meantime for 2d input we could also use `dot`: `x @ y -> da.dot(x, y)` which in dask under the hood uses `tensordot` (and should give the same results as `matmul` for 2d). +1 to starting the discussion about `matmul` using the same trick of blockwise product followed by reduction, and maybe potentially having a common code for `tensordot`, `matmul`, `dot`, `einsum` (it looks like it's also using this trick btw with it's own implementation).
> Do you think it is worth starting a conversation as a dask issue about
changing the matmul implementation @mrocklin <https://github.com/mrocklin>?

Definitely

On Fri, Nov 20, 2020 at 4:41 AM Rafal Wojdyla <notifications@github.com>
wrote:

> Btw in the meantime for 2d input we could also use dot which in dask
> under the hood uses tensordot (and should give the same results as matmul
> for 2d). +1 to starting the discussion about matmul using the same trick
> of blockwise product followed by reduction, and maybe potentially having a
> common code for tensordot, matmul, dot, einsum (it looks like it's also
> using this trick btw).
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/pystatgen/sgkit/issues/375#issuecomment-731147873>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTGCMENG3DBQ2NLTWDTSQZPX5ANCNFSM4TMOJR4Q>
> .
>

Also, nice find

On Fri, Nov 20, 2020 at 8:35 AM Matthew Rocklin <mrocklin@gmail.com> wrote:

> > Do you think it is worth starting a conversation as a dask issue about
> changing the matmul implementation @mrocklin <https://github.com/mrocklin>
> ?
>
> Definitely
>
> On Fri, Nov 20, 2020 at 4:41 AM Rafal Wojdyla <notifications@github.com>
> wrote:
>
>> Btw in the meantime for 2d input we could also use dot which in dask
>> under the hood uses tensordot (and should give the same results as matmul
>> for 2d). +1 to starting the discussion about matmul using the same trick
>> of blockwise product followed by reduction, and maybe potentially having a
>> common code for tensordot, matmul, dot, einsum (it looks like it's also
>> using this trick btw).
>>
>> —
>> You are receiving this because you were mentioned.
>> Reply to this email directly, view it on GitHub
>> <https://github.com/pystatgen/sgkit/issues/375#issuecomment-731147873>,
>> or unsubscribe
>> <https://github.com/notifications/unsubscribe-auth/AACKZTGCMENG3DBQ2NLTWDTSQZPX5ANCNFSM4TMOJR4Q>
>> .
>>
>

I have implemented a blockwise version of pairwise in [this branch](https://github.com/aktech/sgkit/tree/pairwise-new-blockwise) (based on @ravwojdyla’s implementation of matmul in this PR: https://github.com/dask/dask/pull/7000).

The blockwise implementation doesn’t seems to be any better than the [blocks implementation](https://github.com/aktech/sgkit/tree/pairwise-mr-blocks-nop) (the one for which I shared the results above). I ran some benchmarks to compare them, here is the comparison of memory and CPU usage between the two implementations:

<img width="771" alt="Screenshot 2021-01-21 at 14 53 13" src="https://user-images.githubusercontent.com/5647941/105367608-8701ef00-5bf8-11eb-86f4-15cc2ec57d34.png">

### Data:
- Random data of size about 10% of malariaGEN: 2.83 GB,
- 50 MB chunks
- Machine: 48 cores and 64GB memory

In general the blockwise implementation was a bit slower than the blocks implementation, it took the following times:

- Without blockwise: **6min 24s**
- With blockwise: **10min 42s**

Example notebook for this comparison is [here](https://nbviewer.jupyter.org/gist/aktech/072e6a5baa6cfb294fa08e8290b8372a).
@aktech reviewed this [version](https://github.com/aktech/sgkit/tree/cf36c0e784d71bd94fadd4d9268cbeaa987ddc03/sgkit/distance), I believe the blockwise implementation can perform as well as the bespoke blocks implementation, here are some comments:
 * what's the need for the extra `l` dimension [here](https://github.com/aktech/sgkit/blob/cf36c0e784d71bd94fadd4d9268cbeaa987ddc03/sgkit/distance/api.py#L98-L104)?
 * [this](https://github.com/aktech/sgkit/blob/cf36c0e784d71bd94fadd4d9268cbeaa987ddc03/sgkit/distance/api.py#L110-L112) should be implemented as a `da.reduction`

Furthermore (orthogonal to the blockwise investigation):
 * instead of summing two separate `da.triu`, would it be possible to reuse a single result of `da.triu`? (and please use the same trick in both blocks and blockwise)
 * what's the advantage of having `euclidean_reduce` `guvectorize`d?
 * there seem to be a substantial overhead of handling "missing" (negative) values in the current `euclidean_map`, *once* we choose how we orchestrate blocks (blocks vs blockwise etc), we should research if there is a way to optimize it

Btw for future reference:
 * having the performance [notebook](https://nbviewer.jupyter.org/gist/aktech/072e6a5baa6cfb294fa08e8290b8372a) collocated with the code that is required (extra methods etc) and producing the graphs would make things easier to review, I can only assume that you have used this [branch](https://github.com/aktech/sgkit/blob/8ac33501d2ca3c01d254fe0f33fb9f55554a1689/sgkit/distance/api.py) as the "environment" for that notebook?
 * having performance reports linked would help too

Don't hesitate to ping me if something is not clear or problematic.
@ravwojdyla
> what's the need for the extra l dimension here?

When a pair of partial vectors are evaluated for the map step of the pairwise for `correlation` metric, it gives an array of size 6, to capture that in the output of the blockwise and be able to reshape it we need to add that dimension (`l` in this case) which is a function of the metric.

To elaborate on this based on my understanding of blockwise, when we perform blockwise the actual output dimension is not known to dask, since it would never know what the function is doing to the dimension of the output. The only way for blockwise to know its dimension is via its arguments, which is `out_ind` (and `adjust_chunks`), I call that as the anticipated dimension of blockwise’s output. The anticipated dimension is not always same as actual dimension, so if you perform a reshape on the output, it has to be valid w.r.t the anticipated dimension (irrespective of the fact that it's valid for the actual dimension). This makes it important to make the anticipated dimension suitable for reshaping in this case, hence adding the 4th dimension ‘l’, which makes the output dimension a function of nparams.

Also, note that the `np.new_axis` is not really required, in this case it might just work without that, I added that to be able to do `out.compute()` on blockwise’s output prior to reshaping to see what blockwise returned.

Let me know if this doesn’t makes sense or seems inaccurate.

> this should be implemented as a da.reduction

I see, I have never used that. I’ll have a look, do you reckon it will make it more efficient? Or is that a standard way of doing reduction in dask?

> instead of summing two separate da.triu, would it be possible to reuse a single result of da.triu? (and please use the same trick in both blocks and blockwise)

I think so, I'll give it a go.

> what's the advantage of having euclidean_reduce guvectorized?

Not much really, I think it started with having some for loops initially maybe, then its was never removed.

> Btw for future reference:

Yes, sure. Apologies for the missing information. I used this [branch](https://github.com/aktech/sgkit/tree/pairwise-comparison) for comparison. I have now updated the [notebook](https://gist.github.com/aktech/072e6a5baa6cfb294fa08e8290b8372a) to reflect the actual experiment.

> Let me know if this doesn’t makes sense or seems inaccurate.

Ah, I see, it's for the `correlation` metric, ok, if we decide to use blockwise, it should be probably explicitly stated. Thanks for the clarification tho.

>>this should be implemented as a da.reduction

> I see, I have never used that. I’ll have a look, do you reckon it will make it more efficient? Or is that a standard way of doing reduction in dask?

Yes, here's a *sketch* of the `euclidean` distance variant:

```python
    ...
    o = da.blockwise(
        _pairwise,
        'ijk',
        x,
        'ik',
        x,
        'jk',
        adjust_chunks={'k': 1},
        dtype=x.dtype,
        concatenate=False,
        h=np.empty(n_map_param, dtype=x.dtype),
    )

    r = da.reduction(o,
                     chunk=lambda x_chunk, axis, keepdims: x_chunk,
                     combine=lambda x_chunk, axis, keepdims: x_chunk.sum(-1)[..., np.newaxis],
                     aggregate=lambda x_chunk, axis, keepdims: metric_reduce_ufunc(x_chunk),
                     axis=-1,
                     dtype=np.float,
                     name="pairwise")
    t = da.triu(r)
    return t + t.T
```

On my laptop, this performs as well as the custom blocks. Here it's also worth mentioned that `da.reduction` has a handy parameter `split_every`, that could potentially be exposed on the `pairwise_distance` method.

>> what's the advantage of having euclidean_reduce guvectorized?

>Not much really, I think it started with having some for loops initially maybe, then its was never removed.

If there is not need for it, it should be removed?

> Yes, sure. Apologies for the missing information. I used this branch for comparison. I have now updated the notebook to reflect the actual experiment.

Having notebook as GH gist isn't great, because it's hard to rerun it or track it, I assume the gist should be run using the code from the branch, but it would be a lot more clear if the notebook just lived in the branch for the time of the comparison, so it's easy to just clone it and run it or change stuff. Does that make sense @aktech?
> On my laptop, this performs as well as the custom blocks. Here it's also worth mentioned that da.reduction has a handy parameter split_every, that could potentially be exposed on the pairwise_distance method.

Oh, nice! Thanks, let me do the same for correlation too.

> If there is not need for it, it should be removed?

Yes, definitely.

> Having notebook as GH gist isn't great, because it's hard to rerun it or track it, I assume the gist should be run using the code from the branch, but it would be a lot more clear if the notebook just lived in the branch for the time of the comparison, so it's easy to just clone it and run it or change stuff. Does that make sense @aktech?

Ah, I see what you mean. I did that because many times the IPython notebooks in repository are not loaded at all, so if people just want to see the results there isn't a way and nbviewer doesn't seems to load from a repo link. Although I completely agree with your point it makes review difficult. I'll keep that in mind for future, thanks!
@ravwojdyla I have now used the `da.reduction` approach suggested above to compute the pairwise distance in blockwise. I did a quick comparison by running both locally on about 1 GB data and they are performing nearly the same,
roughly:

- blocks implementation: **9 min 5 sec**
- blockwise implementation: **9 min 42 sec**

I'll run it again on a little larger dataset to compare and will share a notebook thereafter, meanwhile the implementation is in [this branch](https://github.com/aktech/sgkit/tree/pairwise-comparison). I have also removed `guvectorize` from `euclidean_reduce`.