Python实现河流图-CFANZ编程社区

以下是进一步优化后的代码，从色彩搭配、元素细节、代码结构等方面进行改进，让生成的小河图更加美观、生动，同时代码也更具可读性：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.patches import Ellipse, Rectangle


# 数据生成函数
def generate_data(years=6, base=10):
    data = []
    for year in range(years):
        seasonal = np.sin(np.linspace(0, 2 * np.pi, 12)) * 0.3 + 1
        trend = base * (1.3 ** year)
        values = trend * seasonal * np.random.normal(1, 0.1, 12)
        data.extend(values.tolist())
        base = trend
    return np.array(data)


# 绘制河流函数
def draw_river(data, ax):
    segments = []
    widths = []
    alphas = []
    max_value = data.max()
    for i in range(len(data) - 1):
        seg = [(i, 0), (i + 1, 0)]
        segments.append(seg)
        widths.append(0.3 + (data[i] / max_value) * 8)
        alphas.append(0.2 + 0.7 * (data[i] / max_value))

    cmap = plt.get_cmap('Greens_r')
    colors = [cmap(a) for a in alphas]
    lc = LineCollection(segments,
                        linewidths=widths,
                        colors=colors,
                        capstyle='round',
                        alpha=0.8)
    ax.add_collection(lc)


# 添加波纹效果函数
def add_wave_effect(ax, data, widths):
    for i in range(0, len(data), 4):
        y = np.random.uniform(-0.1, 0.1)
        ax.plot([i, i + 1], [y, y], color='lightgreen', alpha=0.4, linewidth=widths[i] * 0.2)


# 添加岸边草地函数
def add_grass(ax, years):
    grass_color = '#7CFC00'
    for i in range(0, 12 * years):
        grass_height = np.random.uniform(0.1, 0.3)
        ax.add_patch(Rectangle((i, 0.8), 1, grass_height, color=grass_color, alpha=0.7))
        ax.add_patch(Rectangle((i, -0.8 - grass_height), 1, grass_height, color=grass_color, alpha=0.7))


# 添加石头函数
def add_stones(ax, years):
    stone_color = '#555555'
    num_stones = 15
    for _ in range(num_stones):
        x = np.random.uniform(0, 12 * years)
        y = np.random.uniform(-0.8, 0.8)
        width = np.random.uniform(0.2, 0.6)
        height = np.random.uniform(0.2, 0.6)
        ax.add_patch(Ellipse((x, y), width, height, color=stone_color, alpha=0.5))


# 添加光影效果函数
def add_light_shadow_effect(ax, data, widths):
    light_color = '#90EE90'
    dark_color = '#008000'
    for i in range(len(data)):
        if i % 2 == 0:
            ax.plot([i, i + 1], [0, 0], color=light_color, linewidth=widths[i] * 0.7, alpha=0.2)
        else:
            ax.plot([i, i + 1], [0, 0], color=dark_color, linewidth=widths[i] * 0.7, alpha=0.2)


# 添加小鱼函数
def add_fishes(ax, years):
    fish_color = 'orange'
    num_fishes = 8
    for _ in range(num_fishes):
        x = np.random.uniform(0, 12 * years)
        y = np.random.uniform(-0.6, 0.6)
        ax.plot([x, x + 0.3], [y, y + 0.1], color=fish_color, linewidth=2)
        ax.plot([x, x + 0.3], [y, y - 0.1], color=fish_color, linewidth=2)


# 添加时间刻度函数
def add_time_ticks(ax, years):
    for year in range(years):
        x = year * 12 + 6
        ax.text(x, -1.2, f'{2020 + year}',
                ha='center', va='top',
                color='#2f4f4f', fontstyle='italic')


# 主函数
def main():
    years = 6
    data = generate_data(years)
    fig, ax = plt.subplots(figsize=(36, 2), facecolor='#f0f8ff')
    ax.set_xlim(0, 12 * years)
    ax.set_ylim(-1.5, 1.5)
    ax.axis('off')

    draw_river(data, ax)
    widths = [0.3 + (d / data.max()) * 8 for d in data[:-1]]
    add_wave_effect(ax, data, widths)
    add_grass(ax, years)
    add_stones(ax, years)
    add_light_shadow_effect(ax, data, widths)
    add_fishes(ax, years)
    add_time_ticks(ax, years)

    plt.tight_layout()
    plt.savefig('optimized_river.jpg', dpi=300, bbox_inches='tight')
    plt.show()


if __name__ == "__main__":
    main()