def test_raise_degree(self):
np.random.seed(12345)
x = [0, 1]
k, d = 8, 5
c = np.random.random((k, 1, 2, 3, 4))
bp = BPoly(c, x)
c1 = BPoly._raise_degree(c, d)
bp1 = BPoly(c1, x)
xp = np.linspace(0, 1, 11)
assert_allclose(bp(xp), bp1(xp))
def test_extrapolate_attr(self):
x = [0, 2]
c = [[3], [1], [4]]
bp = BPoly(c, x)
for extrapolate in (True, False, None):
bp = BPoly(c, x, extrapolate=extrapolate)
bp_d = bp.derivative()
if extrapolate is False:
assert_(np.isnan(bp([-0.1, 2.1])).all())
assert_(np.isnan(bp_d([-0.1, 2.1])).all())
else:
assert_(not np.isnan(bp([-0.1, 2.1])).any())
assert_(not np.isnan(bp_d([-0.1, 2.1])).any())
def test_two_intervals(self):
x = [0, 1, 3]
c = [[3, 0], [0, 0], [0, 2]]
bp = BPoly(c, x) # [3*(1-x)**2, 2*((x-1)/2)**2]
assert_allclose(bp(0.4), 3 * 0.6 * 0.6)
assert_allclose(bp(1.7), 2 * (0.7 / 2)**2)
def test_derivative(self):
x = [0, 1, 3]
c = [[3, 0], [0, 0], [0, 2]]
bp = BPoly(c, x) # [3*(1-x)**2, 2*((x-1)/2)**2]
bp_der = bp.derivative()
assert_allclose(bp_der(0.4), -6 * (0.6))
assert_allclose(bp_der(1.7), 0.7)
def test_simple5(self):
x = [0, 1]
c = [[1], [1], [8], [2], [1]]
bp = BPoly(c, x)
assert_allclose(
bp(0.3), 0.7**4 + 4 * 0.7**3 * 0.3 + 8 * 6 * 0.7**2 * 0.3**2 +
2 * 4 * 0.7 * 0.3**3 + 0.3**4)
def Bern(vars, pars):
pars_coef = []
for i in range(len(pars)):
pars_coef.append(pars[i])
pars_coef = np.array(pars_coef).reshape(-1, 1)
return BPoly(pars_coef[0:-2],
[pars_coef[-2][0], pars_coef[-1][0]])(vars[0])
def test_interval_length(self):
x = [0, 2]
c = [[3], [1], [4]]
bp = BPoly(c, x)
xval = 0.1
s = xval / 2 # s = (x - xa) / (xb - xa)
assert_allclose(
bp(xval), 3 * (1 - s) * (1 - s) + 1 * 2 * s * (1 - s) + 4 * s * s)
def test_make_poly_12(self):
np.random.seed(12345)
ya = np.r_[0, np.random.random(5)]
yb = np.r_[0, np.random.random(5)]
c = BPoly._construct_from_derivatives(0, 1, ya, yb)
pp = BPoly(c[:, None], [0, 1])
for j in range(6):
assert_allclose([pp(0.), pp(1.)], [ya[j], yb[j]])
pp = pp.derivative()
def test_pp_from_bp(self):
x = [0, 1, 3]
c = [[3, 3], [1, 1], [4, 2]]
bp = BPoly(c, x)
pp = PPoly.from_bernstein_basis(bp)
bp1 = BPoly.from_power_basis(pp)
xp = [0.1, 1.4]
assert_allclose(bp(xp), pp(xp))
assert_allclose(bp(xp), bp1(xp))
def test_multi_shape(self):
c = np.random.rand(6, 2, 1, 2, 3)
x = np.array([0, 0.5, 1])
p = BPoly(c, x)
assert_equal(p.x.shape, x.shape)
assert_equal(p.c.shape, c.shape)
assert_equal(p(0.3).shape, c.shape[2:])
assert_equal(p(np.random.rand(5, 6)).shape, (5, 6) + c.shape[2:])
dp = p.derivative()
assert_equal(dp.c.shape, (5, 2, 1, 2, 3))
def test_bernstein(self):
# a special knot vector: Bernstein polynomials
k = 3
t = np.asarray([0] * (k + 1) + [1] * (k + 1))
c = np.asarray([1., 2., 3., 4.])
bp = BPoly(c.reshape(-1, 1), [0, 1])
bspl = BSpline(t, c, k)
xx = np.linspace(-1., 2., 10)
assert_allclose(bp(xx, extrapolate=True),
bspl(xx, extrapolate=True),
atol=1e-14)
assert_allclose(splev(xx, (t, c, k)), bspl(xx), atol=1e-14)
def test_derivative_ppoly(self):
# make sure it's consistent w/ power basis
np.random.seed(1234)
m, k = 5, 8 # number of intervals, order
x = np.sort(np.random.random(m))
c = np.random.random((k, m - 1))
bp = BPoly(c, x)
pp = PPoly.from_bernstein_basis(bp)
for d in range(k):
bp = bp.derivative()
pp = pp.derivative()
xp = np.linspace(x[0], x[-1], 21)
assert_allclose(bp(xp), pp(xp))
def test_simple3(self):
x = [0, 1]
c = [[3], [1], [4]]
bp = BPoly(c, x) # 3 * (1-x)**2 + 2 * x (1-x) + 4 * x**2
assert_allclose(bp(0.2),
3 * 0.8 * 0.8 + 1 * 2 * 0.2 * 0.8 + 4 * 0.2 * 0.2)
def run(self):
import ROOT
inp = self.input()
outp = self.output()
interpolation_bins = 1000
# cannot get the function from ROOT, use scipy instead
from scipy.interpolate import PchipInterpolator, BPoly, PPoly
# get categories in which to fit the scale factors
categories = []
for category, _, _ in self.config_inst.walk_categories():
if len(self.category_tags) > 0 and not category.has_tag(self.category_tags, mode=any):
continue
if self.has_c_shift:
if category.get_aux("region", None) == "c" and category.get_aux("phase_space") == "measure":
if category.has_tag(self.b_tagger):
categories.append(category)
else:
if category.has_tag(("merged", self.b_tagger), mode=all) and category.get_aux("phase_space") == "measure":
categories.append(category)
# get scaling factors for normalization
if self.fix_normalization:
norm_factors = inp["norm"].load()[self.effective_shift]
# contents of .csv file for scale factors
fit_results = []
# finely binned histograms to write to the output file
hist_dict = {}
# TF1's to write to the output file
function_dict = {}
with inp["sf"]["scale_factors"].load("r") as input_file:
for category in categories:
region = category.get_aux("region")
category_dir = input_file.GetDirectory(category.name)
# get scale factor histogram
hist_keys = category_dir.GetListOfKeys()
if len(hist_keys) != 1:
raise ValueError("Found more than one histogram in %s, cannot identify scale "
"factor hist." % category_dir)
hist = category_dir.Get(hist_keys[0].GetName())
nbins = hist.GetNbinsX()
if self.fix_normalization:
hist.Scale(norm_factors[category.name])
x_axis = hist.GetXaxis()
interpolation_hist = ROOT.TH1D(hist.GetName() + "_fine", hist.GetTitle(),
interpolation_bins, x_axis.GetXmin(), x_axis.GetXmax())
x_values = ROOT.vector("double")()
y_values = ROOT.vector("double")()
for bin_idx in range(1, nbins + 1):
if hist.GetBinCenter(bin_idx) < 0:
continue
x_values.push_back(hist.GetBinCenter(bin_idx))
y_values.push_back(hist.GetBinContent(bin_idx))
interpolator = PchipInterpolator(x_values, y_values)
# define region in which to use interpolation
first_point, last_point = min(x_values), max(x_values)
# create finely binned histogram from either TF1 or interpolator
for bin_idx in range(interpolation_bins + 2):
bin_center = interpolation_hist.GetBinCenter(bin_idx)
if bin_center < 0:
interpolation_hist.SetBinContent(bin_idx, hist.GetBinContent(1))
elif bin_center < first_point:
interpolation_hist.SetBinContent(bin_idx, interpolator(first_point))
elif bin_center > last_point:
interpolation_hist.SetBinContent(bin_idx, interpolator(last_point))
else:
interpolation_hist.SetBinContent(bin_idx, interpolator(bin_center))
hist_dict[category] = interpolation_hist
# fill .csv file in final iteration (after normalization fix)
# also create piecewise linear TF1's
if self.fix_normalization:
function_pieces = []
results = {}
results["eta_min"], results["eta_max"] = category.get_aux("eta")
pt_range = category.get_aux("pt")
results["pt_min"] = pt_range[0]
results["pt_max"] = min(pt_range[1], 10000.) # replace inf
results["flavor_id"] = self.config_inst.get_aux("flavor_ids")[region]
if self.effective_shift == "nominal":
sysType = "central"
else:
sys_name, direction = self.effective_shift.rsplit("_", 1)
sysType = "{}_{}".format(direction,
sys_name.replace("c_stats", "cferr").replace("lf_stats", "lfstats").replace("hf_stats", "hfstats")
)
results["sysType"] = sysType
# skip unwanted combinations
if "cferr" in sysType and region != "c":
continue
fit_results_tpl = "3, iterativefit, {sysType}, {flavor_id}, {eta_min}, " \
"{eta_max}, {pt_min}, {pt_max}".format(**results)
fit_results.append(fit_results_tpl + ", -15, 0, {}".format(hist.GetBinContent(1)))
fit_results.append(fit_results_tpl + ", 0, {}, {}".format(first_point,
interpolator(first_point)))
function_pieces.append("(x < 0) * {}".format(hist.GetBinContent(1)))
function_pieces.append("(x >= 0) * (x < {}) * {}".format(first_point,
interpolator(first_point)))
# intermediate functions
# change interpolated function from bernstein to power basis
bpoly_interpolation = BPoly(interpolator.c, interpolator.x)
ppoly_interpolation = PPoly.from_bernstein_basis(bpoly_interpolation)
interpolator_idx = 0
for bin_idx in range(1, nbins):
if hist.GetBinCenter(bin_idx) < first_point:
continue
x_min = hist.GetBinCenter(bin_idx)
x_max = hist.GetBinCenter(bin_idx + 1)
interpolator_coefficients = ppoly_interpolation.c[:, interpolator_idx]
interpolator_x = ppoly_interpolation.x[interpolator_idx]
interpolator_idx += 1
formula = ""
for i in xrange(3):
formula += "{}*".format(interpolator_coefficients[i]) + \
"*".join(["(x-{})".format(interpolator_x)] * (3 - i))
formula += "+" if interpolator_coefficients[i + 1] >= 0. else ""
formula += str(interpolator_coefficients[3])
fit_results.append(fit_results_tpl + ", {}, {}, {}".format(x_min,
x_max, formula))
function_pieces.append("(x >= {}) * (x < {}) * ({})".format(
x_min, x_max, formula))
fit_results.append(fit_results_tpl + ", {}, 1.1, {}".format(last_point,
interpolator(last_point)))
function_pieces.append("(x >= {}) * {}".format(last_point,
interpolator(last_point)))
function = ROOT.TF1("sf_{}".format(category.name),
" + ".join(["({})".format(piece) for piece in function_pieces]),
-2., 1.1)
function_dict[category] = function
# sanity check
for val in x_values:
if (abs(function.Eval(val) - interpolator(val)) > 1e-5):
raise Exception("SF Function does not match interpolator values: "
"{} vs {}".format(function.Eval(val), interpolator(val)))
# write to output file
with outp["sf"].localize("w") as tmp:
with tmp.dump("RECREATE") as output_file:
for category, hist in hist_dict.items():
category_dir = output_file.mkdir(category.name)
category_dir.cd()
hist.Write("sf")
if self.fix_normalization:
with outp["csv"].localize("w") as tmp:
with tmp.open("w") as result_file:
result_file.write("\n".join(fit_results))
with outp["functions"].localize("w") as tmp:
with tmp.dump("RECREATE") as output_file:
for category, func in function_dict.items():
category_dir = output_file.mkdir(category.name)
category_dir.cd()
func.Write("sf")
def __init__(self, control_points, intervals):
self.control_points = control_points
self.intervals = intervals
self.poly = BPoly(self.control_points, self.intervals)
def poly_bezier(p_c, n):
p_c = np.expand_dims(p_c, axis=1)
bp = BPoly(p_c, [0, 1])
t = np.linspace(0.0, 1.0, n)
p = bp(t)
return p
def test_simple(self):
x = [0, 1]
c = [[3]]
bp = BPoly(c, x)
assert_allclose(bp(0.1), 3.)
def test_simple2(self):
x = [0, 1]
c = [[3], [1]]
bp = BPoly(c, x) # 3*(1-x) + 1*x
assert_allclose(bp(0.1), 3 * 0.9 + 1. * 0.1)
print(
acfunc(np.array([1, 2, 3]),
np.array([0.15505102572168228, 0.6449489742783174, 1.0]), ks))
# ks = func1(np.array(x2), np.array(y2))
# print(acfunc(np.array(x2), np.array(y2), ks))
S6 = 6**0.5
x = Symbol('x')
Eq = (x**3 * sp.cos(x / 2) + 1 / 2) * sp.sqrt(4 - x**2)
ans = sp.Integral(Eq, (x, -2, 2)).evalf(10)
# print(ans)
Pc = [[13 / 3 + 7 * S6 / 3, -23 / 3 - 22 * S6 / 3, 10 / 3 + 5 * S6],
[13 / 3 - 7 * S6 / 3, -23 / 3 + 22 * S6 / 3, 10 / 3 - 5 * S6],
[1 / 3, -8 / 3, 10 / 3]]
X = [0, 1, 2]
XX = [0, 1, 2, 3]
# print(Pc)
# X = [1, 2, 3]
Y = [0.15505102572168228, 0.6449489742783174, 1.0]
cs = CubicSpline(X, Y, axis=1)
css = PchipInterpolator(X, Y)
csss = CubicHermiteSpline(X, Y, [0, 0, 0])
cs4 = Akima1DInterpolator(X, Y)
cs5 = BPoly(Pc, XX)
# print(cs.c)
# print(css.c)
# print(csss.c)
# print(cs4.c)
# print(cs5)
# print(-0.9797959 + 1.46969385)
def test_simple4(self):
x = [0, 1]
c = [[1], [1], [1], [2]]
bp = BPoly(c, x)
assert_allclose(
bp(0.3), 0.7**3 + 3 * 0.7**2 * 0.3 + 3 * 0.7 * 0.3**2 + 2 * 0.3**3)
def solveProblem():
"""Test the backtrack line search with IP solver."""
prob = 11
if len(sys.argv) > 1:
prob = int(sys.argv[1])
print_purple("Testing on problem: %d" % prob)
tgp = loadTGP(
'dataset/tgp_%d.tgp' % prob
) # every time you have to reload from hard disk since it is modified before
initial_time_allocation = np.array([box.t for box in tgp.getCorridor()])
# we do not limit velocity in the open source implementation
tgp.doLimitVelocity = False
# we minimize jerk
tgp.minimizeOrder = 3
# we use 6-th order piecewise Bezier spline
tgp.trajectoryOrder = 6
#print_green("Use Mosek + Adaptive line search")
solver = IndoorQPProblemMOSEK(tgp, verbose=False)
ts1 = time.time()
solver.solve_once()
tf1 = time.time()
initial_obj = solver.obj
mosek_once_time = tf1 - ts1
t_before_opt, coeff_before_opt = solver.get_output_coefficients()
solver.grad_method = 'ours'
# extract initial trajectory
poly_coef = solver.get_coef_matrix().transpose((1, 0, 2))
#import pdb; pdb.set_trace()
break_points = np.insert(np.cumsum(solver.room_time), 0, 0.0)
initial_trajectory = BPoly(poly_coef, break_points)
tt_initial = np.linspace(0.0, break_points[-1], 100)
ts2 = time.time()
is_okay, converged = solver.refine_time_by_backtrack(
max_iter=100, log=True, adaptiveLineSearch=True)
tf2 = time.time()
#print(solver.room_time)
computation_time = (tf2 - ts2 + mosek_once_time) * 1000
#print('mosek first computation time', mosek_once_time, 'mosek computation time:', mosek_time)
mosek_convergence = solver.log
mosek_convergence[1::2] += mosek_once_time
t_after_opt, coeff_after_opt = solver.get_output_coefficients()
result = [
["Solver", "Mosek"],
["Solved?", is_okay],
#["Converged?", converged],
["Initial Cost", round(initial_obj, 3)],
["Final Cost", round(solver.obj, 3)],
["Solve Time [ms]", round(computation_time, 2)],
["# Major Iterations", solver.major_iteration],
["# Function Evaluation", solver.num_prob_solve]
]
print("Results")
print(
tabulate(result, tablefmt="psql", stralign="right", numalign="center"))
print("Initial time allocation:\n", np.round(initial_time_allocation, 2))
print("Final time allocation:\n", np.round(solver.room_time, 2))
if DO_PLOT_RESULTS:
poly_coef = solver.get_coef_matrix2()
break_points = np.insert(np.cumsum(t_after_opt), 0, 0.0)
final_trajectory = BPoly(poly_coef, break_points)
tt_final = np.linspace(0.0, break_points[-1], 100)
# plot 3D trajectory
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
boxes = tgp.getCorridor()
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
for i in range(len(boxes)):
vertices = boxes[i].vertex
ax.scatter(vertices[:, 0], vertices[:, 1], vertices[:, 2], s=1)
Z = vertices
verts = [[Z[0], Z[1], Z[2], Z[3]], [Z[4], Z[5], Z[6], Z[7]],
[Z[0], Z[1], Z[5], Z[4]], [Z[2], Z[3], Z[7], Z[6]],
[Z[1], Z[2], Z[6], Z[5]], [Z[4], Z[7], Z[3], Z[0]],
[Z[2], Z[3], Z[7], Z[6]]]
# plot safe corridor
pc = Poly3DCollection(verts,
alpha=0.0,
facecolor='gray',
linewidths=0.1,
edgecolors='red')
ax.add_collection3d(pc)
ax.plot(initial_trajectory(tt_initial)[:, 0],
initial_trajectory(tt_initial)[:, 1],
initial_trajectory(tt_initial)[:, 2],
label="Before refinement")
ax.plot(final_trajectory(tt_final)[:, 0],
final_trajectory(tt_final)[:, 1],
final_trajectory(tt_final)[:, 2],
label="After refinement")
#import pdb; pdb.set_trace()
ax.scatter(tgp.position[0, 0],
tgp.position[0, 1],
tgp.position[0, 2],
marker="*",
s=20,
label="Start")
ax.scatter(tgp.position[1, 0],
tgp.position[1, 1],
tgp.position[1, 2],
marker="o",
s=20,
label="Goal")
set_axes_equal(ax)
ax.set_axis_off()
ax.legend()
ax.set_title("3D trajectory")
# plot velocity/acceleration
fig, ax = plt.subplots(3, 1, figsize=(6, 7))
for i in range(3):
ax[i].plot(tt_initial,
initial_trajectory(tt_initial, 1)[:, i],
'-.',
label="Before refinement")
ax[i].plot(tt_final,
final_trajectory(tt_final, 1)[:, i],
label="After refinement")
if i == 0:
ax[i].set_title("X Velocity")
elif i == 1:
ax[i].set_title("Y Velocity")
elif i == 2:
ax[i].set_title("Z Velocity")
else:
pass
ax[i].legend()
ax[i].grid()
plt.tight_layout()
# plot acceleration
fig, ax = plt.subplots(3, 1, figsize=(6, 7))
for i in range(3):
ax[i].plot(tt_initial,
initial_trajectory(tt_initial, 2)[:, i],
'-.',
label="Before refinement")
ax[i].plot(tt_final,
final_trajectory(tt_final, 2)[:, i],
label="After refinement")
if i == 0:
ax[i].set_title("X Acceleration")
elif i == 1:
ax[i].set_title("Y Acceleration")
elif i == 2:
ax[i].set_title("Z Acceleration")
else:
pass
ax[i].legend()
ax[i].grid()
plt.tight_layout()
plt.show()
