240923

CelerisLab/kernels/macros.h

@@ -9,15 +9,15 @@
 
 // flow parameters
 #define LBtype float
-#define UX 10
-#define UY 16
+#define UX 12
+#define UY 20
 #define UZ 1
-#define NX 1280
-#define NY 512
+#define NX 1536
+#define NY 640
 #define NZ 1
 #define DIM 2
 #define NQ 9
-#define VIS 0.004
+#define VIS 0.006
 #define RHO 1.0
 #define U0 0.01
 

configs/config_flowfield.json

@@ -2,8 +2,8 @@
     "data_type": "FP32",
     "dimensionality": 2,
     "lattice": 9,
-    "field_dim_in_U": [10, 16, 1],
-    "viscosity": 0.004,
+    "field_dim_in_U": [12, 20, 1],
+    "viscosity": 0.006,
     "velocity": 0.01,
     "boundary_conditions": {
         "x": ["parabolic", "outflow"],

scripts/d1a3o12_erase.py

@@ -27,16 +27,16 @@ class Sin(Module):
 if __name__ == '__main__':
     
     vec_env = CustomEnv(device_id=1)
-    name = "d1a3o12_re100_erase_1"
+    name = "d1a3o12_re100_erase_c0"
 
-    model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_re100_erase"), env=vec_env, device=torch.device("cuda:1"))
+    # model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_re100_erase_b0"), env=vec_env, device=torch.device("cuda:1"))
 
-    # model = PPO(
-    #     "MlpPolicy", 
-    #     policy_kwargs=dict(activation_fn=Sin),
-    #     env=vec_env, 
-    #     device=torch.device("cuda:1"),
-    #     verbose=0)
+    model = PPO(
+        "MlpPolicy", 
+        policy_kwargs=dict(activation_fn=Sin),
+        env=vec_env, 
+        device=torch.device("cuda:1"),
+        verbose=0)
 
     writer = SummaryWriter(log_dir=os.path.join(parent_dir, "tensorboard", name))
     max_reward = 0

scripts/d1a3o12_imit.py

@@ -0,0 +1,70 @@
+import os
+os.environ['MKL_THREADING_LAYER'] = 'GNU'
+os.environ["OMP_NUM_THREADS"] = "8"
+os.environ["MKL_NUM_THREADS"] = "8"
+import torch
+import numpy as np
+from torch.nn import Module
+import gymnasium as gym
+from gym_env_imit import CustomEnv
+from stable_baselines3 import PPO
+from stable_baselines3.common.vec_env import SubprocVecEnv
+from stable_baselines3.common.vec_env import DummyVecEnv
+from sb3_contrib import RecurrentPPO
+from torch.utils.tensorboard import SummaryWriter
+import pickle
+
+current_dir = os.path.dirname(os.path.abspath("__file__"))
+parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
+
+class Sin(Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.sin(x)
+
+if __name__ == '__main__':
+    
+    vec_env = CustomEnv(device_id=3)
+    name = "d1a3o12_re100_imit_a1"
+
+    model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_re100_imit_a0"), env=vec_env, device=torch.device("cuda:3"))
+
+    # model = PPO(
+    #     "MlpPolicy", 
+    #     policy_kwargs=dict(activation_fn=Sin),
+    #     env=vec_env, 
+    #     device=torch.device("cuda:3"),
+    #     verbose=0)
+
+    writer = SummaryWriter(log_dir=os.path.join(parent_dir, "tensorboard", name))
+    max_reward = 0
+
+    history_data = []
+
+    for i in range(400):
+        model.learn(total_timesteps=360)
+        test_env = model.get_env()
+        test_obs = test_env.reset()
+        list_reward = []
+        episolde_data = {'actions': [], 'observations': [], 'rewards': []}
+
+        for step in range(360):
+            test_action, _states = model.predict(observation=test_obs)
+            test_obs, test_rewards, test_dones, info = test_env.step(test_action)
+            list_reward.append(test_rewards)
+            episolde_data['actions'].append(test_action[0, :])
+            episolde_data['observations'].append(np.array(test_obs))
+            episolde_data['rewards'].append(test_rewards)
+
+        history_data.append(episolde_data)
+
+        avg_reward = np.mean(list_reward[-180:])
+        writer.add_scalar('Reward', np.mean(avg_reward), i)
+        if avg_reward > max_reward:
+            max_reward = avg_reward
+            model.save(os.path.join(parent_dir, "models", name + ".zip"))
+
+    with open(os.path.join(parent_dir, "output", name + ".pkl"), 'wb') as f:
+        pickle.dump(history_data, f)

scripts/gym_env_erase.py

@@ -28,7 +28,7 @@ config_field = utils.load_flow_field_config(
 S_DIM, A_DIM = 12, 3
 U0 = config_field.velocity
 T0 = 1000
-SAMPLE_INTERVAL = 800
+SAMPLE_INTERVAL = 1200
 FIFO_LEN = 120
 CONV_LEN = 60
 MAX_STEPS = 360
@@ -50,13 +50,14 @@ class CustomEnv(gym.Env):
             low=-1, high=1, shape=(S_DIM,), dtype=DATA_TYPE
         )
         self.fifo_states = deque(maxlen=FIFO_LEN)
+        self.fifo_rewards = deque([0.1] * 10, maxlen=10)
         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.flow_field = FlowField(config_field, config_cuda, device_id)
-        L0 = 20
+        L0 = 30
         U0 = config_field.velocity
         NX = self.flow_field.FIELD_SHAPE[0]
         NY = self.flow_field.FIELD_SHAPE[1]
@@ -80,9 +81,9 @@ class CustomEnv(gym.Env):
         self.flow_field.add_cylinder(center, L0)
         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)
+        center: Tuple[float, float, float] = ((30+1.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)
+        center: Tuple[float, float, float] = ((30+1.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(7, dtype=DATA_TYPE))
@@ -113,7 +114,7 @@ class CustomEnv(gym.Env):
             U0 = config_field.velocity
             try:
                 temp = np.zeros(7, dtype=DATA_TYPE)
-                temp[4:7] = np.array((action*8+[0,-4,4])*U0, dtype=DATA_TYPE)
+                temp[4:7] = np.array((action*6+[0,-6,6])*U0, dtype=DATA_TYPE)
                 self.flow_field.run(SAMPLE_INTERVAL, temp)
             finally:
                 self.flow_field.context.pop()
@@ -134,7 +135,9 @@ class CustomEnv(gym.Env):
             reward_cd = np.exp(-np.abs(cd * 20))
             reward_cl = np.exp(-np.abs(cl * 80))
             reward_sim = np.exp(-np.abs(diff * 20))
-            reward = np.minimum(0.3 * reward_cd + 0.3 * reward_cl + 0.4 * reward_sim, 1.0)
+            reward = np.minimum(0.2 * reward_cd + 0.2 * reward_cl + 0.6 * reward_sim, 1.0)
+            self.fifo_rewards.append(reward)
+            reward = np.mean(self.fifo_rewards)
 
             result_queue.put((np.hstack([forces[2:8], sens]), reward))
 

scripts/gym_env_imit.py

@@ -0,0 +1,237 @@
+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 = 1200
+FIFO_LEN = 120
+CONV_LEN = 60
+MAX_STEPS = 360
+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.flow_field = FlowField(config_field, config_cuda, device_id)
+        L0 = 30
+        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 / 2)
+        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(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()
+            self.target_states = np.vstack((self.target_states, new_state))
+
+        self.flow_field.apply_ddf()
+        center: Tuple[float, float, float] = (10 * L0, (NY - 1) / 2, 0)
+        self.flow_field.add_cylinder(center, L0)
+        center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)
+        self.flow_field.add_cylinder(center, L0 / 2)
+        center: Tuple[float, float, float] = ((30+1.3) * L0, (NY - 1) / 2 + 0.75 * L0, 0)
+        self.flow_field.add_cylinder(center, L0 / 2)
+        center: Tuple[float, float, float] = ((30+1.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(8, dtype=DATA_TYPE))
+        self.flow_field.get_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())
+
+        temp_states = np.array(self.fifo_states)
+        self.force_norm_fact = 6 * np.max(np.abs(temp_states[:, 8:16]))
+        for i in range(6):
+            self.sens_deviation[i] = np.mean(temp_states[:, i+2])
+            self.sens_norm_fact[i] = 5 * np.max(np.abs(temp_states[:, i+2] - self.sens_deviation[i]))
+            self.target_states[:, i+2] = (self.target_states[:, i+2] - self.sens_deviation[i]) / self.sens_norm_fact[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(8, dtype=DATA_TYPE)
+                temp[5:8] = 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_states.append(self.flow_field.obs.copy())
+
+        def proc_data():
+            states = np.array(self.fifo_states)
+            forces = states[-1, 8:16] / self.force_norm_fact
+            sens = (states[-1, 2:8] - self.sens_deviation) / self.sens_norm_fact
+            cd = forces[0] + forces[2] + forces[4] + forces[6]
+            cl = forces[1] + forces[3] + forces[5] + forces[7]
+
+            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, 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)
+
+            similarities = 0.0
+            target_seq = self.target_states[:, 2]
+            state_seq = (states[:, 2] - self.sens_deviation[0]) / self.sens_norm_fact[0]
+            lag = calc_lag(target_seq, state_seq)
+            similarities += calc_sim(target_seq, state_seq, lag) / 6
+
+            for i in range(1, 6):
+                target_seq = self.target_states[:, i+2]
+                state_seq = (states[:, i+2] - self.sens_deviation[i]) / self.sens_norm_fact[i]
+                similarities += calc_sim(target_seq, state_seq, lag) / 6
+
+            reward_sim = similarities
+
+            target_seq = self.target_states[:, 0]
+            state_seq = states[:, 8] + states[:, 10] + states[:, 12] + states[:, 14]
+            ave_drag = np.average(state_seq)
+            lag = calc_lag(target_seq, state_seq)
+            similarities += calc_sim(target_seq, state_seq, lag) / 2
+            target_seq = self.target_states[:, 1]
+            state_seq = states[:, 9] + states[:, 11] + states[:, 13] + states[:, 15]
+            similarities += calc_sim(target_seq, state_seq, lag) / 2
+            reward_force = similarities
+
+            reward_cd = np.exp(-np.abs((cd-ave_drag) * 2))
+            reward_cl = np.exp(-np.abs(cl * 8))
+
+            reward = np.minimum(0.0 * reward_cd + 0.0 * reward_cl + 0.4 * reward_force + 0.6 * reward_sim, 1.0)
+            # barrier.wait()
+            result_queue.put((np.hstack([forces[2: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)
+        # truncated = False
+        return observation, float(reward), False, truncated, {}
+    
+    def reset(self, seed=None):
+        self.flow_field.apply_ddf()
+        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__()

scripts/nohup1.out

@@ -0,0 +1,16 @@
+Traceback (most recent call last):
+  File "/home/frank14f/Frank_LBM/scripts/d1a3o12_imit.py", line 55, in <module>
+    test_obs, test_rewards, test_dones, info = test_env.step(test_action)
+  File "/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/base_vec_env.py", line 206, in step
+    return self.step_wait()
+  File "/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/vec_env/dummy_vec_env.py", line 58, in step_wait
+    obs, self.buf_rews[env_idx], terminated, truncated, self.buf_infos[env_idx] = self.envs[env_idx].step(
+  File "/home/frank14f/anaconda3/envs/pycuda_3_10/lib/python3.10/site-packages/stable_baselines3/common/monitor.py", line 94, in step
+    observation, reward, terminated, truncated, info = self.env.step(action)
+  File "/home/frank14f/Frank_LBM/scripts/gym_env_imit.py", line 190, in step
+    run_flow_field(action)
+  File "/home/frank14f/Frank_LBM/scripts/gym_env_imit.py", line 117, in run_flow_field
+    self.flow_field.run(SAMPLE_INTERVAL, temp)
+  File "/home/frank14f/Frank_LBM/CelerisLab/driver.py", line 254, in run
+    cuda.memcpy_htod_async(self.action_gpu, action_pinned, stream)
+KeyboardInterrupt