Skip to main content

hydrogym.jax.flow

FlowConfig Objects​

class FlowConfig(PDEBase)

DEFAULT_OBS_SIZE​

This correlates to a total observation size of 8x8 = 64.

load_mesh​

def load_mesh(name)

Create jax grid given the desired dimensions and spacing in real space

Returns:

jax meshgrid

load_fft_mesh​

def load_fft_mesh()

Create jax grid given desired dimensions and spacing in real Fourier space

Returns:

jax meshgrid

initialize_state​

def initialize_state()

Generate a divergence free velocity field to initialize the state Initializing with divergence free field specified with the following stream function:

φ(x,y) = sin(x)cos(y)

Returns:

fft vorticity field

forcing_function​

def forcing_function(k, x, y)

Sinusoidal forcing function that drives the Kolmogorov flow.

Arguments:

  • k int - forcing wavenumber
  • x jnp.array - spatial coordinates in x
  • y jnp.array - spatial coordinates in y

Returns:

  • tuple - forcing function in (x,y)

num_inputs​

@property
def num_inputs() -> int

Length of the control vector (number of actuators)

num_outputs​

@property
def num_outputs() -> int

Number of scalar observed variables

save_checkpoint​

def save_checkpoint()

Set up mesh, function spaces, state vector, etc

init_bcs​

def init_bcs()

Initialize any boundary conditions for the PDE.

copy_state​

def copy_state(deepcopy=True)

Return a copy of the flow state

evaluate_objective​

def evaluate_objective()

Return a copy of the flow state

render​

def render(**kwargs)

Plot the current PDE state (called by gym.Env)