# Buggy code
def _subdiagrams(X, homology_dimensions, remove_dim=False):
    """For each diagram in a collection, extract the subdiagrams in a given
    list of homology dimensions. It is assumed that all diagrams in X contain
    the same number of points in each homology dimension."""
    n = len(X)
    if len(homology_dimensions) == 1:
        Xs = X[X[:, :, 2] == homology_dimensions[0]].reshape(n, -1, 3)
    else:
        Xs = np.concatenate(
            [X[X[:, :, 2] == dim].reshape(n, -1, 3) for dim in homology_dimensions],
            axis=1,
        )
    if remove_dim:
        Xs = Xs[:, :, :2]
    return Xs


# Traceback (TypeError)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-5-5e233492b878> in <module>
----> 1 Amplitude(metric="persistence_image").fit_transform(Xd)

~\Workspace\giotto-tda_ulupo\gtda\utils\_docs.py in fit_transform_wrapper(*args, **kwargs)
104         @wraps(original_fit_transform)
105         def fit_transform_wrapper(*args, **kwargs):
--> 106             return original_fit_transform(*args, **kwargs)
107         fit_transform_wrapper.__doc__ = \
108             make_fit_transform_docs(fit_docs, transform_docs)

~\Anaconda3\envs\gtda-ulupo-py38\lib\site-packages\sklearn\base.py in fit_transform(self, X, y, **fit_params)
688         if y is None:
689             # fit method of arity 1 (unsupervised transformation)
--> 690             return self.fit(X, **fit_params).transform(X)
691         else:
692             # fit method of arity 2 (supervised transformation)

~\Workspace\giotto-tda_ulupo\gtda\diagrams\features.py in fit(self, X, y)
273         if self.metric == 'persistence_image':
274             self.effective_metric_params_['weights'] = \
--> 275                 _calculate_weights(X, **self.effective_metric_params_)
276
277         return self

TypeError: _calculate_weights() missing 1 required positional argument: 'weight_function'

# Fixed code
def _subdiagrams(X, homology_dimensions, remove_dim=False):
    """For each diagram in a collection, extract the subdiagrams in a given
    list of homology dimensions. It is assumed that all diagrams in X contain
    the same number of points in each homology dimension."""
    n_samples = len(X)
    X_0 = X[0]

    def _subdiagrams_single_homology_dimension(homology_dimension):
        n_features_in_dim = np.sum(X_0[:, 2] == homology_dimension)
        try:
            # In this case, reshape ensures copy
            Xs = X[X[:, :, 2] == homology_dimension].reshape(
                n_samples, n_features_in_dim, 3
            )
            return Xs
        except ValueError as e:
            if e.args[0].lower().startswith("cannot reshape array"):
                raise ValueError(
                    f"All persistence diagrams in the collection must have "
                    f"the same number of birth-death-dimension triples in any "
                    f"given homology dimension. This is not true in homology "
                    f"dimension {homology_dimension}. Trivial triples for "
                    f"which birth = death may be added or removed to fulfill "
                    f"this requirement."
                )
            else:
                raise e

    if len(homology_dimensions) == 1:
        Xs = _subdiagrams_single_homology_dimension(homology_dimensions[0])
    else:
        # np.concatenate will also create a copy
        Xs = np.concatenate(
            [
                _subdiagrams_single_homology_dimension(dim)
                for dim in homology_dimensions
            ],
            axis=1,
        )
    if remove_dim:
        Xs = Xs[:, :, :2]
    return Xs