NeuS-plot.py

from matplotlib import pyplot as plt
import numpy as np


n_samples = 128

# pdf of logistic distribution
def pdf(x, std):
    return 0.25 / std / np.cosh(0.5 * x / std)**2

# cdf of logistic distribution
def cdf(x, std):
    return 1 / (1 + np.exp(-x / std))

def calc_density(f, dfdt, std):
    return (-pdf(f, std) * dfdt / cdf(f, std)).clip(0, 1e7)


def calc_weights(f, dfdt, section_len, std):
    # direct weight construction
    weight_direct = pdf(f, std)
    weight_direct = weight_direct / np.sum(weight_direct)

    # weight construction of NeuS
    density = calc_density(f, dfdt, std)
    alpha = 1 - np.exp(-density * section_len)
    trans = np.cumprod(1 - alpha)
    weight_neus = np.concatenate([alpha[:1], trans[:-1] * alpha[1:]])

    return weight_direct, weight_neus



case = 'single_plane' # or 'multiple surfaces'


if __name__ == '__main__':
    section_len = 1.0 / n_samples * np.ones(n_samples)
    steps = np.linspace(0, 1, n_samples + 1)[:-1]
    if case == 'single_plane':
        f = np.linspace(0.5, -0.5, n_samples + 1)[:-1]  # signed distance function
        dfdt = -np.ones(n_samples)  # df/dt
    else:
        f = np.concatenate([
                np.linspace(0.125, -0.125, n_samples // 4 + 1)[:-1],
                np.linspace(-0.125, 0.125, n_samples // 4 + 1)[:-1],
                np.linspace(0.125, -0.125, n_samples // 4 + 1)[:-1],
                np.linspace(-0.125, 0.125, n_samples // 4 + 1)[:-1]])
        dfdt = np.concatenate([
                -np.ones(n_samples // 4),
                np.ones(n_samples // 4),
                -np.ones(n_samples // 4),
                np.ones(n_samples // 4)
            ])

    for s in range(1, 20):
        weight_direct, weight_neus = calc_weights(f, dfdt, section_len, 1 / (1.5**s))

        plt.plot(steps, weight_direct)
        plt.plot(steps, weight_neus)
        scaled_f = f / f.max() * np.max([weight_direct.max(), weight_neus.max()])
        plt.plot(steps, scaled_f)
        plt.savefig('./tmp_plot_single_plane/{}.png'.format(s))
        plt.cla()