【Python×matplotlib×ChatGPT】稲妻（雷）を描きたい

今日も、こんにちは！
前回は、ChatGPTと稲妻を描こうとしたらこんな感じになってしまいました。

全然稲妻ではありません・・・

今回は稲妻の形を少しでも出したい、ということで以下のサイトを参考にさせていただきました。


emotionexplorer.blog.fc2.com


出来上がったコード

import numpy as np
import matplotlib.pyplot as plt

class DLA:
    def __init__(self, N, view=True, color=True, sharpness=2):
        self.r = 3
        self.N = N
        self.view = view
        self.color = color
        self.sharpness = sharpness
        self.L = int(self.N ** (0.78 + 0.22 * self.sharpness))

        # Update the lattice size to be 5 times larger
        self.lattice = np.zeros([14*self.L+1, 14*self.L+1], dtype=int)
        self.center = 7* self.L
        self.lattice[self.center, self.center] = 1

    def grow_cluster(self):
        rn = np.random.rand

        def reset():
            theta = 2 * np.pi * rn()
            x = int((self.r + 2) * np.cos(theta)) + self.center
            y = int((self.r + 2) * np.sin(theta)) + self.center
            return x, y

        x, y = reset()

        n = 0
        while n < self.N:
            r = np.sqrt((x - self.center) ** 2 + (y - self.center) ** 2)
            if r > self.r + 2:
                l = int(r - self.r - 2)
                if l == 0:
                    l = 1
            else:
                l = 1

            p = rn() * 4
            if p < 1:
                x += l
            elif p < 2:
                x -= l
            elif p < 3:
                y += l
            else:
                y -= l

            r = np.sqrt((x - self.center) ** 2 + (y - self.center) ** 2)

            if r >= 2 * self.r:
                x, y = reset()
                continue

            judge = np.sum(self.lattice[x-1:x+2, y-1:y+2])

            if judge > 0:
                self.lattice[x, y] = 1

                if self.view:
                    if self.color:
                        color = (1, 1, 0)  # カラーを黄色に設定 (RGBでの値は(1, 1, 0))
                    else:
                        color = 'white'
                    size = rn() * 20 + 5  # ランダムな大きさ (5から25までの間)
                    plt.scatter(x - 5*self.L, y - 5*self.L, color=color, s=size)

                if int(r) + 1 > self.r:
                    self.r = int(r) + 1

                x, y = reset()
                n += 1

        if self.view:
            plt.gca().set_facecolor('black')  # 背景を黒に設定
            plt.gca().set_aspect('equal', adjustable='box')
            plt.gca().set_xticks([])
            plt.gca().set_yticks([])
            plt.show()

def main():
    dla = DLA(1000, view=True, color=True, sharpness=1)
    dla.grow_cluster()

if __name__ == '__main__':
    main()

出来上がった図





稲妻に見えます。

次は、色を変えてみました。

よくわからない図になりましたね。

でも、それなりにできたので満足です。