{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.special import jv\n",
    "\n",
    "# 定义网格\n",
    "x = np.linspace(-4, 4, 200)\n",
    "y = np.linspace(-2, 2, 100)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "r = np.sqrt(X**2 + Y**2)\n",
    "theta = np.arctan2(Y, X)\n",
    "\n",
    "# 定义流动参数\n",
    "kappa = 1.0  # 强度参数\n",
    "a = 1.0  # 距离参数\n",
    "b = 3.831705970207512  # 第一个零点\n",
    "n = b / a  # 零点比例\n",
    "u0 = 1\n",
    "rho = 1\n",
    "\n",
    "# 定义分段流函数\n",
    "phi = np.zeros_like(r)\n",
    "psi = np.zeros_like(r)\n",
    "omega = np.zeros_like(r)\n",
    "\n",
    "inside = r <= a\n",
    "outside = r > a\n",
    "\n",
    "phi[inside] = 2 * u0 / n / jv(0, b) * jv(1, n * r[inside]) * np.sin(theta[inside])\n",
    "phi[outside] = u0 * (r[outside] - a**2 / r[outside]) * np.sin(theta[outside])\n",
    "\n",
    "psi[inside] = (2 * u0 / n / jv(0, b) * jv(1, n * r[inside]) - u0 * r[inside]) * np.sin(theta[inside])\n",
    "psi[outside] = -u0 * a**2 / r[outside] * np.sin(theta[outside])\n",
    "\n",
    "# 计算速度场\n",
    "u = np.gradient(phi, axis=0)\n",
    "v = -np.gradient(phi, axis=1)\n",
    "\n",
    "speed = np.sqrt(u**2 + v**2)\n",
    "\n",
    "# 计算涡量\n",
    "vorticity = np.gradient(v, axis=1) - np.gradient(u, axis=0)\n",
    "\n",
    "omega[inside] = 2 * n * u0 / jv(0, b) * jv(1, n * r[inside]) * np.sin(theta[inside])\n",
    "omega[outside] = 0\n",
    "\n",
    "p = rho * (-2 * vorticity * psi - (u**2 + v**2) / 2)\n",
    "\n",
    "# 绘制速度场和流线\n",
    "plt.figure(figsize=(13, 10))\n",
    "plt.streamplot(X, Y, u, v, density=2, linewidth=1, arrowsize=1)\n",
    "plt.contour(X, Y, phi, levels=20, colors='black')\n",
    "plt.contourf(X, Y, speed, levels=10, cmap='jet')\n",
    "plt.colorbar(label='Vorticity')\n",
    "plt.title('Lamb Dipole')\n",
    "plt.xlabel('X')\n",
    "plt.ylabel('Y')\n",
    "plt.axis('equal')\n",
    "plt.show()\n",
    "\n",
    "# 绘制涡量图\n",
    "plt.figure(figsize=(13, 10))\n",
    "plt.contourf(X, Y, p, levels=10, cmap='jet')\n",
    "plt.colorbar(label='Pressure')\n",
    "plt.title('Pressure of Lamb Dipole')\n",
    "plt.xlabel('X')\n",
    "plt.ylabel('Y')\n",
    "plt.axis('equal')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.special import expi\n",
    "\n",
    "# 定义网格\n",
    "x = np.linspace(-100, 100, 200)\n",
    "y = np.linspace(-100, 100, 200)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "r = np.sqrt(X**2 + Y**2)\n",
    "theta = np.arctan2(Y, X)\n",
    "\n",
    "# 定义流动参数\n",
    "strength = 0.01  # 强度参数\n",
    "radius = 10  # 涡量核半径\n",
    "u0 = 1\n",
    "nu = 0.004\n",
    "rho = 1\n",
    "\n",
    "# 定义分段流函数\n",
    "omega = np.zeros_like(r)\n",
    "\n",
    "# 计算速度场\n",
    "kappa = 2 * np.pi * radius **2 * strength\n",
    "v_theta = kappa / (2 * np.pi * r) * (1 - np.exp(-4 * r**2 / radius**2))\n",
    "\n",
    "# M = strength * np.pi * radius**4 / 8 / nu\n",
    "# v_theta = M * r * 4 * nu / radius**4 * np.exp(-r**2 / radius**2)\n",
    "\n",
    "u = -v_theta * np.sin(theta)\n",
    "v = v_theta * np.cos(theta)\n",
    "\n",
    "speed = np.sqrt(u**2 + v**2)\n",
    "\n",
    "# 计算涡量\n",
    "omega = np.gradient(v, axis=1) - np.gradient(u, axis=0)\n",
    "\n",
    "zeta = 4 * r**2 / radius**2\n",
    "p = -rho * kappa**2 / 8 / np.pi**2 / r**2 * (-2*zeta*(expi(-zeta)-expi(-2*zeta)) + (1-np.exp(-zeta))**2)\n",
    "# zeta = r**2 / radius**2\n",
    "# p = -4 * M**2 * nu**2 * (np.exp(-2 * zeta)) / np.pi**2 / radius**6\n",
    "\n",
    "# 绘制速度场和流线\n",
    "plt.figure(figsize=(6, 5))\n",
    "plt.streamplot(X, Y, u, v, density=2, linewidth=1, arrowsize=1)\n",
    "plt.contourf(X, Y, speed, levels=10, cmap='jet')\n",
    "plt.colorbar(label='Vorticity')\n",
    "plt.title('Lamb Dipole')\n",
    "plt.xlabel('X')\n",
    "plt.ylabel('Y')\n",
    "plt.axis('equal')\n",
    "plt.show()\n",
    "\n",
    "# 绘制涡量图\n",
    "plt.figure(figsize=(6, 5))\n",
    "plt.contourf(X, Y, p, levels=10, cmap='jet')\n",
    "plt.colorbar(label='Pressure')\n",
    "plt.title('Pressure of Lamb Dipole')\n",
    "plt.xlabel('X')\n",
    "plt.ylabel('Y')\n",
    "plt.axis('equal')\n",
    "plt.show()\n",
    "plt.figure(figsize=(6, 5))\n",
    "plt.contourf(X, Y, omega, levels=10, cmap='jet')\n",
    "plt.colorbar(label='Vorticity')\n",
    "plt.title('Vorticity of Lamb Dipole')\n",
    "plt.xlabel('X')\n",
    "plt.ylabel('Y')\n",
    "plt.axis('equal')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from collections import deque\n",
    "from typing import Tuple\n",
    "import sys\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "os.environ[\"OMP_NUM_THREADS\"] = \"16\"\n",
    "os.environ[\"MKL_NUM_THREADS\"] = \"16\"\n",
    "\n",
    "current_dir = os.path.dirname(os.path.abspath(\"__file__\"))\n",
    "parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))\n",
    "sys.path.append(parent_dir)\n",
    "from CelerisLab import FlowField\n",
    "from CelerisLab import utils\n",
    "\n",
    "config_cuda = utils.load_cuda_config(\n",
    "    os.path.join(parent_dir, \"configs\", \"config_cuda.json\")\n",
    ")\n",
    "config_field = utils.load_flow_field_config(\n",
    "    os.path.join(parent_dir, \"configs\", \"config_flowfield.json\")\n",
    ")\n",
    "\n",
    "U0 = config_field.velocity\n",
    "T0 = 1000\n",
    "SAMPLE_INTERVAL = 800\n",
    "DATA_TYPE = np.float32\n",
    "L0 = 20\n",
    "\n",
    "flow_field = FlowField(config_field, config_cuda, 1)\n",
    "NX = flow_field.FIELD_SHAPE[0]\n",
    "NY = flow_field.FIELD_SHAPE[1]\n",
    "center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 + 2 * L0, 0)\n",
    "flow_field.add_sensor(center, L0 / 4)\n",
    "center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2, 0)\n",
    "flow_field.add_sensor(center, L0 / 4)\n",
    "center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 - 2 * L0, 0)\n",
    "flow_field.add_sensor(center, L0 / 4)\n",
    "center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)\n",
    "flow_field.add_cylinder(center, L0 / 2)\n",
    "center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 + 0.75 * L0, 0)\n",
    "flow_field.add_cylinder(center, L0 / 2)\n",
    "center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 - 0.75 * L0, 0)\n",
    "flow_field.add_cylinder(center, L0 / 2)\n",
    "flow_field.run(int(3*NX/U0), np.zeros(6, dtype=DATA_TYPE))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 377924 characters]",
      "text/plain": [
       "<Figure size 1000x1500 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "flow_field.get_ddf()\n",
    "ddf_plot = flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)\n",
    "ux = (ddf_plot[:, :, 1] + ddf_plot[:, :, 5] + ddf_plot[:, :, 8] - ddf_plot[:, :, 3] - ddf_plot[:, :, 6] - ddf_plot[:, :, 7]) / U0\n",
    "uy = (ddf_plot[:, :, 2] + ddf_plot[:, :, 5] + ddf_plot[:, :, 6] - ddf_plot[:, :, 4] - ddf_plot[:, :, 7] - ddf_plot[:, :, 8]) / U0\n",
    "pre = np.sum(ddf_plot, axis=2) / 3\n",
    "speed = np.sqrt(ux**2 + uy**2)\n",
    "vorticity = np.gradient(uy, axis=0) - np.gradient(ux, axis=1)\n",
    "plt.figure(figsize=(10, 15))\n",
    "plt.subplot(3, 1, 1)\n",
    "plt.streamplot(np.arange(NX), np.arange(NY), ux.T, uy.T, density=2, linewidth=1, arrowsize=1)\n",
    "plt.imshow(speed.T, origin='lower', cmap='jet', extent=[0, NX, 0, NY])\n",
    "plt.colorbar(label='Speed')\n",
    "plt.subplot(3, 1, 2)\n",
    "plt.imshow(vorticity.T, origin='lower', cmap='jet', extent=[0, NX, 0, NY])\n",
    "plt.colorbar(label='Vorticity')\n",
    "plt.subplot(3, 1, 3)\n",
    "plt.imshow(pre.T, origin='lower', cmap='jet', extent=[0, NX, 0, NY])\n",
    "plt.colorbar(label='Pressure')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# center: Tuple[float, float, float] = (50 * L0, (NY - 1) / 2, 0)\n",
    "# flow_field.add_vortex(center, L0 * 2, 2*U0, 0, \"lamb\")\n",
    "# center: Tuple[float, float, float] = (10 * L0, (NY - 1) / 2, 0)\n",
    "# flow_field.add_vortex(center, L0 * 2, 0.05*U0, 0, \"oseen\")\n",
    "\n",
    "center: Tuple[float, float, float] = (15 * L0, (NY - 1) / 2, 0)\n",
    "flow_field.add_vortex(center, L0 * 2, 0.5*U0, 0, \"lamb\")\n",
    "# flow_field.run(int(1*NX/U0), np.zeros(6, dtype=DATA_TYPE))\n",
    "# flow_field.get_ddf()\n",
    "# flow_field.save_ddf()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "flow_field.run(int(0.1*NX/U0), np.zeros(6, dtype=DATA_TYPE))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "flow_field.restore_ddf()\n",
    "flow_field.apply_ddf()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(150):\n",
    "    flow_field.run(800, np.zeros(1, dtype=DATA_TYPE))\n",
    "    if i % 2 == 0:\n",
    "        file_index = i // 2 + 1\n",
    "        file_name = f\"bg.dat.{file_index:02d}\"\n",
    "\n",
    "        flow_field.get_ddf()\n",
    "        ddf_plot = flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)\n",
    "        flag_plot = flow_field.flag.copy().reshape((NY, NX)).transpose(1, 0)\n",
    "        ux = (ddf_plot[:, :, 1] + ddf_plot[:, :, 5] + ddf_plot[:, :, 8] - ddf_plot[:, :, 3] - ddf_plot[:, :, 6] - ddf_plot[:, :, 7]) / U0\n",
    "        uy = (ddf_plot[:, :, 2] + ddf_plot[:, :, 5] + ddf_plot[:, :, 6] - ddf_plot[:, :, 4] - ddf_plot[:, :, 7] - ddf_plot[:, :, 8]) / U0\n",
    "        with open(os.path.join(parent_dir, \"output\", file_name), \"w\") as f:\n",
    "            f.write(\"Title= \\\"LBM 2D\\\"\\r\\n\")\n",
    "            f.write(\"VARIABLES= \\\"X\\\",\\\"Y\\\",\\\"flag\\\",\\\"U\\\",\\\"V\\\",\\r\\n\")\n",
    "            f.write(f\"ZONE T= \\\"BOX\\\",I= {NX},J= {NY},F=POINT\\r\\n\")\n",
    "            for j in range(NY):\n",
    "                for i in range(NX):\n",
    "                    f.write(f\"{i},{j},{flag_plot[i, j]},{ux[i, j]},{uy[i, j]}\\r\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import re\n",
    "\n",
    "def read_flow_field(file_path):\n",
    "    with open(file_path, 'r') as f:\n",
    "        # 读取前三行元数据\n",
    "        title_line = f.readline()    # 标题行（忽略）\n",
    "        variables_line = f.readline()  # 变量声明（忽略）\n",
    "        zone_line = f.readline()     # 区域信息（需要解析NX和NY）\n",
    "        \n",
    "        # 从ZONE行解析NX和NY\n",
    "        nx = int(re.search(r'I=\\s*(\\d+)', zone_line).group(1))\n",
    "        ny = int(re.search(r'J=\\s*(\\d+)', zone_line).group(1))\n",
    "        \n",
    "        # 初始化ux和uy数组\n",
    "        ux = np.zeros((nx, ny), dtype=np.float64)\n",
    "        uy = np.zeros((nx, ny), dtype=np.float64)\n",
    "        \n",
    "        # 按坐标顺序读取数据（与保存时的双重循环顺序一致）\n",
    "        for j in range(ny):\n",
    "            for i in range(nx):\n",
    "                line = f.readline().strip()\n",
    "                if not line:\n",
    "                    continue\n",
    "                parts = line.split(',')\n",
    "                # 直接按读取顺序填充，无需解析i/j\n",
    "                ux_val = float(parts[3])\n",
    "                uy_val = float(parts[4])\n",
    "                ux[i, j] = ux_val\n",
    "                uy[i, j] = uy_val\n",
    "        \n",
    "        return ux, uy\n",
    "\n",
    "# 使用示例\n",
    "list_index = np.array([])\n",
    "list_ux = np.array([])\n",
    "list_uy = np.array([])\n",
    "list_ref = np.array([])\n",
    "x_slice = 20 * 30\n",
    "for i in range(100):\n",
    "    file_name = os.path.join(parent_dir, \"output\", \"d0a3o12\", f\"field.dat.{i+1:02d}\")\n",
    "    ux, uy = read_flow_field(file_name)\n",
    "    list_index = np.concatenate((list_index, np.ones(ux[x_slice, :].shape) * i))\n",
    "    list_ux = np.concatenate((list_ux, np.array(ux[x_slice, :])))\n",
    "    list_uy = np.concatenate((list_uy, np.array(uy[x_slice, :])))\n",
    "    list_ref = np.concatenate((list_ref, np.array(ux[1, :])))\n",
    "data = {}\n",
    "data['index'] = list_index\n",
    "data['ux'] = list_ux\n",
    "data['uy'] = list_uy\n",
    "data['ref'] = list_ref\n",
    "df = pd.DataFrame(data)\n",
    "df.to_csv(os.path.join(parent_dir, \"output\", \"d0a3o12\", \"silce\", \"test.csv\"), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os\n",
    "def save_slice(self, x_slice, filename):\n",
    "    NX = self.flow_field.FIELD_SHAPE[0]\n",
    "    NY = self.flow_field.FIELD_SHAPE[1]\n",
    "    self.flow_field.get_ddf()\n",
    "    ddf_plot = self.flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)\n",
    "    ux = (ddf_plot[x_slice, :, 1] + ddf_plot[x_slice, :, 5] + ddf_plot[x_slice, :, 8] - ddf_plot[x_slice, :, 3] - ddf_plot[x_slice, :, 6] - ddf_plot[x_slice, :, 7]) / U0\n",
    "    uy = (ddf_plot[x_slice, :, 2] + ddf_plot[x_slice, :, 5] + ddf_plot[x_slice, :, 6] - ddf_plot[x_slice, :, 4] - ddf_plot[x_slice, :, 7] - ddf_plot[x_slice, :, 8]) / U0\n",
    "    data = {}\n",
    "    data[\"ux\"] = ux\n",
    "    data[\"uy\"] = uy\n",
    "    df = pd.DataFrame(data)\n",
    "    df.to_csv(filename, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple, Union\n",
    "from collections import deque\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from stable_baselines3 import PPO\n",
    "import pandas as pd\n",
    "import pickle\n",
    "import sys\n",
    "import os\n",
    "\n",
    "current_dir = os.path.dirname(os.path.abspath(\"__file__\"))\n",
    "parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))\n",
    "sys.path.append(parent_dir)\n",
    "\n",
    "from gym_env_250326 import CustomEnv\n",
    "\n",
    "vec_env = CustomEnv(device_id=1)\n",
    "model = PPO.load(os.path.join(parent_dir, \"models\", \"250326\", \"d1a3o12_250326_240.zip\"), env=vec_env, device=\"cuda:2\")\n",
    "model_env = model.get_env()\n",
    "\n",
    "L0 = 20\n",
    "U0 = 0.01\n",
    "NY = 512"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best seed:  0\n"
     ]
    }
   ],
   "source": [
    "best_seed = 0\n",
    "max_reward = 0\n",
    "for seed in range(100):\n",
    "    model.set_random_seed(seed=seed)\n",
    "    model_env = model.get_env()\n",
    "    obs = model_env.reset()\n",
    "    reward_list = []\n",
    "    for _ in range(150):\n",
    "        action, _states = model.predict(observation=obs, deterministic=True)\n",
    "        obs, rewards, dones, info = model_env.step(action)\n",
    "        reward_list.append(rewards)\n",
    "    if np.mean(reward_list[:100]) > max_reward:\n",
    "        max_reward = np.mean(np.array(reward_list))\n",
    "        best_seed = seed\n",
    "print(\"Best seed: \", best_seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.set_random_seed(seed=best_seed)\n",
    "model.set_random_seed(seed=0)\n",
    "states_list = []\n",
    "reward_list = []\n",
    "action_list = []\n",
    "obs_list = []\n",
    "r_cl_list, r_cd_list, r_sim_list = [], [], []\n",
    "obs = model_env.reset()\n",
    "for i in range(300):\n",
    "    if i < -50:\n",
    "        action = np.array([0.0, 1.0, -1.0], dtype=np.float32)[np.newaxis, :]\n",
    "    else:\n",
    "        action, _states = model.predict(observation=obs, deterministic=True)\n",
    "    obs, rewards, dones, info = model_env.step(action)\n",
    "    states = model_env.envs[0].flow_field.obs.copy()\n",
    "\n",
    "    # if i == 180:\n",
    "    #     center: Tuple[float, float, float] = (15 * L0, (NY - 1) / 2, 0)\n",
    "    #     model_env.envs[0].flow_field.add_vortex(center, L0 * 2, 0.5*U0, 0, \"lamb\")\n",
    "\n",
    "    if i % 2 == 0 and i > 200:\n",
    "    # if i == 299:\n",
    "        file_index = i // 2 + 1\n",
    "        file_name = f\"d1a3o12.dat.{file_index:02d}\"\n",
    "        model_env.envs[0].save_field(os.path.join(parent_dir, \"output\", \"250326\", \"episode\", \"240\", file_name))\n",
    "    \n",
    "    reward_list.append(rewards)\n",
    "    action_list.append(action[0,:])\n",
    "    states_list.append(np.array(states))\n",
    "    obs_list.append(obs[0,:])\n",
    "    r_cl_list.append(model_env.envs[0].reward_cl.copy())\n",
    "    r_cd_list.append(model_env.envs[0].reward_cd.copy())\n",
    "    r_sim_list.append(model_env.envs[0].reward_sim.copy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 66604 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "obs = model_env.reset()\n",
    "# model_env = model.get_env()\n",
    "model_env.envs[0].render()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(np.array(model_env.envs[0].fifo_states))\n",
    "plt.plot(np.array(model_env.envs[0].target_states))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Target States')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 400752 characters]",
      "text/plain": [
       "<Figure size 1500x1200 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "offset = np.array([0, -4, +4], dtype=np.float32)\n",
    "action_offset = [action*4+offset for action in action_list]\n",
    "# target_states = model_env.envs[0].target_states.copy()\n",
    "sense_list = np.array(obs_list)[:, 6:12]\n",
    "for i in range(6):\n",
    "    norm_factor = model_env.envs[0].sens_norm_fact[i]\n",
    "    diviation = model_env.envs[0].sens_deviation[i]\n",
    "    # target_states[:, i] = (target_states[:, i] * norm_factor + diviation) / 78 / U0\n",
    "    sense_list[:, i] = (sense_list[:, i] * norm_factor + diviation) / 78 / U0\n",
    "force_list = np.array(obs_list)[:, 0:6] * model_env.envs[0].force_norm_fact / (L0*U0*U0)\n",
    "force_list = np.array(force_list)\n",
    "drag = force_list[:, 0] + force_list[:, 2] + force_list[:, 4]\n",
    "lift = force_list[:, 1] + force_list[:, 3] + force_list[:, 5]\n",
    "plt.figure(figsize=(15, 12))\n",
    "plt.subplot(2, 3, 1)\n",
    "plt.plot(np.array(sense_list))\n",
    "# plt.plot(np.array(model_env.envs[0].target_states))\n",
    "# plt.plot(np.array(obs_list)[:, 6:12])\n",
    "plt.title(\"Sensors\")\n",
    "plt.subplot(2, 3, 2)\n",
    "plt.plot(drag, label=\"Drag\")\n",
    "plt.plot(lift, label=\"Lift\")\n",
    "plt.title(\"Forces\")\n",
    "plt.subplot(2, 3, 3)\n",
    "plt.plot(np.array(action_offset))\n",
    "plt.title(\"Actions\")\n",
    "plt.subplot(2, 3, 4)\n",
    "plt.plot(reward_list)\n",
    "plt.title(\"Rewards\")\n",
    "plt.subplot(2, 3, 5)\n",
    "plt.plot(np.array(r_sim_list))\n",
    "plt.title(\"Similarity\")\n",
    "plt.subplot(2, 3, 6)\n",
    "plt.plot(np.array(model_env.envs[0].target_states) / 78 / U0)\n",
    "plt.title(\"Target States\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = {}\n",
    "for i in range(3):\n",
    "    data[\"action\" + str(i+1)] = np.array(action_offset)[:, i]\n",
    "for i in range(3):\n",
    "    data[\"cd\" + str(i+1)] = np.array(force_list)[:, 2*i]\n",
    "    data[\"cl\" + str(i+1)] = np.array(force_list)[:, 2*i+1]\n",
    "for i in range(6):\n",
    "    data[\"sensor\" + str(i+1)] = np.array(sense_list)[:, i]\n",
    "    data[\"target\" + str(i+1)] = np.array(model_env.envs[0].target_states)[:, i] / 78 / U0\n",
    "data[\"reward\"] = np.array(reward_list)[:, 0]\n",
    "data[\"sim\"] = np.array(r_sim_list)\n",
    "\n",
    "df = pd.DataFrame(data)\n",
    "df.to_csv(os.path.join(parent_dir, \"output\", \"250329\", \"lamb.csv\"), index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pycuda_3_10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
