diff --git a/examples/jaxfluids/test_jaxfluids_env.py b/examples/jaxfluids/test_jaxfluids_env.py index e841df65..89e0bcc1 100644 --- a/examples/jaxfluids/test_jaxfluids_env.py +++ b/examples/jaxfluids/test_jaxfluids_env.py @@ -8,7 +8,7 @@ which serves as the common base class. JAXFluidsFlowEnv has the following arguments: - - environment_name: Required. Name of the enviroment. + - environment_name: Required. Name of the environment. - hf_repo_id: Hugging Face repository (default: 'dynamicslab/HydroGym-environments') - use_clean_cache: Use clean cache directory (default: True) @@ -46,7 +46,7 @@ are part of the observation - is_scale_observations: Optional. Boolean indicating whether observations are scaled to [0, 1]. - target_fn: Optional. Target thrust vector function. Choose either 'sine' or 'step'. - + The Nozzle3D environment has the following additional arguments: - num_actuators: Required. Integer number of actuators. Must be between 4 and 12. - secondary_pressure_ratio: Optional. Float, must be between 0.7 and 0.9. @@ -57,7 +57,7 @@ are part of the observation - is_scale_observations: Optional. Boolean indicating whether observations are scaled to [0, 1]. - target_fn: Optional. Target thrust vector function. Choose either 'sine' or 'step'. - + """ import os @@ -68,7 +68,7 @@ def main(): env_config = { "environment_name": "Nozzle2D_coarse", - "configuration_file": os.path.abspath("environment_config.yaml") + "configuration_file": os.path.abspath("environment_config.yaml"), } env = Nozzle2D(env_config=env_config) @@ -77,12 +77,11 @@ def main(): env.render() for i in range(1000): - # Random action - # action = env.action_space.sample() + # action = env.action_space.sample() # Fixed action - action = [0.0, 0.5] + action = [0.0, 0.5] observation, reward, terminated, truncated, info = env.step(action) @@ -96,4 +95,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/examples/nek/getting_started/3_pettingzoo/run_pettingzoo_docker.sh b/examples/nek/getting_started/3_pettingzoo/run_pettingzoo_docker.sh index 8f729473..56e5128a 100755 --- a/examples/nek/getting_started/3_pettingzoo/run_pettingzoo_docker.sh +++ b/examples/nek/getting_started/3_pettingzoo/run_pettingzoo_docker.sh @@ -1,6 +1,7 @@ #!/usr/bin/env bash # # Run NEK5000 PettingZoo tests with MPMD coupling. +# Config is loaded automatically from HuggingFace (environment_config.yaml). # # Usage: # ./run_pettingzoo_docker.sh # Test only diff --git a/examples/nek/getting_started/6_zeroshot_wing_demo/README.md b/examples/nek/getting_started/6_zeroshot_wing_demo/README.md index 5d7f6eed..77d9df91 100644 --- a/examples/nek/getting_started/6_zeroshot_wing_demo/README.md +++ b/examples/nek/getting_started/6_zeroshot_wing_demo/README.md @@ -9,7 +9,6 @@ __This is a deployment/evaluation demo only (no training). The template and cont ## What the script does `test_nek_pettingzoo.py`: - - loads a base `NekEnv` via `NekEnv.from_hf(...)` and wraps it with `make_pettingzoo_env(...)` - builds one controller per entry in `POLICY_SPECS` (from `meta_policy_small_wing_template.py`) - assigns each controller to actuator agents by `x_range` and `side` (`SS` means `y > 0`, `PS` means `y < 0`) @@ -42,7 +41,6 @@ obs_dict, rewards_dict, terminations, truncations, infos = env.step(actions) ## Usage ### Recommended: use the runner script - From `6_zeroshot_wing_demo/`: ```bash @@ -58,7 +56,6 @@ mpirun -np 1 python test_nek_pettingzoo.py : -np 12 nek5000 ``` Legacy policy template + run root: - ```bash mpirun -np 1 python test_nek_pettingzoo.py \ --policy-template ./meta_policy_small_wing_template.py \ @@ -68,7 +65,6 @@ mpirun -np 1 python test_nek_pettingzoo.py \ ``` Useful overrides: - - `--policy-template PATH` (defaults to `./meta_policy_small_wing_template.py`) - `--env ENV_NAME` (defaults from template `ENV_NAME`) - `--nproc NPROC` (defaults from template `NPROC`) @@ -82,7 +78,6 @@ Useful overrides: The template defines a lightweight legacy-`MetaPolicy.py`-style configuration. Required top-level variables: - - `ENV_NAME` - `NPROC` - `NUM_STEPS` @@ -90,7 +85,6 @@ Required top-level variables: - `POLICY_SPECS` (list of policy group dicts) Each `POLICY_SPECS` entry supports: - - `name` - `x_range: [x_min, x_max]` - `side: "SS"` (y>0) or `"PS"` (y<0) @@ -101,7 +95,6 @@ Each `POLICY_SPECS` entry supports: - RL algorithms only: `agent_run_name`, `policy`, and/or `model_path` Algorithm semantics: - - `ZERO` outputs an all-zero action (no model needed) - `BL` outputs a constant action equal to `action_max` (no model needed) - `PPO`/`TD3`/`DDPG` load a Stable-Baselines3 model from `model_path`/`POLICY_ROOT` diff --git a/examples/nek/getting_started/6_zeroshot_wing_demo/meta_policy_small_wing_template.py b/examples/nek/getting_started/6_zeroshot_wing_demo/meta_policy_small_wing_template.py old mode 100755 new mode 100644 diff --git a/examples/nek/getting_started/6_zeroshot_wing_demo/run_pettingzoo_docker.sh b/examples/nek/getting_started/6_zeroshot_wing_demo/run_pettingzoo_docker.sh old mode 100755 new mode 100644 diff --git a/examples/nek/getting_started/6_zeroshot_wing_demo/zeroshot_demo_pettingzoo.py b/examples/nek/getting_started/6_zeroshot_wing_demo/zeroshot_demo_pettingzoo.py old mode 100755 new mode 100644 diff --git a/hydrogym/jaxfluids/__init__.py b/hydrogym/jaxfluids/__init__.py index e34f630c..7814e973 100644 --- a/hydrogym/jaxfluids/__init__.py +++ b/hydrogym/jaxfluids/__init__.py @@ -1 +1 @@ -from .envs.nozzle import Nozzle2D, Nozzle3D \ No newline at end of file +from .envs.nozzle import Nozzle2D, Nozzle3D diff --git a/hydrogym/jaxfluids/env_core.py b/hydrogym/jaxfluids/env_core.py index 7ba530e6..afbed9e0 100644 --- a/hydrogym/jaxfluids/env_core.py +++ b/hydrogym/jaxfluids/env_core.py @@ -3,12 +3,11 @@ from pathlib import Path from typing import Dict, Optional, Tuple, Union +from jaxfluids_rl.jxf_env import JAXFluidsEnv, RenderMode from omegaconf import OmegaConf from hydrogym.data_manager import HFDataManager -from jaxfluids_rl.jxf_env import JAXFluidsEnv, RenderMode - class ConfigError(Exception): """Exception raised for configuration-related errors.""" @@ -21,7 +20,7 @@ class JAXFluidsFlowEnv(JAXFluidsEnv): Base JAXFluidsFlowEnv with Hugging Face Hub integration for configuration management. Arguments: - - environment_name: Required. Name of the enviroment. + - environment_name: Required. Name of the environment. - hf_repo_id: Hugging Face repository (default: 'dynamicslab/HydroGym-environments') - use_clean_cache: Use clean cache directory (default: True) @@ -35,7 +34,6 @@ class JAXFluidsFlowEnv(JAXFluidsEnv): """ def _init_from_hf(self, env_config: dict) -> None: - # Initialize HF data manager self.hf_repo_id = env_config.get("hf_repo_id", "dynamicslab/HydroGym-environments") self.local_fallback_dir = env_config.get("local_fallback_dir", None) @@ -45,7 +43,7 @@ def _init_from_hf(self, env_config: dict) -> None: repo_id=self.hf_repo_id, local_fallback_dir=self.local_fallback_dir, use_clean_cache=self.use_clean_cache, - fallback_profile="JAXFLUIDS" + fallback_profile="JAXFLUIDS", ) # Environment identification @@ -53,7 +51,7 @@ def _init_from_hf(self, env_config: dict) -> None: if not self.environment_name: raise ConfigError("'environment_name' must be specified in env_config") - + # Download/get environment configuration self.env_data_path = self._setup_environment_data() @@ -65,11 +63,10 @@ def _init_from_hf(self, env_config: dict) -> None: f"No configuration file found for environment '{self.environment_name}'. " f"Expected config.yaml in: {self.env_data_path}" ) - + # Load configuration from HF self.conf = OmegaConf.load(self.configuration_file) - def _setup_environment_data(self) -> str: """ Download and setup environment data from HF Hub. @@ -95,7 +92,6 @@ def _setup_environment_data(self) -> str: return env_path except Exception as e: raise ConfigError(f"Failed to setup environment data for {self.environment_name}: {e}") - def _resolve_configuration_file(self, config_file_input: Optional[str]) -> Optional[str]: """ @@ -154,7 +150,7 @@ def _resolve_configuration_file(self, config_file_input: Optional[str]) -> Optio f" - Current directory: {os.getcwd()}\n" f" - Environment directory: {self.env_data_path}" ) - + def _find_configuration_file(self) -> Optional[str]: """ Auto-detect configuration file in the environment data directory. @@ -192,4 +188,4 @@ def _find_configuration_file(self) -> Optional[str]: if os.path.exists(self.env_data_path): print(f"Available files: {os.listdir(self.env_data_path)}") - return None \ No newline at end of file + return None diff --git a/hydrogym/jaxfluids/envs/nozzle.py b/hydrogym/jaxfluids/envs/nozzle.py index de5707ee..1df5db52 100644 --- a/hydrogym/jaxfluids/envs/nozzle.py +++ b/hydrogym/jaxfluids/envs/nozzle.py @@ -1,39 +1,37 @@ from abc import abstractmethod from dataclasses import dataclass from pathlib import Path -from typing import Any, ClassVar, Callable, NamedTuple +from typing import Any, Callable, ClassVar, NamedTuple +import gymnasium as gym import jax -from jax import Array import jax.numpy as jnp import matplotlib.pyplot as plt -from mpl_toolkits.axes_grid1 import make_axes_locatable import numpy as np -import gymnasium as gym - +from jax import Array from jaxfluids.data_types import JaxFluidsBuffers from jaxfluids.data_types.ml_buffers import ( CallablesSetup, - ParametersSetup, - LevelSetSetup, InterfaceFluxCallablesSetup, InterfaceFluxParametersSetup, + LevelSetSetup, + ParametersSetup, ) from jaxfluids.domain.helper_functions import ( reassemble_buffer_np, reassemble_cell_centers, - reassemble_cell_sizes + reassemble_cell_sizes, ) from jaxfluids_rl.jxf_env import RenderMode - +from mpl_toolkits.axes_grid1 import make_axes_locatable from hydrogym.jaxfluids.env_core import JAXFluidsFlowEnv from hydrogym.jaxfluids.utils.nozzle import ( InjectorGeometry, ObsData, PressureRatios, - TVCSpec, TargetThrustAngleFn, + TVCSpec, build_tvc_env_options, build_tvc_runtime_setup, compute_thrust, @@ -42,16 +40,11 @@ ) -Array = jax.Array - - class NozzleBase(JAXFluidsFlowEnv): - SPEC: ClassVar[TVCSpec] TARGET_FNS: ClassVar[dict[str, TargetThrustAngleFn]] = {} def __init__(self, env_config: dict) -> None: - self._init_from_hf(env_config) env_options = build_tvc_env_options( @@ -95,7 +88,7 @@ def __init__(self, env_config: dict) -> None: ) self.default_action_reset = np.zeros(self.num_actuators) - + if self.is_pressure_probes: self.probe_locations = self._compute_probe_locations() self.num_probes = self.probe_locations.shape[0] @@ -113,7 +106,7 @@ def __init__(self, env_config: dict) -> None: def _is_terminated(self, action: np.ndarray, jxf_buffers: JaxFluidsBuffers, info: dict) -> bool: physical_simulation_time = jxf_buffers.time_control_variables.physical_simulation_time return physical_simulation_time >= self.SPEC.t_end - + def _is_truncated(self, jxf_buffers: JaxFluidsBuffers, info: dict) -> bool: return False @@ -131,9 +124,7 @@ def _build_action_callable_setup(self) -> CallablesSetup: specific_gas_constant=self.SPEC.specific_gas_constant, sim_manager=self.sim_manager, ) - levelset_setup = LevelSetSetup( - fluid_solid=InterfaceFluxCallablesSetup(interface_flux_fn) - ) + levelset_setup = LevelSetSetup(fluid_solid=InterfaceFluxCallablesSetup(interface_flux_fn)) return CallablesSetup(levelset=levelset_setup) def _build_action_space(self) -> gym.spaces.Box: @@ -156,7 +147,10 @@ def _build_observation_space(self) -> gym.spaces.Box: high = np.array([1.0] * (2 * num_angles) + [1.0] * self.num_probes, dtype=np.float32) else: low = np.array([-np.pi] * (2 * num_angles) + [0.0] * self.num_probes, dtype=np.float32) - high = np.array([np.pi] * (2 * num_angles) + [np.inf] * self.num_probes, dtype=np.float32) + high = np.array( + [np.pi] * (2 * num_angles) + [np.inf] * self.num_probes, + dtype=np.float32, + ) return gym.spaces.Box( low=low, @@ -169,7 +163,6 @@ def _convert_action_for_jxf(self, action: np.ndarray) -> ParametersSetup: return ParametersSetup(levelset=levelset_setup) def _get_obs(self) -> np.ndarray: - jxf_buffers, _ = self._require_state() obs_data = self.compute_obs(jxf_buffers) @@ -205,16 +198,15 @@ def _get_info(self) -> dict[str, Any]: } def _after_step( - self, - action: np.ndarray, - observation: np.ndarray, - reward: float, - terminated: bool, - truncated: bool, - info: dict, - jxf_buffers: JaxFluidsBuffers - ) -> None: - + self, + action: np.ndarray, + observation: np.ndarray, + reward: float, + terminated: bool, + truncated: bool, + info: dict, + jxf_buffers: JaxFluidsBuffers, + ) -> None: t = float(jxf_buffers.time_control_variables.physical_simulation_time) self._append_history( time=t, @@ -229,14 +221,13 @@ def compute_obs(self, jxf_buffers: JaxFluidsBuffers) -> ObsData: self._compute_obs, axis_name="i", in_axes=(JaxFluidsBuffers(0, None, None, None),), - out_axes=None + out_axes=None, )(jxf_buffers) else: return jax.jit(self._compute_obs)(jxf_buffers) def _compute_obs(self, jxf_buffers: JaxFluidsBuffers) -> ObsData: - current_angle = self.compute_thrust_angle(jxf_buffers) sim_time = jxf_buffers.time_control_variables.physical_simulation_time @@ -279,20 +270,22 @@ def compute_thrust_angle(self, jxf_buffers: JaxFluidsBuffers) -> Array: if self.SPEC.dim == 2: current_angle = jnp.atan2(thrust[1], thrust[0]) else: - current_angle = jnp.stack([ - jnp.atan2(thrust[1], thrust[0]), # Pitch - jnp.atan2(thrust[2], thrust[0]), # Yaw - ]) + current_angle = jnp.stack( + [ + jnp.atan2(thrust[1], thrust[0]), # Pitch + jnp.atan2(thrust[2], thrust[0]), # Yaw + ] + ) return current_angle def compute_pressure_probes(self, jxf_buffers: JaxFluidsBuffers) -> Array: - x_p = self.probe_locations[:,0] - y_p = self.probe_locations[:,1] - z_p = self.probe_locations[:,2] + x_p = self.probe_locations[:, 0] + y_p = self.probe_locations[:, 1] + z_p = self.probe_locations[:, 2] domain_information = self.sim_manager.domain_information - + cell_centers = domain_information.get_device_cell_centers() x, y, z = [xi.flatten() for xi in cell_centers] @@ -302,7 +295,7 @@ def compute_pressure_probes(self, jxf_buffers: JaxFluidsBuffers) -> Array: if self.SPEC.dim == 3: z_id = jnp.searchsorted(z, z_p, side="left", method="scan_unrolled") else: - z_id = 0 + z_id = 0 nhx, nhy, nhz = domain_information.domain_slices_conservatives pressure = jxf_buffers.simulation_buffers.material_fields.primitives[4, nhx, nhy, nhz] @@ -310,10 +303,10 @@ def compute_pressure_probes(self, jxf_buffers: JaxFluidsBuffers) -> Array: if domain_information.is_parallel: device_domain_size = domain_information.get_device_domain_size() - + mask = 1 for i in range(domain_information.dim): - xi = self.probe_locations[:,i] + xi = self.probe_locations[:, i] mask *= (device_domain_size[i][0] <= xi) & (xi < device_domain_size[i][1]) pressure_probes = jax.lax.psum(mask * pressure_probes, axis_name="i") @@ -335,11 +328,8 @@ def _get_fields_for_plotting(self, jxf_buffers: JaxFluidsBuffers) -> tuple[np.nd ] if domain_information.is_parallel: - fields = [ - reassemble_buffer_np(field, domain_information.split_factors) - for field in fields - ] - + fields = [reassemble_buffer_np(field, domain_information.split_factors) for field in fields] + fields = [field.squeeze() for field in fields] return tuple(fields) @@ -348,7 +338,7 @@ def _get_meshgrid_for_plotting(self) -> tuple[np.ndarray, np.ndarray]: cell_centers = domain_information.get_global_cell_centers() if domain_information.is_parallel: cell_centers = reassemble_cell_centers(cell_centers, domain_information.split_factors) - + x, y, _ = [xi.flatten() for xi in cell_centers] return np.meshgrid(x, y, indexing="ij") @@ -359,7 +349,7 @@ def _compute_probe_locations(self) -> np.ndarray: def render(self) -> None: if self.render_mode is None: return - + self._plot_flow_field() self._plot_observations() @@ -396,7 +386,7 @@ class Nozzle2D(NozzleBase): TARGET_FNS = { "sine": lambda t: (t > 5e-4) * (10.0 / 180.0 * jnp.pi) * jnp.sin(2 * jnp.pi * (t - 5e-4) / 4e-3), - "step": lambda t: (t > 5e-4) * (5.0 / 180.0 * jnp.pi) + "step": lambda t: (t > 5e-4) * (5.0 / 180.0 * jnp.pi), } def _compute_probe_locations(self) -> np.ndarray: @@ -411,11 +401,9 @@ def _compute_probe_locations(self) -> np.ndarray: x_probes = np.array([x, x]) y_probes = np.array([y, -y]) z_probes = np.zeros_like(x_probes) - probe_locations.append( - np.stack([x_probes, y_probes, z_probes], axis=1) - ) + probe_locations.append(np.stack([x_probes, y_probes, z_probes], axis=1)) - return np.concatenate(probe_locations, axis=0) + return np.concatenate(probe_locations, axis=0) def _get_reward(self, action: np.ndarray) -> float: error = np.abs(self.target_angle - self.thrust_angle) @@ -429,7 +417,7 @@ def _plot_flow_field(self) -> None: D_throat = self.SPEC.nozzle_geometry.D_throat X = X / D_throat Y = Y / D_throat - + physical_simulation_time = jxf_buffers.time_control_variables.physical_simulation_time rho = primitives[0] @@ -446,14 +434,15 @@ def _plot_flow_field(self) -> None: fig.suptitle(f"Env Step: {self.env_step}, Time: {physical_simulation_time * 1e3:.3f} ms") quants = (M, p / self.SPEC.p0) - vmins = (0.0, 0.0); vmaxs = (3.0, 1.0) + vmins = (0.0, 0.0) + vmaxs = (3.0, 1.0) for axi, quant, vmin, vmax in zip(ax, quants, vmins, vmaxs): pci = axi.pcolormesh(X, Y, quant, cmap="Spectral_r", vmin=vmin, vmax=vmax, shading="auto") axi.contour(X, Y, M, levels=[1.0], linewidths=0.5, colors="k", linestyles="-") divider = make_axes_locatable(axi) - cax = divider.append_axes('right', size='5%', pad=0.05) - fig.colorbar(pci, cax=cax, orientation='vertical') + cax = divider.append_axes("right", size="5%", pad=0.05) + fig.colorbar(pci, cax=cax, orientation="vertical") titles = ("Mach number", "Pressure p / p_0") for axi, title in zip(ax, titles): @@ -466,7 +455,12 @@ def _plot_flow_field(self) -> None: if self.is_pressure_probes: for axi in ax: - axi.scatter(self.probe_locations[:,0] / D_throat, self.probe_locations[:,1] / D_throat, s=2, c="black") + axi.scatter( + self.probe_locations[:, 0] / D_throat, + self.probe_locations[:, 1] / D_throat, + s=2, + c="black", + ) if self.render_mode is RenderMode.SHOW: plt.show() @@ -476,7 +470,6 @@ def _plot_flow_field(self) -> None: raise ValueError(f"RenderMode {self.render_mode} is not valid.") plt.close(fig) - def _plot_observations(self) -> None: jxf_buffers, _ = self._require_state() physical_simulation_time = jxf_buffers.time_control_variables.physical_simulation_time @@ -497,7 +490,7 @@ def _plot_observations(self) -> None: if len(t) > 0: ax[0].plot(t * 1e3, np.rad2deg(thrust_angle), "k", label="current") for actuator_i in range(self.num_actuators): - ax[1].plot(t * 1e3, action[:,actuator_i], label=f"Inj. {actuator_i:02d}") + ax[1].plot(t * 1e3, action[:, actuator_i], label=f"Inj. {actuator_i:02d}") ax[1].set_ylim(-0.05, 1.05) @@ -506,7 +499,11 @@ def _plot_observations(self) -> None: pressure_probes /= self.SPEC.p0 for probe_i in range(self.num_probes // 2): ax[2].plot(t * 1e3, pressure_probes[:, probe_i], label=f"Probe 0{probe_i:d}") - ax[3].plot(t * 1e3, pressure_probes[:, probe_i + self.num_probes // 2], label=f"Probe 1{probe_i:d}") + ax[3].plot( + t * 1e3, + pressure_probes[:, probe_i + self.num_probes // 2], + label=f"Probe 1{probe_i:d}", + ) ax[2].set_ylim(0.0, 0.5) ax[3].set_ylim(0.0, 0.5) @@ -557,14 +554,18 @@ class Nozzle3D(NozzleBase): ) TARGET_FNS = { - "sine": lambda t: jnp.array([ - (t > 1e-3) * (10.0 / 180.0 * jnp.pi) * jnp.sin(2 * jnp.pi * (t - 1e-3) / 4e-3), - jnp.zeros_like(t), - ]), - "step": lambda t: jnp.array([ - (t > 1e-3) * (5.0 / 180.0 * jnp.pi), - jnp.zeros_like(t), - ]), + "sine": lambda t: jnp.array( + [ + (t > 1e-3) * (10.0 / 180.0 * jnp.pi) * jnp.sin(2 * jnp.pi * (t - 1e-3) / 4e-3), + jnp.zeros_like(t), + ] + ), + "step": lambda t: jnp.array( + [ + (t > 1e-3) * (5.0 / 180.0 * jnp.pi), + jnp.zeros_like(t), + ] + ), } def _compute_probe_locations(self) -> np.ndarray: @@ -581,21 +582,17 @@ def _compute_probe_locations(self) -> np.ndarray: x_probes = np.full_like(theta, x) y_probes = R * np.cos(theta) z_probes = -R * np.sin(theta) - probe_locations.append( - np.stack([x_probes, y_probes, z_probes], axis=1) - ) - return np.concatenate(probe_locations, axis=0) + probe_locations.append(np.stack([x_probes, y_probes, z_probes], axis=1)) + return np.concatenate(probe_locations, axis=0) def _get_reward(self, action: np.ndarray) -> float: - error = jnp.sqrt(jnp.sum((self.target_angle - self.thrust_angle)**2)) + error = jnp.sqrt(jnp.sum((self.target_angle - self.thrust_angle) ** 2)) return float(-error) def _plot_flow_field(self) -> None: jxf_buffers, _ = self._require_state() primitives, levelset, _ = self._get_fields_for_plotting(jxf_buffers) - physical_simulation_time = jxf_buffers.time_control_variables.physical_simulation_time - domain_information = self.sim_manager.domain_information cell_centers = domain_information.get_global_cell_centers() cell_sizes = domain_information.get_global_cell_sizes() @@ -644,11 +641,10 @@ def _plot_observations(self) -> None: t = np.array(self.history["time"]) thrust_angle = np.array(self.history["thrust_angle"]) - action = np.array(self.history["action"]) if len(t) > 0: - ax[0].plot(t * 1e3, np.rad2deg(thrust_angle[:,0]), "k", label="current") - ax[1].plot(t * 1e3, np.rad2deg(thrust_angle[:,1]), "k", label="current") + ax[0].plot(t * 1e3, np.rad2deg(thrust_angle[:, 0]), "k", label="current") + ax[1].plot(t * 1e3, np.rad2deg(thrust_angle[:, 1]), "k", label="current") ax[0].set_ylim(-10.0, 10.0) ax[1].set_ylim(-10.0, 10.0) @@ -663,7 +659,11 @@ def _plot_observations(self) -> None: pressure_probes /= self.SPEC.p0 for probe_i in range(self.num_probes // 2): ax[2].plot(t * 1e3, pressure_probes[:, probe_i], label=f"Probe 0{probe_i:d}") - ax[3].plot(t * 1e3, pressure_probes[:, probe_i + self.num_probes // 2], label=f"Probe 1{probe_i:d}") + ax[3].plot( + t * 1e3, + pressure_probes[:, probe_i + self.num_probes // 2], + label=f"Probe 1{probe_i:d}", + ) ax[2].set_ylim(0.0, 0.5) ax[3].set_ylim(0.0, 0.5) @@ -671,7 +671,7 @@ def _plot_observations(self) -> None: "Thrust angle" + r"$\delta_0$" + "[deg]", "Thrust angle" + r"$\delta_1$" + "[deg]", "Pressure probes downstream loc 0", - "Pressure probes downstream loc 1" + "Pressure probes downstream loc 1", ) for axi, title in zip(ax, titles): axi.set_box_aspect(1.0) diff --git a/hydrogym/jaxfluids/utils/nozzle.py b/hydrogym/jaxfluids/utils/nozzle.py index e855d06a..8ed08601 100644 --- a/hydrogym/jaxfluids/utils/nozzle.py +++ b/hydrogym/jaxfluids/utils/nozzle.py @@ -1,32 +1,31 @@ -from dataclasses import dataclass import os -from pathlib import Path import platform +from dataclasses import dataclass +from pathlib import Path + if platform.machine().lower() in {"aarch64", "arm64"}: os.environ["VTK_DEFAULT_OPENGL_WINDOW"] = "vtkOSOpenGLRenderWindow" +import json from typing import Callable, NamedTuple -from jax import Array import jax.numpy as jnp -import json -from numpy import ndarray import numpy as np import pyvista as pv -pv.global_theme.allow_empty_mesh = True - - +from jax import Array from jaxfluids import SimulationManager -from jaxfluids_thirdparty.gas_dynamics.isentropic import ( - pressure_ratio_isentropic, - density_ratio_isentropic, - mach_number_from_pressure_ratio_isentropic, -) from jaxfluids_thirdparty.gas_dynamics.core import ( + density_from_pressure_temperature, speed_of_sound, total_energy, - density_from_pressure_temperature, ) +from jaxfluids_thirdparty.gas_dynamics.isentropic import ( + density_ratio_isentropic, + mach_number_from_pressure_ratio_isentropic, + pressure_ratio_isentropic, +) +from numpy import ndarray +pv.global_theme.allow_empty_mesh = True TargetThrustAngle = Array | float TargetThrustAngleFn = Callable[[Array | float], TargetThrustAngle] @@ -37,7 +36,8 @@ class NozzleGeometry: """Fixed nozzle geometry based on the publication Das et al. 2025 AIAA """ - A: tuple[float, float] = (0.0,-0.01559) + + A: tuple[float, float] = (0.0, -0.01559) B: tuple[float, float] = (0.0, 0.0352) C: tuple[float, float] = (0.02329, 0.02954) D: tuple[float, float] = (0.05049, 0.01552) @@ -54,14 +54,14 @@ def area_ratio_inlet(self, dim: int) -> float: """ if dim not in (2, 3): raise ValueError(f"Invalid dim. Got {dim}.") - + ratio = self.B[1] / self.R if dim == 2: return ratio else: return ratio**2 - + @property def D_exit(self) -> float: return 2 * self.H[1] @@ -73,9 +73,9 @@ def D_throat(self) -> float: @dataclass(frozen=True, slots=True) class InjectorGeometry: - X: float # position - IW: float # width - N: int # count + X: float # position + IW: float # width + N: int # count @dataclass(frozen=True, slots=True) @@ -87,8 +87,8 @@ class InjectorPlaneParameters: @dataclass(frozen=True, slots=True) class PressureRatios: - NPR: float # nozzle pressure ratio - SPR: float # secondary pressure ratio + NPR: float # nozzle pressure ratio + SPR: float # secondary pressure ratio class ObsData(NamedTuple): @@ -104,7 +104,7 @@ class TVCSpec: fixed_num_actuators: int | None = None min_num_actuators: int | None = None max_num_actuators: int | None = None - ambient_pressure: float = 1e+5 + ambient_pressure: float = 1e5 ambient_temperature: float = 300.0 specific_gas_constant: float = 287.14 specific_heat_ratio: float = 1.4 @@ -140,19 +140,16 @@ class TVCRuntimeSetup: def build_tvc_env_options( - *, - env_config: dict, - spec: TVCSpec, - target_fns: dict[str, TargetThrustAngleFn], - cls_name: str, - ) -> TVCEnvOptions: - + *, + env_config: dict, + spec: TVCSpec, + target_fns: dict[str, TargetThrustAngleFn], + cls_name: str, +) -> TVCEnvOptions: num_actuators = env_config.get("num_actuators") if spec.fixed_num_actuators is not None: if num_actuators is not None: - raise ValueError( - f"{cls_name} requires {spec.fixed_num_actuators} actuators." - ) + raise ValueError(f"{cls_name} requires {spec.fixed_num_actuators} actuators.") num_actuators = spec.fixed_num_actuators else: if num_actuators is None: @@ -162,30 +159,21 @@ def build_tvc_env_options( max_num_actuators = spec.max_num_actuators if min_num_actuators is None or max_num_actuators is None: raise ValueError( - f"{cls_name} must define either a fixed actuator count " - "or both min_num_actuators and max_num_actuators." + f"{cls_name} must define either a fixed actuator count or both min_num_actuators and max_num_actuators." ) if not (min_num_actuators <= num_actuators <= max_num_actuators): raise ValueError( - f"num_actuators must be in " - f"[{min_num_actuators}, {max_num_actuators}]. " - f"Got {num_actuators}." + f"num_actuators must be in [{min_num_actuators}, {max_num_actuators}]. Got {num_actuators}." ) secondary_pressure_ratio = env_config.get("secondary_pressure_ratio", 0.7) if secondary_pressure_ratio < 0.7 or secondary_pressure_ratio > 0.9: - raise ValueError( - f"secondary_pressure_ratio must be >= 0.7 and <= 0.9." - f"Got {secondary_pressure_ratio}." - ) + raise ValueError(f"secondary_pressure_ratio must be >= 0.7 and <= 0.9.Got {secondary_pressure_ratio}.") resolution = env_config.get("resolution", "fine") if resolution not in spec.grid_resolutions: - raise ValueError( - f"Resolution {resolution} is not supported. " - f"Please choose from {spec.grid_resolutions}." - ) + raise ValueError(f"Resolution {resolution} is not supported. Please choose from {spec.grid_resolutions}.") ngpus = env_config.get("ngpus", 1) if not isinstance(ngpus, int): @@ -195,24 +183,15 @@ def build_tvc_env_options( is_pressure_probes = env_config.get("is_pressure_probes", False) if not isinstance(is_pressure_probes, bool): - raise ValueError( - "is_pressure_probes must be of type bool. " - f"Got {type(is_pressure_probes)}" - ) + raise ValueError(f"is_pressure_probes must be of type bool. Got {type(is_pressure_probes)}") is_scale_observations = env_config.get("is_scale_observations", True) if not isinstance(is_scale_observations, bool): - raise ValueError( - f"is_scale_observations needs to be of type bool. " - f"Got {type(is_scale_observations)}." - ) + raise ValueError(f"is_scale_observations needs to be of type bool. Got {type(is_scale_observations)}.") target_key = env_config.get("target_fn", "sine") if target_key not in target_fns: - raise ValueError( - f"Unknown target_fn {target_key!r}. " - f"Please choose from {tuple(target_fns)}." - ) + raise ValueError(f"Unknown target_fn {target_key!r}. Please choose from {tuple(target_fns)}.") return TVCEnvOptions( num_actuators=num_actuators, @@ -226,13 +205,12 @@ def build_tvc_env_options( def build_tvc_runtime_setup( - *, - base_path: Path, - dim: int, - resolution: str, - ngpus: int, - ) -> TVCRuntimeSetup: - + *, + base_path: Path, + dim: int, + resolution: str, + ngpus: int, +) -> TVCRuntimeSetup: env_name = f"Nozzle{dim}D_{resolution}" # env_dir = base_path / env_name env_dir = base_path @@ -265,12 +243,12 @@ def build_tvc_runtime_setup( def compute_thrust( - primitives: Array, # shape (Np,Nx,Ny,Nz) - p_infty: float, - apertures_x: Array, - cell_centers: tuple[Array, ...], - cell_sizes: tuple[Array, ...], - ) -> tuple[Array, Array, Array]: + primitives: Array, # shape (Np,Nx,Ny,Nz) + p_infty: float, + apertures_x: Array, + cell_centers: tuple[Array, ...], + cell_sizes: tuple[Array, ...], +) -> tuple[Array, Array, Array]: """Computes the thrust of the nozzle. F_x = mdot_e * u_e + (p_e - p_infty) * A_e @@ -288,12 +266,10 @@ def compute_thrust( cell_face_area = dx_min if DIM == 2 else dx_min**2 # interpolate primitives to cell face - primitives_cf = jnp.concatenate([ - primitives[:,0:1], primitives, primitives[:,-1:] - ], axis=1) - primitives_cf = (primitives_cf[:,1:] + primitives_cf[:,:-1]) / 2 + primitives_cf = jnp.concatenate([primitives[:, 0:1], primitives, primitives[:, -1:]], axis=1) + primitives_cf = (primitives_cf[:, 1:] + primitives_cf[:, :-1]) / 2 - x_cf = jnp.concatenate([x - dx/2, x[-1:] + dx[-1:]/2], axis=0) + x_cf = jnp.concatenate([x - dx / 2, x[-1:] + dx[-1:] / 2], axis=0) # x ids throat and exit xid_e = jnp.searchsorted(x_cf, nozzle_geometry.H[0]) - 1 @@ -303,7 +279,7 @@ def compute_thrust( if DIM == 2: mask = jnp.abs(Y[xid_e]) <= nozzle_geometry.H[1] else: - mask = jnp.sqrt(Y[xid_e]**2 + Z[xid_e]**2) <= nozzle_geometry.H[1] + mask = jnp.sqrt(Y[xid_e] ** 2 + Z[xid_e] ** 2) <= nozzle_geometry.H[1] # states at nozzle exit A_e = apertures_x[xid_e] * cell_face_area @@ -320,21 +296,21 @@ def compute_thrust( mdot_e = rho_e * vel_e[0] * A_e * mask mdot_t = rho_t * u_t * A_t * mask - # thrust + # thrust thrust = mdot_e * vel_e - thrust = thrust.at[0].add( (p_e - p_infty) * A_e ) - thrust = jnp.sum(thrust * mask, axis=(-1,-2)) + thrust = thrust.at[0].add((p_e - p_infty) * A_e) + thrust = jnp.sum(thrust * mask, axis=(-1, -2)) # metrics - mdot_e = jnp.sum(mdot_e, axis=(-1,-2)) - mdot_t = jnp.sum(mdot_t, axis=(-1,-2)) + mdot_e = jnp.sum(mdot_e, axis=(-1, -2)) + mdot_t = jnp.sum(mdot_t, axis=(-1, -2)) return thrust, mdot_t, mdot_e def _compute_injector_plane_params( - injector_geometry: InjectorGeometry, - ) -> InjectorPlaneParameters: + injector_geometry: InjectorGeometry, +) -> InjectorPlaneParameters: """Computes position, normal and tangent of the injector. :param injector_geometry: _description_ @@ -346,12 +322,12 @@ def _compute_injector_plane_params( nozzle_geometry = NozzleGeometry() # compute vectors for base injector - H = jnp.array(nozzle_geometry.H) # end point of convergent linear nozzle section - E = jnp.array(nozzle_geometry.E) # start point of convergent linear nozzle section + H = jnp.array(nozzle_geometry.H) # end point of convergent linear nozzle section + E = jnp.array(nozzle_geometry.E) # start point of convergent linear nozzle section EH = H - E tangent = EH / jnp.linalg.norm(EH) x_position = E[0] + injector_geometry.X * (H[0] - E[0]) - y_position = (H[1] - E[1])/(H[0] - E[0]) * (x_position - E[0]) + E[1] + y_position = (H[1] - E[1]) / (H[0] - E[0]) * (x_position - E[0]) + E[1] t0 = jnp.array([tangent[0], tangent[1], 0.0]) n0 = jnp.array([tangent[1], -tangent[0], 0.0]) @@ -361,16 +337,18 @@ def _compute_injector_plane_params( tangents = [] normals = [] - theta = jnp.linspace(0, 2*jnp.pi, injector_geometry.N, endpoint=False) + theta = jnp.linspace(0, 2 * jnp.pi, injector_geometry.N, endpoint=False) - # rotate the vectors around the circumference of the nozzle + # rotate the vectors around the circumference of the nozzle # to get the corresponding vectors of the remaining injectors for th in theta: - R = jnp.array([ - [1, 0, 0], - [0, jnp.cos(th), -jnp.sin(th)], - [0, jnp.sin(th), jnp.cos(th)], - ]) + R = jnp.array( + [ + [1, 0, 0], + [0, jnp.cos(th), -jnp.sin(th)], + [0, jnp.sin(th), jnp.cos(th)], + ] + ) positions.append(R @ pos0) tangents.append(R @ t0) @@ -380,19 +358,16 @@ def _compute_injector_plane_params( tangents = jnp.stack(tangents) normals = jnp.stack(normals) - return InjectorPlaneParameters( - positions, tangents, normals - ) + return InjectorPlaneParameters(positions, tangents, normals) def _compute_injector_mask( - mesh_grid: tuple[Array], - IW: float, - injector_params: InjectorPlaneParameters, - dim: int, - dx: float - ) -> Array: - + mesh_grid: tuple[Array], + IW: float, + injector_params: InjectorPlaneParameters, + dim: int, + dx: float, +) -> Array: position = injector_params.positions tangent = injector_params.tangents normal = injector_params.normals @@ -407,43 +382,32 @@ def _compute_injector_mask( if dim == 2: X, Y = mesh_grid - R = jnp.stack([ - X - position[0], - Y - position[1] - ], axis=0) + R = jnp.stack([X - position[0], Y - position[1]], axis=0) - s = jnp.sum(R*tangent[:2], axis=0) - d = jnp.sum(R*normal[:2], axis=0) + s = jnp.sum(R * tangent[:2], axis=0) + d = jnp.sum(R * normal[:2], axis=0) - mask_injector = ( - (jnp.abs(s) <= IW / 2) & - (jnp.abs(d) <= dx) - ) + mask_injector = (jnp.abs(s) <= IW / 2) & (jnp.abs(d) <= dx) elif dim == 3: X, Y, Z = mesh_grid - R = jnp.stack([ - X - position[0], - Y - position[1], - Z - position[2] - ], axis=0) + R = jnp.stack([X - position[0], Y - position[1], Z - position[2]], axis=0) # compute circumferential tangent t_theta = jnp.cross(normal, tangent, axis=0) t_theta = t_theta / jnp.linalg.norm(t_theta, keepdims=True) # projections - s_x = jnp.sum(R*tangent, axis=0) - s_theta = jnp.sum(R*t_theta, axis=0) - d = jnp.sum(R*normal, axis=0) + s_x = jnp.sum(R * tangent, axis=0) + s_theta = jnp.sum(R * t_theta, axis=0) + d = jnp.sum(R * normal, axis=0) R_inj = IW / 2 - mask_injector = ( - (s_x**2 + s_theta**2 <= R_inj**2) & - (jnp.abs(d) <= 5*dx) # NOTE large safety offset in normal direction given that nozzle exit surface is curved around circumference - ) + mask_injector = (s_x**2 + s_theta**2 <= R_inj**2) & ( + jnp.abs(d) <= 5 * dx + ) # Large safety offset in normal direction given that nozzle exit surface is curved around circumference else: raise ValueError @@ -452,45 +416,41 @@ def _compute_injector_mask( def _compute_choked_state( - p_total_injector: float, - rho_total_injector: float, - specific_heat_ratio: float, - ) -> tuple[float, float, float, float]: + p_total_injector: float, + rho_total_injector: float, + specific_heat_ratio: float, +) -> tuple[float, float, float, float]: p_choked = p_total_injector * pressure_ratio_isentropic(1.0, specific_heat_ratio) rho_choked = rho_total_injector * density_ratio_isentropic(1.0, specific_heat_ratio) u_choked = speed_of_sound(p_choked, rho_choked, specific_heat_ratio) E_choked = total_energy(p_choked, rho_choked, u_choked, specific_heat_ratio) - return p_choked, rho_choked, u_choked, E_choked + return p_choked, rho_choked, u_choked, E_choked def _compute_unchoked_state( - p_local: float, - rho_total_injector: float, - p_ratio: float, - specific_heat_ratio: float - ) -> tuple[float, float, float]: + p_local: float, + rho_total_injector: float, + p_ratio: float, + specific_heat_ratio: float, +) -> tuple[float, float, float]: # NOTE if p_ratio_unchoked <= 1.0 -> no injection. We clip to prevent negative sqrt. p_ratio_unchoked = jnp.clip(p_ratio, 0.0, 1.0) M_unchoked = mach_number_from_pressure_ratio_isentropic(p_ratio_unchoked, specific_heat_ratio) rho_unchoked = rho_total_injector * density_ratio_isentropic(M_unchoked, specific_heat_ratio) u_unchoked = M_unchoked * speed_of_sound(p_local, rho_unchoked, specific_heat_ratio) - E_unchoked = total_energy(p_local, rho_unchoked, u_unchoked, specific_heat_ratio) - return rho_unchoked, u_unchoked, E_unchoked + E_unchoked = total_energy(p_local, rho_unchoked, u_unchoked, specific_heat_ratio) + return rho_unchoked, u_unchoked, E_unchoked def initialize_injector_flux_fn( - injector_geometry: InjectorGeometry, - pressure_ratios: PressureRatios, - p_infty: float, - T_infty: float, - specific_heat_ratio: float, - specific_gas_constant: float, - sim_manager: SimulationManager, - ) -> Callable[ - [Array, Array, Array, Array, Array], - tuple[Array, Array, Array] - ]: - + injector_geometry: InjectorGeometry, + pressure_ratios: PressureRatios, + p_infty: float, + T_infty: float, + specific_heat_ratio: float, + specific_gas_constant: float, + sim_manager: SimulationManager, +) -> Callable[[Array, Array, Array, Array, Array], tuple[Array, Array, Array]]: domain_information = sim_manager.domain_information dim = domain_information.dim dx = domain_information.smallest_cell_size @@ -506,15 +466,14 @@ def initialize_injector_flux_fn( injector_params = _compute_injector_plane_params(injector_geometry) def compute_interface_flux( - primitives: Array, - interface_length: Array, - normal: Array, - mesh_grid: Array, - actuator: Array - ) -> tuple[Array, Array, Array]: - """Computes the interface flux for each injector. - """ - + primitives: Array, + interface_length: Array, + normal: Array, + mesh_grid: Array, + actuator: Array, + ) -> tuple[Array, Array, Array]: + """Computes the interface flux for each injector.""" + if len(actuator) != num_injectors: raise ValueError("Number of actions unequal to injector count") @@ -529,7 +488,6 @@ def compute_interface_flux( energy_flux = 0.0 for i in range(num_injectors): - actuator_i = actuator[i] # actuator [0.0, 1.0] adjusts total pressure @@ -538,7 +496,7 @@ def compute_interface_flux( p=p_total_injector, T=T_infty, R=specific_gas_constant, - ) + ) # we assume isentropic expansion to either M=1 (choked) # or to local pressure to compute injector state @@ -549,14 +507,14 @@ def compute_interface_flux( rho_total_injector=rho_total_injector, specific_heat_ratio=specific_heat_ratio, ) - + # unchoked injector state p_ratio = p_local / p_total_injector rho_unchoked, u_unchoked, E_unchoked = _compute_unchoked_state( p_local=p_local, rho_total_injector=rho_total_injector, p_ratio=p_ratio, - specific_heat_ratio=specific_heat_ratio + specific_heat_ratio=specific_heat_ratio, ) # checking if injector is choked @@ -570,41 +528,39 @@ def compute_interface_flux( E_injector = jnp.where(mask_choked, E_choked, E_unchoked) mask_injector = _compute_injector_mask( - mesh_grid, IW, + mesh_grid, + IW, InjectorPlaneParameters( injector_params.positions[i], injector_params.tangents[i], - injector_params.normals[i] + injector_params.normals[i], ), - dim, dx + dim, + dx, ) - u_injector_vec = u_injector * injector_params.normals[i][:,None] + u_injector_vec = u_injector * injector_params.normals[i][:, None] u_injector_n = jnp.sum(u_injector_vec * normal, axis=0) mass_flux += rho_injector * u_injector_n * interface_length * mask_injector - momentum_flux_i = ( - rho_injector * u_injector_n * u_injector_vec + p_injector * normal - ) * interface_length + momentum_flux_i = (rho_injector * u_injector_n * u_injector_vec + p_injector * normal) * interface_length momentum_flux = momentum_flux * (1 - mask_injector) + momentum_flux_i * mask_injector energy_flux += (E_injector + p_injector) * u_injector_n * interface_length * mask_injector - + return mass_flux, momentum_flux, energy_flux return compute_interface_flux def plot_flowfield_3d( - primitives: ndarray, - levelset: ndarray, - cell_centers: tuple[ndarray, ...], - cell_sizes: tuple[ndarray, ...], - nozzle_pressure: float, - injector_geometry: InjectorGeometry, - specific_heat_ratio: float, - ) -> pv.Plotter: - - + primitives: ndarray, + levelset: ndarray, + cell_centers: tuple[ndarray, ...], + cell_sizes: tuple[ndarray, ...], + nozzle_pressure: float, + injector_geometry: InjectorGeometry, + specific_heat_ratio: float, +) -> pv.Plotter: injector_params = _compute_injector_plane_params(injector_geometry) mesh_grid = np.meshgrid(*cell_centers, indexing="ij") @@ -616,23 +572,25 @@ def plot_flowfield_3d( for i in range(injector_geometry.N): mask_injector = _compute_injector_mask( - mesh_grid.reshape(3,-1), IW, + mesh_grid.reshape(3, -1), + IW, InjectorPlaneParameters( injector_params.positions[i], injector_params.tangents[i], - injector_params.normals[i] + injector_params.normals[i], ), - 3, min_dx + 3, + min_dx, ) mask_injector_list.append(mask_injector.reshape(mesh_grid[0].shape)) # compute fields density = primitives[0] - velocity = np.linalg.norm(primitives[1:4],axis=0,ord=2) + velocity = np.linalg.norm(primitives[1:4], axis=0, ord=2) pressure = primitives[-1] speed_of_sound = np.sqrt(specific_heat_ratio * pressure / density) - mach_number = velocity/speed_of_sound - schlieren = np.linalg.norm(np.gradient(density),axis=0,ord=2) + mach_number = velocity / speed_of_sound + schlieren = np.linalg.norm(np.gradient(density), axis=0, ord=2) min_dx = np.min(cell_sizes[0]) @@ -655,7 +613,7 @@ def plot_flowfield_3d( grid_levelset.cell_data["pressure"] = pressure.ravel(order="F") / nozzle_pressure grid_levelset.cell_data["total_mask"] = total_mask.ravel(order="F") - grid.cell_data["levelset"] = levelset.ravel(order="F")/min_dx + grid.cell_data["levelset"] = levelset.ravel(order="F") / min_dx grid.cell_data["schlieren"] = schlieren.ravel(order="F") grid.cell_data["mach_number"] = mach_number.ravel(order="F") @@ -671,7 +629,6 @@ def plot_flowfield_3d( def clip_grids(grid, grid_levelset): - # --- compute shared geometry once --- points = grid.points # use one grid (same topology assumed) x = points[:, 0] @@ -682,32 +639,19 @@ def clip_grids(grid, grid_levelset): # --- build masks --- nozzle_x_limit = NozzleGeometry().H[0] - 5e-4 - mask_levelset = ( - (r < 0.044) & - (x < nozzle_x_limit) - ) + mask_levelset = (r < 0.044) & (x < nozzle_x_limit) - mask_grid = ( - (r < 0.044) & - (x < 0.25) & - (grid.point_data["levelset"] > 2) - ) + mask_grid = (r < 0.044) & (x < 0.25) & (grid.point_data["levelset"] > 2) # --- extract in one pass --- - grid_levelset_clipped = grid_levelset.extract_points( - mask_levelset, - adjacent_cells=True - ) + grid_levelset_clipped = grid_levelset.extract_points(mask_levelset, adjacent_cells=True) - grid_clipped = grid.extract_points( - mask_grid, - adjacent_cells=True - ) + grid_clipped = grid.extract_points(mask_grid, adjacent_cells=True) return grid_clipped, grid_levelset_clipped -def plot_slice(grid_levelset, grid) -> pv.Plotter: +def plot_slice(grid_levelset, grid) -> pv.Plotter: plotter = pv.Plotter(off_screen=True, window_size=(3000, 2200)) # CMAP = "Spectral_r" @@ -775,13 +719,9 @@ def plot_slice(grid_levelset, grid) -> pv.Plotter: lighting=False, ) - # mach number offset = 0.05 - slice_xy = grid.slice( - normal=(0, 0, 1), - origin=(0.0, 0.0, 0.0) - ) + slice_xy = grid.slice(normal=(0, 0, 1), origin=(0.0, 0.0, 0.0)) slice_xy = slice_xy.threshold( value=0.0, scalars="levelset", @@ -799,10 +739,7 @@ def plot_slice(grid_levelset, grid) -> pv.Plotter: # specular=0.1, ) - slice_xz = grid.slice( - normal=(0, 1, 0), - origin=(0.0, 0.0, 0.0) - ) + slice_xz = grid.slice(normal=(0, 1, 0), origin=(0.0, 0.0, 0.0)) slice_xz = slice_xz.threshold( value=0.0, scalars="levelset", @@ -826,10 +763,7 @@ def plot_slice(grid_levelset, grid) -> pv.Plotter: n = 20 values = np.linspace(low, high, n) - contours = grid.contour( - isosurfaces=values, - scalars="schlieren" - ) + contours = grid.contour(isosurfaces=values, scalars="schlieren") plotter.add_mesh( contours, color="darkgray", @@ -842,16 +776,16 @@ def plot_slice(grid_levelset, grid) -> pv.Plotter: specular_power=100, ) - plotter.add_light(pv.Light(light_type='headlight', intensity=0.3)) + plotter.add_light(pv.Light(light_type="headlight", intensity=0.3)) # plotter.add_axes() # plotter.show_bounds() # plotter.show_grid() plotter.camera_position = ( - (0.25,0.1,0.15), - (0.12,0.0,0.0), - (0,1,0), + (0.25, 0.1, 0.15), + (0.12, 0.0, 0.0), + (0, 1, 0), ) # plotter.enable_eye_dome_lighting() # plotter.set_background("black") - return plotter \ No newline at end of file + return plotter diff --git a/hydrogym/nek/integrate.py b/hydrogym/nek/integrate.py index 5544856b..216c748d 100644 --- a/hydrogym/nek/integrate.py +++ b/hydrogym/nek/integrate.py @@ -6,6 +6,7 @@ import scipy.io as sio from hydrogym.core import CallbackBase +from hydrogym.nek.nek_lib.nek_utils import show_end, show_title def integrate(