Test of scipy.odr regressor

scipy_odr_test.py

from __future__ import print_function

import numpy as np
import scipy.linalg
from scipy.odr import *
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import pyplot as plt
import sys
import time

np.set_printoptions(formatter={'float': lambda f: "%12.8f" % f})

####################
# Helper functions #
####################

def explicit_to_implicit_params(p, normalize=True):
    assert p.shape[0] == 3
    return np.r_[p[0], p[1], -1, p[2]] / (p[2] if normalize else 1.0)

def implicit_to_explicit_params(p):
    assert p.shape[0] == 4
    return -np.array([p[0], p[1], p[3]]) / p[2]

###################
# Fitted function #
###################

def explicit_f(p, input):
    assert p.shape[0] == 3
    input = np.array(input)
    x = input[0]
    y = input[1]
    return p[0] * x + p[1] * y + p[2]

def implicit_f(p, input, normalize=True):
    assert p.shape[0] == 4
    input = np.array(input)
    x = input[0]
    y = input[1]
    z = input[2]
    if normalize:
        p = p/p[3]
    return p[0] * x + p[1] * y + p[2] * z + p[3]

######################
# Fitting procedures #
######################

def lsq_fit(data):
    t = time.time()
    a = np.c_[data[:, 0], data[:, 1], np.ones(data.shape[0])]
    c, _, _, _ = scipy.linalg.lstsq(a, data[:, 2])
    return c, time.time() - t

def odr_fit_explicit(data, initial_guess_explicit):
    t = time.time()
    myModel = Model(explicit_f)
    myData = Data([data[:, 0], data[:, 1]], y=data[:, 2])
    myOdr = ODR(myData, myModel, beta0=initial_guess_explicit)
    myOdr.maxit = 500
    out = myOdr.run()
    iters = out.iwork[22]
    return out.beta, out.stopreason, iters, time.time() - t

def odr_fit_implicit(data, initial_guess_implicit, normalize=True):
    t = time.time()
    myModel = Model(implicit_f, implicit=True, extra_args=[normalize])
    myData = Data([data[:, 0], data[:, 1], data[:, 2]], y=1)
    myOdr = ODR(myData, myModel, beta0=initial_guess_implicit)
    myOdr.maxit = 500
    out = myOdr.run()
    iters = out.iwork[25]
    return out.beta, out.stopreason, iters, time.time() - t

###########################
# Data loading/generation #
###########################

def get_data_generated():
    """Generate test data
    :return: (n data items of shape (n, 3) , x grid coordinates for evaluation, y grid coordinates for evaluation, value used as the bad initial guess (implicit params))
    :rtype: np.ndarray, np.ndarray, np.ndarray, np.ndarray
    """
    P = np.array([1.0, 2.0, 8.0, 3.0])  # 1x + 2y + 8z + 3 = 0    or     z = -1/8x - 1/4y -3/8
    print("P:                                 ", P / P[3])

    X, Y = np.meshgrid(np.arange(-3.0, 3.0, 0.5), np.arange(-3.0, 3.0, 0.5))
    x_grid = X.flatten()
    y_grid = Y.flatten()
    z_true = np.array(list(map(lambda data: explicit_f(implicit_to_explicit_params(P), data), np.c_[x_grid, y_grid])))
    data = np.c_[x_grid, y_grid, z_true]

    # specify noise parameters
    mean = np.array([0.0, 0.0, 0.0])
    noise_magnitude = 0.001
    cov = np.ones((3,3)) * noise_magnitude
    cov[2, 2] *= 10  # higher z noise
    noise = np.random.multivariate_normal(mean, cov, x_grid.shape[0])
    noisy_data = data + noise

    return noisy_data, X, Y, np.array([1, 1, 1, 1])

def get_data_pickle(pickle_path):
    """Load test data
    :param str pickle_path: Path to the pickle to load. Tuned for the data downloaded from the question at
                            https://stackoverflow.com/q/48115464/1076564
    :return: (n data items of shape (n, 3) , x grid coordinates for evaluation, y grid coordinates for evaluation, value used as the bad initial guess (implicit params))
    :rtype: np.ndarray, np.ndarray, np.ndarray
    """
    try:
        import cPickle
        with open(pickle_path, 'rb') as f:
            d = cPickle.load(f)
    except:
        import pickle
        with open(pickle_path, 'rb') as f:
            d = pickle.load(f)

    d = np.array(d)
    X, Y = np.meshgrid(np.arange(-30.0, 30.0, 5.), np.arange(-135.0, -125.0, 5.))

    return d, X, Y, np.array([1, 1, 1, 1])

##################
# Data selection #
##################

# generate or load data (uncomment what you want to try, not both at the same time)
# noisy_data, X, Y, bad_initial_guess = get_data_generated()
noisy_data, X, Y, bad_initial_guess = get_data_pickle(sys.argv[1])

###########
# Fitting #
###########

def print_results(name, beta, stop_reason, iters, t):
    print(
        "%-35s" % (name + ":",),
        beta / beta[3],
        "stop reason: %-30s," % stop_reason[0] if stop_reason is not None else " " * 43,
        "iters: %3i," % iters if iters is not None else " " * 12,
        "%8.5f" % t, "secs, ",
        "beta denormalized magnitudes: ",
        ["E%+4i" % int(("%8.0e" % x).split('e')[1]) for x in beta.flatten()],
    )

# LSQ fit
lsq_params, t = lsq_fit(noisy_data)
# evaluate it on grid
lsq_Z = explicit_f(lsq_params, [X, Y])
print_results("LSQ", explicit_to_implicit_params(lsq_params, normalize=False), None, None, t)

# explicit ODR fit
odr_ex_params, odr_ex_stop_reason, iters, t = odr_fit_explicit(noisy_data, lsq_params)
odr_ex_Z = explicit_f(odr_ex_params, [X, Y])
print_results("ODR explicit (LSQ initial guess)", explicit_to_implicit_params(odr_ex_params, normalize=False), odr_ex_stop_reason, iters, t)

# implicit ODR fit
odr_im_params, odr_im_stop_reason, iters, t = odr_fit_implicit(noisy_data, explicit_to_implicit_params(lsq_params))
odr_im_Z = explicit_f(implicit_to_explicit_params(odr_im_params), [X, Y])
print_results("ODR implicit (LSQ initial guess)", odr_im_params, odr_im_stop_reason, iters, t)

# implicit ODR fit without normalization
odr_im_no_norm_params, odr_im_no_norm_stop_reason, iters, t = odr_fit_implicit(noisy_data, explicit_to_implicit_params(lsq_params), normalize=False)
odr_im_no_norm_Z = explicit_f(implicit_to_explicit_params(odr_im_no_norm_params), [X, Y])
print_results("ODR implicit (LSQ, no norm.)", odr_im_no_norm_params, odr_im_no_norm_stop_reason, iters, t)

# explicit ODR fit with bad initial params
odr_ex_bad_params, odr_ex_bad_stop_reason, iters, t = odr_fit_explicit(noisy_data, implicit_to_explicit_params(bad_initial_guess))
odr_ex_bad_Z = explicit_f(odr_ex_bad_params, [X, Y])
print_results("ODR explicit (bad initial guess)", explicit_to_implicit_params(odr_ex_bad_params, normalize=False), odr_ex_bad_stop_reason, iters, t)

# implicit ODR fit with bad initial params
odr_im_bad_params, odr_im_bad_stop_reason, iters, t = odr_fit_implicit(noisy_data, bad_initial_guess)
odr_im_bad_Z = explicit_f(implicit_to_explicit_params(odr_im_bad_params), [X, Y])
print_results("ODR implicit (bad initial guess)", odr_im_bad_params, odr_im_bad_stop_reason, iters, t)

# implicit ODR fit with bad initial params and without normalization
odr_im_bad_no_norm_params, odr_im_bad_no_norm_stop_reason, iters, t = odr_fit_implicit(noisy_data, bad_initial_guess, normalize=False)
odr_im_bad_no_norm_Z = explicit_f(implicit_to_explicit_params(odr_im_bad_no_norm_params), [X, Y])
print_results("ODR implicit (bad i.g., no norm.)", odr_im_bad_no_norm_params, odr_im_bad_no_norm_stop_reason, iters, t)

############
# Plotting #
############

fig = plt.figure(figsize=(20,20))
ax = fig.gca(projection='3d')

# 3D surfaces cannot have legend, so we create fake lines to show their legend instead
legend = [
    [mpl.lines.Line2D([0], [0], linestyle="none", c='b', marker='o'), "LSQ"],
    [mpl.lines.Line2D([0], [0], linestyle="none", c='r', marker='o'), "ODR explicit (LSQ initial guess)"],
    [mpl.lines.Line2D([0], [0], linestyle="none", c='g', marker='o'), "ODR implicit (LSQ initial guess)"],
    [mpl.lines.Line2D([0], [0], linestyle="none", c='y', marker='o'), "ODR implicit (LSQ initial guess) (w/o norm.)"],
    [mpl.lines.Line2D([0], [0], linestyle="none", c='c', marker='o'), "ODR explicit (bad initial guess)"],
    [mpl.lines.Line2D([0], [0], linestyle="none", c='m', marker='o'), "ODR implicit (bad initial guess)"],
    [mpl.lines.Line2D([0], [0], linestyle="none", c=(0.5, 0.5, 0.5), marker='o'), "ODR implicit (bad initial guess) (w/o norm.)"]
]

ax.legend([x[0] for x in legend], [y[1] for y in legend], numpoints=1)

# plot the input data
ax.scatter(noisy_data[1:100:, 0], noisy_data[1:100:, 1], noisy_data[1:100:, 2], alpha=1.0, c='k', s=20)

# plot the fits
ax.plot_surface(X, Y, lsq_Z, rstride=1, cstride=1, alpha=0.2, color=legend[0][0]._color)
ax.plot_surface(X, Y, odr_ex_Z, rstride=1, cstride=1, alpha=0.5, color=legend[1][0]._color)
ax.plot_surface(X, Y, odr_im_Z, rstride=1, cstride=1, alpha=0.5, color=legend[2][0]._color)
ax.plot_surface(X, Y, odr_im_no_norm_Z, rstride=1, cstride=1, alpha=0.5, color=legend[3][0]._color)
ax.plot_surface(X, Y, odr_ex_bad_Z, rstride=1, cstride=1, alpha=0.5, color=legend[4][0]._color)
ax.plot_surface(X, Y, odr_im_bad_Z, rstride=1, cstride=1, alpha=0.5, color=legend[5][0]._color)
ax.plot_surface(X, Y, odr_im_bad_no_norm_Z, rstride=1, cstride=1, alpha=0.5, color=legend[6][0]._color)



plt.xlabel('X')
plt.ylabel('Y')
ax.set_zlabel('Z')
ax.axis('equal')
ax.axis('tight')

plt.show()