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 = 600 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, 6), 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_u = 0.0 self.reward_v = 0.0 self.reward_sim = 0.0 self.current_step = 0 self.reset_cont = 0 self.weight_r = [0.3, 0.7, 0.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.time_delay = int(18 * L0 / U0 / SAMPLE_INTERVAL) self.time_delay = 63 self.ddf_ave = np.zeros((NX, NY, 9), dtype=DATA_TYPE) self.ddf_ave_cont = 0 center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 + 2 * L0, 0) self.flow_field.add_sensor(center, L0 / 4) center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2, 0) self.flow_field.add_sensor(center, L0 / 4) center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 - 2 * L0, 0) self.flow_field.add_sensor(center, L0 / 4) self.flow_field.run(int(2*NX/U0), np.zeros(3, dtype=DATA_TYPE)) for i in range(FIFO_LEN): self.flow_field.run(SAMPLE_INTERVAL, np.zeros(3, dtype=DATA_TYPE)) new_state = self.flow_field.obs.copy()[0:6] self.target_states = np.vstack((self.target_states, new_state)) self.target_sensors = np.mean(self.target_states, axis=0) # self.flow_field.apply_ddf() center: Tuple[float, float, float] = (10 * L0, (NY - 1) / 2, 0) self.flow_field.add_cylinder(center, L0 / 2) center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0) self.flow_field.add_cylinder(center, L0 / 2) center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 + 0.75 * L0, 0) self.flow_field.add_cylinder(center, L0 / 2) center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 - 0.75 * L0, 0) self.flow_field.add_cylinder(center, L0 / 2) center: Tuple[float, float, float] = (32 * L0, (NY - 1) / 2, 0) self.flow_field.add_sensor(center, L0 / 4) self.flow_field.run(int(4*NX/U0), np.zeros(8, 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(8, dtype=DATA_TYPE)) self.fifo_states.append(self.flow_field.obs.copy()[0:16]) temp_states = np.array(self.fifo_states) self.force_norm_fact = 50 * np.max(np.abs(temp_states[:, 6:14])) for i in range(6): self.sens_deviation[i] = np.mean(temp_states[:, i]) self.sens_norm_fact[i] = 10 * np.max(np.abs(temp_states[:, i] - self.target_sensors[i])) self.sens_norm_fact = np.max(self.sens_norm_fact) 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, 0.0, -8*U0, 8*U0, 0.0], dtype=DATA_TYPE)) self.fifo_states.append(self.flow_field.obs.copy()[0:16]) self.save_states = self.fifo_states.copy() self.flow_field.apply_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(8, dtype=DATA_TYPE) temp[4:7] = np.array((action*8+[0,-8,8])*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:16]) def proc_data(): states = np.array(self.fifo_states) forces = states[-1, 6:14] / self.force_norm_fact forces_delay = states[-1-self.time_delay, 6:14] / self.force_norm_fact cd = (forces[2] + forces[4] + forces[6]) / 3 cl = (forces[3] + forces[5] + forces[7]) / 3 sens = (states[-1, 0:6] - self.target_sensors) / self.sens_norm_fact sens_near = states[-1, 15] / 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): # 计算幅值差异权重 target_std = np.std(target) if np.std(target) > 1e-8 else 1e-8 state_std = np.std(state) if np.std(state) > 1e-8 else 1e-8 amplitude_ratio = min(target_std, state_std) / max(target_std, state_std) # 计算均值差异 mean_diff = abs(np.mean(target) - np.mean(state)) max_scale = max(abs(np.mean(target)), abs(np.mean(state)), 1e-8) mean_similarity = 1 / (1 + mean_diff / max_scale * 10) # DTW计算 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 # 改进的归一化方法 max_possible_cost = max(np.max(np.abs(target)), np.max(np.abs(state)), 1e-8) dtw_distance = dtw_matrix[n, m] / (len(target) * max_possible_cost) DTW_similarity = max(0, 1 - dtw_distance) # 综合相似度：形状相似度 * 幅值相似度 * 均值相似度 total_similarity = 0.8 * DTW_similarity + 0.1 * amplitude_ratio + 0.1 * mean_similarity return total_similarity # target_seq = -states[CONV_LEN:2*CONV_LEN, 7] # state_seq = states[-CONV_LEN:, 9] # lag = calc_lag(target_seq, state_seq) # for i in range(0, 2): # target_seq = -np.roll(states[:, i+6], -lag)[CONV_LEN:2*CONV_LEN] # state_seq = states[-CONV_LEN:, i+8] + states[-CONV_LEN:, i+10] + states[-CONV_LEN:, i+12] # similarities += calc_sim(target_seq, state_seq) / 2 diff_u = (np.abs(sens[0]) + np.abs(sens[2]) + np.abs(sens[4]))/3 # diff_v = (np.abs(sens[1]) + np.abs(sens[3]) + np.abs(sens[5]))/3 diff_v = 0 for i in range(1, 19): diff_v += 1/(3.15*i**1.2) * (np.abs(states[-i, 1] - self.target_sensors[1]) + np.abs(states[-i, 3] - self.target_sensors[3]) + np.abs(states[-i, 5] - self.target_sensors[5])) / self.sens_norm_fact / 3 amp_v = np.std(states[-36:, 15]) / self.sens_norm_fact # diff_near = np.abs(sens_near) self.reward_u = np.exp(-np.abs(diff_u * 70)) self.reward_v = 0.5 * np.exp(-np.abs(amp_v * 70)) + 0.5 * np.exp(-np.abs(diff_v * 70)) self.reward_sim = 0.4*np.exp(-140*np.abs(forces_delay[0]+forces[2]+forces[4]+forces[6])) + 0.6*np.exp(-140*np.abs(forces_delay[1]+forces[3]+forces[5]+forces[7])) # self.reward_sim = similarities reward = np.minimum(self.weight_r[0] * self.reward_u + self.weight_r[1] * self.reward_v + self.weight_r[2] * self.reward_sim, 1.0) result_queue.put((np.hstack([forces[0:8], sens]), 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 self.reset_cont += 1 # if self.reset_cont % 10 == 0: # weight = np.array([[0.6, 0.3, 0.1], [0.3, 0.6, 0.1], [0.3, 0.3, 0.4], [0.8, 0.1, 0.1], [0.1, 0.8, 0.1]]) # self.weight_r = weight[np.random.randint(0, 5)].tolist() # print(f"Reset count: {self.reset_cont}, weight: {self.weight_r}") 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__()