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__()