Consider clusters of segments with similar naive averages as segment candidates in their own right

wise_pizza/cluster.py

@@ -0,0 +1,56 @@
+import numpy as np
+import pandas as pd
+from sklearn.preprocessing import PowerTransformer
+from sklearn.cluster import KMeans, kmeans_plusplus
+from sklearn.metrics import silhouette_score
+
+
+def guided_kmeans(X: np.ndarray, power_transform: bool = True) -> np.ndarray:
+    """
+    Cluster segment averages to calculate aggregated segments
+    @param X: Segment mean minus global mean, for each dimension value
+    @param power_transform: Do we power transform before clustering
+    @return: cluster labels and the transformed values
+    """
+    if isinstance(X, pd.Series):
+        X = X.values.reshape(-1, 1)
+    elif isinstance(X, pd.DataFrame):
+        X = X.values
+
+    if power_transform:
+        if len(X[X > 0] > 1):
+            X[X > 0] = (
+                PowerTransformer(standardize=False)
+                .fit_transform(X[X > 0].reshape(-1, 1))
+                .reshape(-1)
+            )
+        if len(X[X < 0] > 1):
+            X[X < 0] = (
+                -PowerTransformer(standardize=False)
+                .fit_transform(-X[X < 0].reshape(-1, 1))
+                .reshape(-1)
+            )
+
+    best_score = -1
+    best_labels = None
+    # If we allow 2 clusters, it almost always just splits positive vs negative - boring!
+    for n_clusters in range(3, 10):
+        cluster_labels = KMeans(
+            n_clusters=n_clusters, init="k-means++", n_init=10
+        ).fit_predict(X)
+        score = silhouette_score(X, cluster_labels)
+        print(n_clusters, score)
+        if score > best_score:
+            best_score = score
+            best_labels = cluster_labels
+            best_n = n_clusters
+
+    print(best_n)
+    return best_labels, X
+
+
+def to_matrix(labels: np.ndarray) -> np.ndarray:
+    out = np.zeros((len(labels), len(labels.unique())))
+    for i in labels.unique():
+        out[labels == i, i] = 1.0
+    return out

wise_pizza/explain.py

@@ -200,14 +200,12 @@ def explain_changes_in_totals(
         sf_size.final_size = final_size
         sf_avg.final_size = final_size
         sp = SlicerPair(sf_size, sf_avg)
-        sp.plot = (
-            lambda plot_is_static=False, width=2000, height=500: plot_split_segments(
-                sp.s1,
-                sp.s2,
-                plot_is_static=plot_is_static,
-                width=width,
-                height=height,
-            )
+        sp.plot = lambda plot_is_static=False, width=2000, height=500: plot_split_segments(
+            sp.s1,
+            sp.s2,
+            plot_is_static=plot_is_static,
+            width=width,
+            height=height,
         )
         return sp
 
@@ -266,8 +264,7 @@ def explain_levels(
     @param solver: If this equals to "lp" uses the LP solver, else uses the (recommended) Lasso solver
     @param verbose: If set to a truish value, lots of debug info is printed to console
     @param force_add_up: Force the contributions of chosen segments to add up to zero
-    @param constrain_signs: Whether to constrain weights of segments to have the same
-    sign as naive segment averages
+    @param constrain_signs: Whether to constrain weights of segments to have the same sign as naive segment averages
     @return: A fitted object
     """
     df = copy.copy(df)

wise_pizza/make_matrix.py

@@ -1,5 +1,5 @@
 import itertools
-from typing import Optional, List, Dict
+from typing import Optional, List, Dict, Sequence
 
 import numpy as np
 import scipy
@@ -133,9 +133,42 @@ def sparse_dummy_matrix(
             else:
                 used_dims = [force_dim] + list(these_dims)
 
-            these_defs = segment_defs(dim_df, used_dims, verbose=verbose)
-            this_mat = construct_dummies(these_defs, dummy_cache)
+            segment_constraints = segment_defs_new(dims_dict, used_dims)
+            this_mat, these_defs = construct_dummies_new(used_dims, segment_constraints, dummy_cache)
+
+            # these_defs = segment_defs(dim_df, used_dims, verbose=verbose)
+            # this_mat = construct_dummies(these_defs, dummy_cache)
             mats.append(this_mat)
             defs += these_defs
     mat = hstack(mats)
     return mat, defs
+
+
+def segment_defs_new(dims_dict: Dict[str, Sequence[str]], used_dims) -> List[Dict[str, str]]:
+    if len(used_dims) == 1:
+        return np.array(dims_dict[used_dims[0]]).reshape(-1, 1)
+    else:
+        tmp = segment_defs_new(dims_dict, used_dims[:-1])
+        this_dim_values = np.array(dims_dict[used_dims[-1]])
+        repeated_values = np.tile(this_dim_values, len(tmp)).reshape(-1, 1)
+        pre_out = np.tile(tmp, (len(this_dim_values), 1))
+        out = np.concatenate(pre_out, repeated_values)
+        return out
+
+
+def construct_dummies_new(
+    used_dims: List[str], segment_defs: np.ndarray, cache: Dict[str, Dict[str, np.ndarray]]
+) -> scipy.sparse.csc_matrix:
+    dummies = []
+    segments = []
+    for sgdf in segment_defs:
+        tmp = None
+        for i, d in enumerate(used_dims):
+            if tmp is None:
+                tmp = cache[d][sgdf[i]]
+            else:
+                tmp = tmp.multiply(cache[d][sgdf[i]])
+        if tmp.sum() > 0:
+            dummies.append(tmp)
+            segments.append(dict(zip(used_dims, sgdf)))
+    return hstack(dummies), segments