import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures

# 生成随机数作为x变量，范围在-5到5之间，共500个样本
x = np.random.uniform(-5, 5, size=100)
X = x.reshape(-1, 1)

# 根据条件设置y变量，当x>0时，y设置为2；否则符合二次方程的值加上正态分布的噪声
y = np.where(x > 0, 2, 0.5 * x ** 2 + x + 2 - (0.3 * x) ** 3 + np.random.normal(0, 1, size=100))

# 用散点图形式展示生成的数据
plt.scatter(x, y)
plt.show()

# 实例化一个线性回归模型
lr = LinearRegression()

# 使用原始特征进行拟合
# lr.fit(X, y)  # 如果不想展示线性回归结果，可以注释掉这行和下面两行

# 得到模型预测值
# y_predict = lr.predict(X)  # 同上，注释掉

# 在图上展示原始数据和线性模型的拟合效果
# plt.scatter(x, y)  # 同上，注释掉
# plt.plot(x, y_predict, color='red')  # 同上，注释掉
# plt.show()  # 同上，注释掉

# 将原始特征转换为3次多项式特征
poly = PolynomialFeatures(degree=3)
poly.fit(X)
X_poly = poly.transform(X)

# 使用3次多项式特征进行拟合
lr.fit(X_poly, y)

# 得到模型预测值
y_predict_poly = lr.predict(X_poly)

# 只在图上展示原始数据和多项式模型的拟合效果
plt.scatter(x, y)
plt.plot(np.sort(x), y_predict_poly[np.argsort(x)], color='red')
plt.show()
print(x)
print(y)

# 预测结果与2比较，产生布尔数组
is_greater_or_equal_to_two = y_predict_poly >= 2
# 打印布尔数组
print(is_greater_or_equal_to_two)
# 如果您想要在布尔数组中看到对应的x和y的值，可以这样打印
for i in range(len(x)):
    print(f"x: {x[i]}, y: {y_predict_poly[i]}, >=2: {is_greater_or_equal_to_two[i]}")

python3.6 环境