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 = 12, 3 U0 = config_field.velocity T0 = 1000 SAMPLE_INTERVAL = 800 FIFO_LEN = 120 CONV_LEN = 60 MAX_STEPS = 720 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.fifo_target = deque(maxlen=FIFO_LEN) self.fifo_forces = deque(maxlen=FIFO_LEN) 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_sim_now = 0.0 self.reward_yaw = 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] center: Tuple[float, float, float] = (10 * L0, (NY - 1) / 2, 0) self.flow_field.add_cylinder(center, L0) center: Tuple[float, float, float] = (25 * L0, (NY - 1) / 2 + 2 * L0, 0) self.flow_field.add_sensor(center, L0 / 4) center: Tuple[float, float, float] = (25 * L0, (NY - 1) / 2, 0) self.flow_field.add_sensor(center, L0 / 4) center: Tuple[float, float, float] = (25 * 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 + 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(4*NX/U0), np.zeros(7, dtype=DATA_TYPE)) for i in range(FIFO_LEN): self.flow_field.run(SAMPLE_INTERVAL, np.zeros(7, dtype=DATA_TYPE)) self.fifo_target.append(self.flow_field.obs.copy()[2:8]) self.fifo_states.append(self.flow_field.obs.copy()[8:14]) # target = np.array(self.fifo_target)[:, 0] # state = np.array(self.fifo_states)[:, 0] # 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)) # self.LAG = lags[np.argmax(correlation)] self.LAG = -9 self.flow_field.apply_ddf() 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) self.flow_field.run(int(4*NX/U0), np.zeros(10, dtype=DATA_TYPE)) self.flow_field.get_ddf() for i in range(FIFO_LEN): self.flow_field.run(SAMPLE_INTERVAL, np.zeros(10, dtype=DATA_TYPE)) self.fifo_target.append(self.flow_field.obs.copy()[2:8]) self.fifo_states.append(self.flow_field.obs.copy()[8:14]) self.fifo_forces.append(self.flow_field.obs.copy()[14:20]) self.save_target = self.fifo_target.copy() self.save_states = self.fifo_states.copy() self.save_forces = self.fifo_forces.copy() self.flow_field.get_ddf() temp_states = np.array(self.fifo_states) self.force_norm_fact = 6 * np.max(np.abs(np.array(self.fifo_forces))) 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])) 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(10, dtype=DATA_TYPE) temp[7:10] = np.array((action*8+[0,-4,4])*U0, dtype=DATA_TYPE) self.flow_field.run(SAMPLE_INTERVAL, temp) finally: self.flow_field.context.pop() # barrier.wait() self.fifo_target.append(self.flow_field.obs.copy()[2:8]) self.fifo_states.append(self.flow_field.obs.copy()[8:14]) self.fifo_forces.append(self.flow_field.obs.copy()[14:20]) def proc_data(): target = np.array(self.fifo_target) states = np.array(self.fifo_states) forces = np.array(self.fifo_forces)[-1, :] / self.force_norm_fact cd = (forces[0] + forces[2] + forces[4]) / 3 cl = (forces[1] + forces[3] + forces[5]) / 3 ave_v = np.mean(states[:, 1] + states[:, 3] + states[:, 5]) / 3 targ = (target[-1, :] - self.sens_deviation) / self.sens_norm_fact sens = (states[-1, :] - self.sens_deviation) / self.sens_norm_fact similarities = 0.0 sim_now = 0.0 def calc_sim(target, state, lag): target_mean = np.mean(target) state_mean = np.mean(state) target_std = np.std(target) aligned_state = np.roll(state, lag) if lag >= 0: seq_target = target[-CONV_LEN:]-target_mean seq_state = aligned_state[-CONV_LEN:]-state_mean else: seq_target = target[:CONV_LEN]-target_mean seq_state = aligned_state[:CONV_LEN]-state_mean seq_diff = seq_target - seq_state sim_cor = 10*(np.corrcoef(seq_target, seq_state)[0, 1] - 1) sim_div = -np.abs((target_mean - state_mean) / target_std * 0.75) sim_amp = -np.abs(np.std(seq_diff) / target_std * 2) return np.exp((sim_cor + sim_div + sim_amp) / 3) for i in range(0, 6): target_seq = (target[:, i] - self.sens_deviation[i]) / self.sens_norm_fact[i] state_seq = (states[:, i] - self.sens_deviation[i]) / self.sens_norm_fact[i] similarities += calc_sim(target_seq, state_seq, -self.LAG) / 6 sim_now += np.abs(target_seq[self.LAG-1] - state_seq[-1]) / 6 self.reward_sim_now = np.exp(-sim_now*10) self.reward_yaw = 1 - np.exp(-np.abs(ave_v * 10)) self.reward_sim = similarities reward = np.clip(0.5 * self.reward_sim_now - 0.3 * self.reward_yaw + 0.8 * self.reward_sim, 0, 1) # barrier.wait() result_queue.put((np.hstack([targ, 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 return observation, float(reward), done, truncated, {} def reset(self, seed=None): self.flow_field.apply_ddf() self.fifo_target = self.save_target.copy() self.fifo_states = self.save_states.copy() self.fifo_forces = self.save_forces.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 close(self): self.flow_field.__del__()