PARCtorch.integrator package

Submodules

PARCtorch.integrator.datadrivenintegrator module

class PARCtorch.integrator.datadrivenintegrator.DataDrivenIntegrator(n_io, n_base_features, kernel_size=1, padding_mode='constant', spade_random_noise=True, **kwarg)

Bases: Module

forward(update, current)

Args update (torch.tensor): The change reported by the numerical integrator. Number of channels must equal to `n_io`. current (torch.tensor): The state and/or velocity variables at current time step. Number of channels must equal to `n_io`.

Returns out (torch.tensor): The predicted state and/or velocity variables at the next time step. Shape is the same as `update` and `current`

PARCtorch.integrator.euler module

class PARCtorch.integrator.euler.Euler(**kwarg)

Bases: NumericalIntegrator

forward(f, current, step_size)

Euler integration. Fixed step, 1st order.

Args f (callable): the function that returns time deriviative current (torch.tensor): the current state and velocity variables step_size (float): integration step size

Returns final_state (torch.tensor): the same shape as `current`, the next state and velocity varaibles update (torch.tensor): the same shape as `current`, the updates at this step

PARCtorch.integrator.heun module

class PARCtorch.integrator.heun.Heun(**kwarg)

Bases: NumericalIntegrator

forward(f, current, step_size)

Heun integration. Fixed step, explicit, 2nd order.

Parameters:

  • f (function, the function that returns time deriviative)

  • current (tensor, the current state and velocity variables)

  • step_size (float, integration step size)

Returns:

  • final_state (tensor with the same shape of `current`, the next state and velocity varaibles)

  • update (tensor with the same shape of `current`, the update in this step)

PARCtorch.integrator.integrator module

class PARCtorch.integrator.integrator.Integrator(clip, list_poi_idx, num_int, list_dd_int, padding_mode, finite_difference_method, poi_kernel_size=3, n_poi_features=64, **kwarg)

Bases: Module

Constructor of Integrator

Args clip (bool): whether to clip the state or velocity variable before each integration step list_poi_idx (list[tuple(int)]): List of channel indices for I/O of PoissonBlock.

The last element in each tuple will be the channel index to output to, the 2nd and 3rd last will be assumed as a vector, and the rest as state variables.

num_int (nn.Module): Numerical integrator list_dd_int (list[nn.Module]): List of data driven integrator. One per channel for state variables and one for the velocity as a whole. padding_mode (str): Padding mode finite_difference_method (nn.Module): Numercial method for finite difference calculation. poi_kernel_size (int, optional): Kernel size of the PoissonBlock. Default value is 3. n_poi_features (int, optional): Number of features in the PoissonBlock. Default value is 64. **kwarg: Other arguments that will be passed to nn.Module during initialization.

forward(f, ic, t0, t1)

Forward of Integrator. It will clip the current state and velocity variable (if necessary), go through the numerical integrator and then datadriven integrator.

Args f (callable): Callable that returns time derivative ic (torch.tensor): 4-d tensor of Float with the shape (batch_size, channels, y, x), the initial condition to start integrating from t0 (float): Starting time t1 (torch.tensor): 1-d tensor of Float, the time points to integrate to

Returns res (torch.tensor): 5-d tensor of Float with the shape (timesteps, batch_size, channels, y, x), the predicted state and velocity variables at each time in t1

PARCtorch.integrator.numintegrator module

class PARCtorch.integrator.numintegrator.NumericalIntegrator(**kwarg)

Bases: Module

forward(f, current, step_size)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

PARCtorch.integrator.poisson module

class PARCtorch.integrator.poisson.Poisson(finite_difference_method)

Bases: Module

Computes Poisson terms based on the vector field using finite difference filters.

Parameters:

  • channel_size (int) – Number of input channels for each vector component.

  • cd_filter_1d (np.array) – 1D filter for finite difference (e.g., np.array([-1.0, 1.0])).

  • padding (str) – Padding mode. “SYMMETRIC” in TensorFlow corresponds to “reflect” in PyTorch.

forward(vector_field)

Forward pass to compute Poisson terms from the vector field.

Parameters:

vector_field (torch.Tensor) – Tensor of shape [N, 2, H, W].

Returns:

Element-wise square of the x-derivative of the first component, shape [N, C, H, W] vy2 (torch.Tensor): Element-wise square of the y-derivative of the second component, shape [N, C, H, W] uyvx (torch.Tensor): Element-wise product of the y-derivative of the first component and the x-derivative of the second component, shape [N, C, H, W]

Return type:

ux2 (torch.Tensor)

class PARCtorch.integrator.poisson.PoissonBlock(n_input_channel, finite_difference_method, kernel_size=3, n_base_features=64, padding_mode='constant')

Bases: Module

Poisson Block to compute pressure terms from input vector fields using existing UnetDownBlock and ResNetBlock.

Parameters:

  • n_input_channel (int) – Number of input channels. We always assume the last 2 channels will be the one to run the Poisson operator on.

  • finite_difference_method (nn.Module) – Numerical method for finite difference.

  • kernel_size (int, optional) – Kernel size. Default is 3.

  • n_base_features (int, optional) – Number of channels of the convolutional layers. Default is 64.

  • padding_mode (str, optional) – Padding mode for the convolutional layers. Default is “constant”.

forward(x)

Forward pass to compute pressure terms.

Parameters:

x (torch.Tensor) – Input tensor of shape [N, n_input_channel, y, x].

Returns:

Output tensor of shape [N, 1, y, x].

Return type:

torch.Tensor

PARCtorch.integrator.rk4 module

class PARCtorch.integrator.rk4.RK4(**kwarg)

Bases: NumericalIntegrator

forward(f, current, step_size)

RK4 integration. Fixed step, 4th order.

Parameters:

  • f (function, the function that returns time deriviative)

  • current (tensor, the current state and velocity variables)

  • step_size (float, integration step size)

Returns:

  • final_state (tensor with the same shape of `current`, the next state and velocity varaibles)

  • update (tensor with the same shape of `current`, the update in this step)

Module contents