Persistence landscape discretization

landscape.py

# Task: convert persistence landscape given in the form (N, k, n, 2)
# to its another representation of the shape (N, k, n, p), where N is
# the batch size, k - number of feature dimensionalities, n - number of
# layers. And the last dimension is 2 dates of event's birth and death
# in the former brackets and p is number of discretization steps
# in the latter, on the interval where these event dates are defined.
import torch
import matplotlib.pyplot as plt

def triangle_hat(x, a, b):
    out = torch.zeros_like(x)
    left_slope = (x > a) & (x < .5*(a+b))
    right_slope = (x >= .5*(a+b)) & (x < b)
    out[left_slope] = (x - a)[left_slope]
    out[right_slope] = (b - x)[right_slope]
    return out

def discretize(data, lims, n_bins=100):
    sh = data.shape[:-1]
    a = data[..., 0][..., None].expand(*sh, n_bins)
    b = data[..., 1][..., None].expand(*sh, n_bins)
    x = torch.linspace(*lims, n_bins)
    return x, triangle_hat(x.expand(*sh, -1), a, b)

if __name__ == '__main__':
    data,_ = torch.rand(3, 4, 2).sort(1)
    print(data[0, 0])
    x, hats = discretize(data, (0, 1), 200)
    plt.plot(x, hats[0, 0]);