# Buggy code
def persistence_images(diagrams, sampling, step_size, weights, sigma):
    persistence_images_ = np.zeros(
        (diagrams.shape[0], sampling.shape[0], sampling.shape[0])
    )
    # Transform diagrams from (birth, death, dim) to (birth, persistence, dim)
    diagrams[:, :, 1] = diagrams[:, :, 1] - diagrams[:, :, 0]

    for axis in [0, 1]:
        # Set the values outside of the sampling range to the sampling range.
        diagrams[:, :, axis][diagrams[:, :, axis] < sampling[0, axis]] = sampling[
            0, axis
        ]
        diagrams[:, :, axis][diagrams[:, :, axis] > sampling[-1, axis]] = sampling[
            -1, axis
        ]
        # Convert into pixel
        diagrams[:, :, axis] = np.array(
            (diagrams[:, :, axis] - sampling[0, axis]) / step_size[axis], dtype=int
        )
    # Sample the image
    [
        _sample_image(persistence_images_[i], sampled_diag)
        for i, sampled_diag in enumerate(diagrams)
    ]

    # Apply the weights
    persistence_images_ *= weights / np.max(weights)

    # Smoothen the weighted-image
    for i, image in enumerate(persistence_images_):
        persistence_images_[i] = gaussian_filter(image, sigma, mode="reflect")

    persistence_images_ = np.rot90(persistence_images_, k=1, axes=(1, 2))
    return persistence_images_


# 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_images(diagrams, sampling, step_size, sigma, weights):
    # 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
    # WARNING: modifies `diagrams` in place
    persistence_images_ = np.zeros(
        (len(diagrams), len(sampling), len(sampling)), dtype=float
    )
    # If both step sizes are zero, we return a trivial image
    if (step_size == 0).all():
        return persistence_images_

    # Transform diagrams from (birth, death, dim) to (birth, persistence, dim)
    diagrams[:, :, 1] -= diagrams[:, :, 0]

    sigma_pixel = []
    first_samplings = sampling[0]
    last_samplings = sampling[-1]
    for ax in [0, 1]:
        diagrams_ax = diagrams[:, :, ax]
        # Set the values outside of the sampling range
        diagrams_ax[diagrams_ax < first_samplings[ax]] = first_samplings[ax]
        diagrams_ax[diagrams_ax > last_samplings[ax]] = last_samplings[ax]
        # Calculate the value of the component of `sigma` in pixel units
        sigma_pixel.append(sigma / step_size[ax])

    # Sample the image, apply the weights, smoothen
    for i, diagram in enumerate(diagrams):
        nontrivial_points_idx = np.flatnonzero(diagram[:, 1])
        diagram_nontrivial_pixel_coords = np.array(
            (diagram - first_samplings) / step_size, dtype=int
        )[nontrivial_points_idx]
        image = persistence_images_[i]
        _sample_image(image, diagram_nontrivial_pixel_coords)
        image *= weights
        gaussian_filter(image, sigma_pixel, mode="constant", output=image)

    persistence_images_ = np.rot90(persistence_images_, k=1, axes=(1, 2))
    persistence_images_ /= np.product(step_size)
    return persistence_images_