Source code for rofunc.utils.datalab.data_generator.bezier

# Copyright 2023, Junjia LIU, jjliu@mae.cuhk.edu.hk
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import matplotlib.pyplot as plt
import numpy as np
from math import factorial


[docs]def comb(n, k): return factorial(n) // (factorial(k) * factorial(n - k))
[docs]def get_bezier_curve(points): n = len(points) - 1 return lambda t: sum(comb(n, i) * t ** i * (1 - t) ** (n - i) * points[i] for i in range(n + 1))
[docs]def evaluate_bezier(points, total): """ Generate the Bezier curve from points. :param points: points for shaping the Bezier curve :param total: the number of points on the Bezier curve :return: """ bezier = get_bezier_curve(points) new_points = np.array([bezier(t) for t in np.linspace(0, 1, total)]) return new_points
[docs]def plot_bezier(bx, by, x, y, bz=None, z=None, ax=None): """ Plot the Bezier curve, either in 2D or 3D. Example:: >>> from rofunc.utils.datalab.data_generator.bezier import plot_bezier >>> import numpy as np >>> demo_points = np.array([[0, 0], [-1, 8], [4, 3], [2, 1], [4, 3]]) >>> bx, by = evaluate_bezier(demo_points, 50)[:, 0], evaluate_bezier(demo_points, 50)[:, 1] >>> x, y = demo_points[:, 0], demo_points[:, 1] >>> plot_bezier(bx, by, x, y) :param bx: the x coordinates of the Bezier curve :param by: the y coordinates of the Bezier curve :param x: the x coordinates of the demonstration points :param y: the y coordinates of the demonstration points :param bz: the z coordinates of the Bezier curve, default None :param z: the z coordinates of the demonstration points, default None :param ax: the axis to plot the Bezier curve, default None :return: """ assert isinstance(bx, list) or isinstance(bx, np.ndarray) if bz is not None and z is not None: for i in range(len(bx)): ax.plot(bx[i], by[i], bz[i]) # ax.plot(x[i], y[i], z[i], 'r.') else: for i in range(len(bx)): plt.plot(bx[i], by[i])
# plt.plot(x[i], y[i], 'r.')
[docs]def multi_bezier_demos(demo_points, ax=None): """ Generate multiple Bezier curves as demonstrations. Example:: >>> from rofunc.utils.datalab.data_generator.bezier import multi_bezier_demos >>> import numpy as np >>> demo_points = np.array([[[0, 0], [-1, 8], [4, 3], [2, 1], [4, 3]], ... [[0, -2], [-1, 7], [3, 2.5], [2, 1.6], [4, 3]], ... [[0, -1], [-1, 8], [4, 5.2], [2, 1.1], [4, 3.5]]]) >>> demos_x = multi_bezier_demos(demo_points) :param demo_points: points for shaping the Bezier curve :param ax: the axis to plot the Bezier curves, default None :return: """ bx_lst = [evaluate_bezier(demo_point, 50)[:, 0] for demo_point in demo_points] by_lst = [evaluate_bezier(demo_point, 50)[:, 1] for demo_point in demo_points] x_lst = [demo_point[:, 0] for demo_point in demo_points] y_lst = [demo_point[:, 1] for demo_point in demo_points] if len(demo_points[0][0]) == 3: bz_lst = [evaluate_bezier(demo_point, 50)[:, 2] for demo_point in demo_points] z_lst = [demo_point[:, 2] for demo_point in demo_points] plot_bezier(bx_lst, by_lst, x_lst, y_lst, bz_lst, z_lst, ax) demos_x = np.concatenate( (np.array(bx_lst).reshape((len(demo_points), -1, 1)), np.array(by_lst).reshape((len(demo_points), -1, 1)), np.array(bz_lst).reshape((len(demo_points), -1, 1))), axis=2) else: plot_bezier(bx_lst, by_lst, x_lst, y_lst) demos_x = np.concatenate( (np.array(bx_lst).reshape((len(demo_points), -1, 1)), np.array(by_lst).reshape((len(demo_points), -1, 1))), axis=2) return demos_x