zmc
2023-08-08 e792e9a60d958b93aef96050644f369feb25d61b
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
from __future__ import annotations
 
import numpy as np
 
from pandas._typing import NumpyIndexT
 
from pandas.core.dtypes.common import is_list_like
 
 
def cartesian_product(X) -> list[np.ndarray]:
    """
    Numpy version of itertools.product.
    Sometimes faster (for large inputs)...
 
    Parameters
    ----------
    X : list-like of list-likes
 
    Returns
    -------
    product : list of ndarrays
 
    Examples
    --------
    >>> cartesian_product([list('ABC'), [1, 2]])
    [array(['A', 'A', 'B', 'B', 'C', 'C'], dtype='<U1'), array([1, 2, 1, 2, 1, 2])]
 
    See Also
    --------
    itertools.product : Cartesian product of input iterables.  Equivalent to
        nested for-loops.
    """
    msg = "Input must be a list-like of list-likes"
    if not is_list_like(X):
        raise TypeError(msg)
    for x in X:
        if not is_list_like(x):
            raise TypeError(msg)
 
    if len(X) == 0:
        return []
 
    lenX = np.fromiter((len(x) for x in X), dtype=np.intp)
    cumprodX = np.cumprod(lenX)
 
    if np.any(cumprodX < 0):
        raise ValueError("Product space too large to allocate arrays!")
 
    a = np.roll(cumprodX, 1)
    a[0] = 1
 
    if cumprodX[-1] != 0:
        b = cumprodX[-1] / cumprodX
    else:
        # if any factor is empty, the cartesian product is empty
        b = np.zeros_like(cumprodX)
 
    # error: Argument of type "int_" cannot be assigned to parameter "num" of
    # type "int" in function "tile_compat"
    return [
        tile_compat(
            np.repeat(x, b[i]),
            np.prod(a[i]),  # pyright: ignore[reportGeneralTypeIssues]
        )
        for i, x in enumerate(X)
    ]
 
 
def tile_compat(arr: NumpyIndexT, num: int) -> NumpyIndexT:
    """
    Index compat for np.tile.
 
    Notes
    -----
    Does not support multi-dimensional `num`.
    """
    if isinstance(arr, np.ndarray):
        return np.tile(arr, num)
 
    # Otherwise we have an Index
    taker = np.tile(np.arange(len(arr)), num)
    return arr.take(taker)