-
Notifications
You must be signed in to change notification settings - Fork 0
/
spin-decomp.py
150 lines (129 loc) · 4.05 KB
/
spin-decomp.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
# pyright: basic
import glob
import shutil
import os
import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt
from PIL import Image
from scipy.ndimage import laplace
SEED: int = 42 # change this for different outcomes
class CahnHilliard:
def __init__(
self,
N: int = 100,
p: float = 0,
D: float = 1,
gamma: float = 0.5,
num_iter: int = 5000,
dt: float = 0.001,
frame_iter: int = 10,
) -> None:
self.N = N
self.p = p
self.D = D
self.gamma = gamma
self.num_iter = num_iter
self.dt = dt
self.frame_iter = frame_iter
# Cahn-Hilliard equation
def CH_eqn(
self, c: npt.NDArray[np.float64], D: float, gamma: float
) -> npt.NDArray[np.float64]:
return D * laplace(
input=(c**3) - c - (gamma * laplace(input=c, mode="wrap")), mode="wrap"
)
# Initialize random data on a 2d square lattice
def initialize_lattice(self) -> npt.NDArray[np.float64]:
rng = np.random.default_rng(seed=SEED)
return rng.choice(
[-1, 1], (self.N, self.N), p=[0.5 + (self.p / 2), 0.5 - (self.p / 2)]
)
# forward Euler method for solving ODEs
def forward_euler(self, x, func, dt, *fargs) -> npt.NDArray[np.float64]:
return x + func(*fargs) * dt
# solver function
def solve(self) -> npt.NDArray[np.float64]:
c: npt.NDArray[np.float64] = self.initialize_lattice()
t = 0
_, ax = plt.subplots()
self.plot_lattice(c, ax, t, i=0)
for i in range(1, self.num_iter + 1):
fargs: tuple[npt.NDArray[np.float64], float, float] = (
c,
self.D,
self.gamma,
)
ax.clear()
c: npt.NDArray[np.float64] = self.forward_euler(
c, self.CH_eqn, self.dt, *fargs
)
t = np.round(t + self.dt, 6)
if i % self.frame_iter == 0:
self.plot_lattice(c, ax, t, i)
return c
# create snapshot at time t
def plot_lattice(self, c, ax, t, i) -> None:
ax.imshow(c)
ax.set_xticks([], [])
ax.set_yticks([], [])
ax.annotate(
f"$t = {t:.3f}$",
xy=(0.05, 0.95),
xycoords="axes fraction",
bbox={
"color": "white",
"alpha": 0.8,
"boxstyle": "Round",
"edgecolor": None,
},
ha="left",
va="top",
fontsize=14,
)
plt.tight_layout()
# save snapshot
if not os.path.exists("gifs/tmp"):
os.makedirs("gifs/tmp")
plt.savefig(f"gifs/tmp/{i:06d}.png")
# create animation from snapshots
def animate(self) -> None:
path_in: str = "gifs/tmp/*.png"
path_out: str = f"gifs/spin-decomp-D_{self.D}-gamma_{self.gamma}-p_{self.p}.gif"
imgs: list[Image.Image] = []
# grab all snapshots
for f in sorted(glob.glob(path_in)):
img: Image.Image = Image.open(f)
imgs.append(img.copy())
img.close()
# convert snapshots to GIF
imgs[0].save(
fp=path_out,
format="GIF",
append_images=imgs[1:],
save_all=True,
optimize=True,
duration=1,
loop=0,
)
# delete snapshots
if os.path.exists("gifs/tmp"):
shutil.rmtree("gifs/tmp/")
def main(
N: int, p: float, D: float, gamma: float, num_iter: int, dt: float, frame_iter: int
) -> None:
ch: CahnHilliard = CahnHilliard(N, p, D, gamma, num_iter, dt, frame_iter)
ch.solve()
ch.animate()
if __name__ == "__main__":
N: int = 500
p: float = 0.2
D: float = 200
gamma: float = 0.5
num_iter: int = 10000
dt: float = 0.0002
frame_iter: int = 100
print("Beginning calculation ...\n")
print(f"N: {N}\np: {p}\nD: {D}\ngamma: {gamma}\ndt: {dt}\n")
main(N, p, D, gamma, num_iter, dt, frame_iter)
print("\nDone!")