# Buggy code
def persistence_image_distances(
    diagrams_1,
    diagrams_2,
    sampling,
    step_size,
    weight_function=lambda x: x,
    sigma=1.0,
    p=2.0,
    **kwargs,
):
    sampling_ = np.copy(sampling.reshape((-1,)))
    weights = weight_function(sampling_ - sampling_[0])
    persistence_image_1 = persistence_images(
        diagrams_1, sampling_, step_size, weights, sigma
    ).reshape(diagrams_1.shape[0], -1)
    if np.array_equal(diagrams_1, diagrams_2):
        unnorm_dist = squareform(pdist(persistence_image_1, "minkowski", p=p))
        return (step_size ** (1 / p)) * unnorm_dist
    persistence_image_2 = persistence_images(
        diagrams_2,
        sampling_,
        step_size,
        weights,
        sigma,
    ).reshape(diagrams_2.shape[0], -1)
    unnorm_dist = cdist(persistence_image_1, persistence_image_2, "minkowski", p=p)
    return (step_size ** (1 / p)) * unnorm_dist


# 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 persistence_image_distances(
    diagrams_1,
    diagrams_2,
    sampling,
    step_size,
    sigma=0.1,
    weight_function=np.ones_like,
    p=2.0,
    **kwargs,
):
    # For persistence images, `sampling` is a tall matrix with two columns
    # (the first for birth and the second for persistence), and `step_size` is
    # a 2d array
    weights = weight_function(sampling[:, 1])
    step_sizes_factor = np.product(step_size) ** (1 / p)
    # WARNING: `persistence_images` modifies `diagrams` in place
    are_arrays_equal = np.array_equal(diagrams_1, diagrams_2)
    persistence_images_1 = persistence_images(
        diagrams_1, sampling, step_size, sigma, weights
    ).reshape(len(diagrams_1), -1)
    if are_arrays_equal:
        distances = pdist(persistence_images_1, "minkowski", p=p)
        distances *= step_sizes_factor
        return squareform(distances)
    persistence_images_2 = persistence_images(
        diagrams_2, sampling, step_size, sigma, weights
    ).reshape(len(diagrams_2), -1)
    distances = cdist(persistence_images_1, persistence_images_2, "minkowski", p=p)
    distances *= step_sizes_factor
    return distances