315 lines
13 KiB
Python
315 lines
13 KiB
Python
import gymnasium as gym
|
|
import numpy as np
|
|
from gymnasium import spaces
|
|
import ctypes
|
|
from collections import deque
|
|
from typing import Tuple
|
|
import sys
|
|
import os
|
|
import matplotlib.pyplot as plt
|
|
import queue
|
|
|
|
os.environ["OMP_NUM_THREADS"] = "1"
|
|
os.environ["MKL_NUM_THREADS"] = "1"
|
|
|
|
current_dir = os.path.dirname(os.path.abspath("__file__"))
|
|
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
|
|
sys.path.append(parent_dir)
|
|
from CelerisLab import FlowField
|
|
from CelerisLab import utils
|
|
|
|
config_cuda = utils.load_cuda_config(
|
|
os.path.join(parent_dir, "configs", "config_cuda.json")
|
|
)
|
|
config_field = utils.load_flow_field_config(
|
|
os.path.join(parent_dir, "configs", "config_flowfield.json")
|
|
)
|
|
|
|
S_DIM, A_DIM = 14, 3
|
|
U0 = config_field.velocity
|
|
T0 = 1000
|
|
SAMPLE_INTERVAL = 600
|
|
FIFO_LEN = 150
|
|
CONV_LEN = 36
|
|
MAX_STEPS = 500
|
|
if config_field.data_type == "FP32":
|
|
DATA_TYPE = np.float32
|
|
else:
|
|
raise ValueError(f"Unsupported data type {config_field.data_type}.")
|
|
|
|
|
|
class CustomEnv(gym.Env):
|
|
"""Custom Environment that follows gym interface."""
|
|
|
|
metadata = {"render_modes": ["human"], "render_fps": T0 / SAMPLE_INTERVAL}
|
|
|
|
def __init__(self, device_id=0):
|
|
super().__init__()
|
|
self.action_space = spaces.Box(low=-1, high=1, shape=(A_DIM,), dtype=DATA_TYPE)
|
|
self.observation_space = spaces.Box(
|
|
low=-1, high=1, shape=(S_DIM,), dtype=DATA_TYPE
|
|
)
|
|
self.fifo_states = deque(maxlen=FIFO_LEN)
|
|
self.target_states = np.empty((0, 8), dtype=DATA_TYPE)
|
|
self.force_norm_fact = 1.0
|
|
self.sens_norm_fact = np.ones(6, dtype=DATA_TYPE)
|
|
self.sens_deviation = np.zeros(6, dtype=DATA_TYPE)
|
|
self.reward_cd = 0.0
|
|
self.reward_cl = 0.0
|
|
self.reward_sim = 0.0
|
|
self.current_step = 0
|
|
|
|
self.flow_field = FlowField(config_field, config_cuda, device_id)
|
|
L0 = 20
|
|
U0 = config_field.velocity
|
|
NX = self.flow_field.FIELD_SHAPE[0]
|
|
NY = self.flow_field.FIELD_SHAPE[1]
|
|
|
|
self.ddf_ave = np.zeros((NX, NY, 9), dtype=DATA_TYPE)
|
|
self.ddf_ave_cont = 0
|
|
|
|
center: Tuple[float, float, float] = (20 * L0, (NY - 1) / 2, 0)
|
|
self.flow_field.add_cylinder(center, 1.5*L0)
|
|
center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2 + 2 * L0, 0)
|
|
self.flow_field.add_sensor(center, L0 / 4)
|
|
center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)
|
|
self.flow_field.add_sensor(center, L0 / 4)
|
|
center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2 - 2 * L0, 0)
|
|
self.flow_field.add_sensor(center, L0 / 4)
|
|
self.flow_field.run(int(4*NX/U0), np.zeros(4, dtype=DATA_TYPE))
|
|
|
|
for i in range(FIFO_LEN):
|
|
self.flow_field.run(SAMPLE_INTERVAL, np.zeros(4, dtype=DATA_TYPE))
|
|
new_state = self.flow_field.obs.copy()[0:8]
|
|
self.target_states = np.vstack((self.target_states, new_state))
|
|
|
|
def analyze_harmonics(states, n_harmonics):
|
|
N, D = states.shape
|
|
result = []
|
|
for d in range(D):
|
|
y = states[:, d]
|
|
fft_coef = np.fft.rfft(y)
|
|
freqs = np.fft.rfftfreq(N, d=1)
|
|
amps = 2 * np.abs(fft_coef) / N
|
|
phases = np.angle(fft_coef)
|
|
idx = np.argsort(amps[1:])[::-1][:n_harmonics] + 1
|
|
harmonics = {
|
|
'dc': np.real(fft_coef[0]) / N,
|
|
'amps': amps[idx],
|
|
'freqs': freqs[idx],
|
|
'phases': phases[idx]
|
|
}
|
|
result.append(harmonics)
|
|
return result
|
|
|
|
self.target_harmonics = analyze_harmonics(self.target_states, n_harmonics=5)
|
|
|
|
del self.flow_field
|
|
self.flow_field = FlowField(config_field, config_cuda, device_id)
|
|
|
|
center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2 + 2 * L0, 0)
|
|
self.flow_field.add_sensor(center, L0 / 4)
|
|
center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)
|
|
self.flow_field.add_sensor(center, L0 / 4)
|
|
center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2 - 2 * L0, 0)
|
|
self.flow_field.add_sensor(center, L0 / 4)
|
|
center: Tuple[float, float, float] = (19 * L0, (NY - 1) / 2, 0)
|
|
self.flow_field.add_cylinder(center, L0 / 2)
|
|
center: Tuple[float, float, float] = (20.3 * L0, (NY - 1) / 2 + 0.75 * L0, 0)
|
|
self.flow_field.add_cylinder(center, L0 / 2)
|
|
center: Tuple[float, float, float] = (20.3 * L0, (NY - 1) / 2 - 0.75 * L0, 0)
|
|
self.flow_field.add_cylinder(center, L0 / 2)
|
|
self.flow_field.run(int(4*NX/U0), np.zeros(6, dtype=DATA_TYPE))
|
|
self.flow_field.get_ddf()
|
|
self.flow_field.save_ddf()
|
|
|
|
for i in range(FIFO_LEN):
|
|
self.flow_field.run(SAMPLE_INTERVAL, np.zeros(6, dtype=DATA_TYPE))
|
|
self.fifo_states.append(self.flow_field.obs.copy()[0:12])
|
|
|
|
temp_states = np.array(self.fifo_states)
|
|
self.force_norm_fact = 6 * np.max(np.abs(temp_states[:, 6:12]))
|
|
|
|
for i in range(6):
|
|
self.sens_deviation[i] = np.mean(temp_states[:, i])
|
|
self.sens_norm_fact[i] = 5 * np.max(np.abs(temp_states[:, i] - self.sens_deviation[i]))
|
|
|
|
self.flow_field.apply_ddf()
|
|
|
|
for i in range(FIFO_LEN):
|
|
self.flow_field.run(SAMPLE_INTERVAL, np.array([0.0, 0.0, 0.0, 0.0, -1*U0, 1*U0], dtype=DATA_TYPE))
|
|
self.fifo_states.append(self.flow_field.obs.copy()[0:12])
|
|
|
|
self.save_states = self.fifo_states.copy()
|
|
self.flow_field.get_ddf()
|
|
self.flow_field.save_ddf()
|
|
|
|
def step(self, action):
|
|
assert self.action_space.contains(action), "%r (%s) invalid" % (
|
|
action,
|
|
type(action),
|
|
)
|
|
|
|
# barrier = threading.Barrier(2)
|
|
result_queue = queue.Queue()
|
|
|
|
def run_flow_field(action):
|
|
self.flow_field.context.push()
|
|
U0 = config_field.velocity
|
|
try:
|
|
temp = np.zeros(6, dtype=DATA_TYPE)
|
|
temp[3:6] = np.array((action*8+[0,-2,2])*U0, dtype=DATA_TYPE)
|
|
self.flow_field.run(SAMPLE_INTERVAL, temp)
|
|
finally:
|
|
self.flow_field.context.pop()
|
|
# barrier.wait()
|
|
self.fifo_states.append(self.flow_field.obs.copy()[0:12])
|
|
|
|
def proc_data():
|
|
states = np.array(self.fifo_states)
|
|
forces = states[-1, 6:12] / self.force_norm_fact
|
|
cd = forces[0] + forces[2] + forces[4]
|
|
cl = forces[1] + forces[3] + forces[5]
|
|
sens = (states[-1, 0:6] - self.sens_deviation) / self.sens_norm_fact
|
|
|
|
similarities = 0.0
|
|
|
|
def calc_lag(target, state):
|
|
target_mean = np.mean(target)
|
|
state_mean = np.mean(state)
|
|
|
|
correlation = np.correlate(target - target_mean, state - state_mean, "full")
|
|
lags = np.arange(-len(target) + 1, len(target))
|
|
max_lag = lags[np.argmax(correlation)]
|
|
return max_lag
|
|
|
|
def calc_sim(target, state):
|
|
|
|
n = len(target)
|
|
m = len(state)
|
|
|
|
dtw_matrix = np.full((n + 1, m + 1), np.inf)
|
|
dtw_matrix[0, 0] = 0
|
|
|
|
for i in range(1, n + 1):
|
|
for j in range(1, m + 1):
|
|
cost = abs(target[i - 1] - state[j - 1])
|
|
last_min = min(dtw_matrix[i - 1, j],
|
|
dtw_matrix[i, j - 1],
|
|
dtw_matrix[i - 1, j - 1])
|
|
dtw_matrix[i, j] = cost + last_min
|
|
|
|
return 1 - (dtw_matrix[n, m] / len(target))
|
|
|
|
def gen_target_states_at(t, harmonics):
|
|
t = np.asarray(t)
|
|
D = len(harmonics)
|
|
result = np.zeros((t.size, D), dtype=np.float32)
|
|
for d, h in enumerate(harmonics):
|
|
val = np.full(t.shape, h['dc'], dtype=np.float32)
|
|
for amp, freq, phase in zip(h['amps'], h['freqs'], h['phases']):
|
|
val += amp * np.cos(2 * np.pi * freq * t + phase)
|
|
result[:, d] = val
|
|
if result.shape[0] == 1:
|
|
return result[0]
|
|
return result
|
|
|
|
id_sens = 1
|
|
target_seq = self.target_states[CONV_LEN:2*CONV_LEN, id_sens+2]
|
|
state_seq = states[-CONV_LEN:, id_sens]
|
|
lag = calc_lag(target_seq, state_seq)
|
|
|
|
for i in range(0, 6):
|
|
target_seq = np.roll(self.target_states[:, i+2], -lag)[CONV_LEN:2*CONV_LEN]
|
|
state_seq = states[-CONV_LEN:, i]
|
|
similarities += calc_sim(target_seq, state_seq) / 6
|
|
|
|
target_states = gen_target_states_at(self.current_step, self.target_harmonics)
|
|
target_cd = target_states[0] / self.force_norm_fact
|
|
target_cl = target_states[1] / self.force_norm_fact
|
|
|
|
self.reward_cd = np.exp(-np.abs((cd-target_cd) * 10))
|
|
self.reward_cl = np.exp(-np.abs((cl-target_cl) * 10))
|
|
self.reward_sim = np.exp(-10*np.abs(similarities - 1))
|
|
reward = np.minimum(0.3 * self.reward_cd + 0.3 * self.reward_cl + 0.4 * self.reward_sim, 1.0)
|
|
result_queue.put((np.hstack([forces, sens, target_cd, target_cl]), reward))
|
|
|
|
run_flow_field(action)
|
|
proc_data()
|
|
observation, reward = result_queue.get()
|
|
|
|
truncated = bool(np.any(observation > 1) or np.any(observation < -1))
|
|
observation = np.clip(observation, -1, 1)
|
|
self.current_step += 1
|
|
# done = self.current_step >= MAX_STEPS
|
|
done = False
|
|
return observation, float(reward), done, truncated, {}
|
|
|
|
def reset(self, seed=None):
|
|
self.flow_field.restore_ddf()
|
|
self.flow_field.apply_ddf()
|
|
self.fifo_states = self.save_states.copy()
|
|
self.current_step = 0
|
|
return np.zeros(S_DIM, dtype=np.float32), {}
|
|
|
|
def render(self, mode="human"):
|
|
NX = self.flow_field.FIELD_SHAPE[0]
|
|
NY = self.flow_field.FIELD_SHAPE[1]
|
|
self.flow_field.get_ddf()
|
|
ddf_plot = self.flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)
|
|
ux = (ddf_plot[:, :, 1] + ddf_plot[:, :, 5] + ddf_plot[:, :, 8] - ddf_plot[:, :, 3] - ddf_plot[:, :, 6] - ddf_plot[:, :, 7]) / U0
|
|
uy = (ddf_plot[:, :, 2] + ddf_plot[:, :, 5] + ddf_plot[:, :, 6] - ddf_plot[:, :, 4] - ddf_plot[:, :, 7] - ddf_plot[:, :, 8]) / U0
|
|
speed = np.sqrt(ux**2 + uy**2)
|
|
plt.figure(figsize=(10, 5))
|
|
plt.imshow(speed.T, origin='lower', cmap='viridis', extent=[0, NX, 0, NY])
|
|
plt.colorbar(label='Speed')
|
|
plt.title('Scalar Velocity Field')
|
|
plt.xlabel('X')
|
|
plt.ylabel('Y')
|
|
plt.tight_layout()
|
|
plt.show()
|
|
|
|
def save_field(self, filename):
|
|
NX = self.flow_field.FIELD_SHAPE[0]
|
|
NY = self.flow_field.FIELD_SHAPE[1]
|
|
self.flow_field.get_ddf()
|
|
ddf_plot = self.flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)
|
|
flag_plot = self.flow_field.flag.copy().reshape((NY, NX)).transpose(1, 0)
|
|
ux = (ddf_plot[:, :, 1] + ddf_plot[:, :, 5] + ddf_plot[:, :, 8] - ddf_plot[:, :, 3] - ddf_plot[:, :, 6] - ddf_plot[:, :, 7]) / U0
|
|
uy = (ddf_plot[:, :, 2] + ddf_plot[:, :, 5] + ddf_plot[:, :, 6] - ddf_plot[:, :, 4] - ddf_plot[:, :, 7] - ddf_plot[:, :, 8]) / U0
|
|
with open(os.path.join(parent_dir, "output", filename), "w") as f:
|
|
f.write("Title= \"LBM 2D\"\r\n")
|
|
f.write("VARIABLES= \"X\",\"Y\",\"flag\",\"U\",\"V\",\r\n")
|
|
f.write(f"ZONE T= \"BOX\",I= {NX},J= {NY},F=POINT\r\n")
|
|
for j in range(NY):
|
|
for i in range(NX):
|
|
f.write(f"{i},{j},{flag_plot[i, j]},{ux[i, j]},{uy[i, j]}\r\n")
|
|
|
|
def average_field(self, mode=["add", "save", "clear"], filename="average_field.dat"):
|
|
NX = self.flow_field.FIELD_SHAPE[0]
|
|
NY = self.flow_field.FIELD_SHAPE[1]
|
|
self.flow_field.get_ddf()
|
|
ddf_new = self.flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)
|
|
if "add" in mode:
|
|
self.ddf_ave = self.ddf_ave + ddf_new
|
|
self.ddf_ave_cont += 1
|
|
if "save" in mode:
|
|
if self.ddf_ave_cont == 0:
|
|
raise ValueError("No data to save. Please run 'add' mode first.")
|
|
ux = (self.ddf_ave[:, :, 1] + self.ddf_ave[:, :, 5] + self.ddf_ave[:, :, 8] - self.ddf_ave[:, :, 3] - self.ddf_ave[:, :, 6] - self.ddf_ave[:, :, 7]) / U0 / self.ddf_ave_cont
|
|
uy = (self.ddf_ave[:, :, 2] + self.ddf_ave[:, :, 5] + self.ddf_ave[:, :, 6] - self.ddf_ave[:, :, 4] - self.ddf_ave[:, :, 7] - self.ddf_ave[:, :, 8]) / U0 / self.ddf_ave_cont
|
|
flag_plot = self.flow_field.flag.copy().reshape((NY, NX)).transpose(1, 0)
|
|
with open(os.path.join(parent_dir, "output", filename), "w") as f:
|
|
f.write("Title= \"LBM 2D\"\r\n")
|
|
f.write("VARIABLES= \"X\",\"Y\",\"flag\",\"U\",\"V\",\r\n")
|
|
f.write(f"ZONE T= \"BOX\",I= {NX},J= {NY},F=POINT\r\n")
|
|
for j in range(NY):
|
|
for i in range(NX):
|
|
f.write(f"{i},{j},{flag_plot[i, j]},{ux[i, j]},{uy[i, j]}\r\n")
|
|
print(f"Average field amount: {self.ddf_ave_cont}")
|
|
if "clear" in mode:
|
|
self.ddf_ave = np.zeros((NX, NY, 9), dtype=DATA_TYPE)
|
|
self.ddf_ave_cont = 0
|
|
|
|
def close(self):
|
|
self.flow_field.__del__() |