PARCtorch.differentiator package

Submodules

PARCtorch.differentiator.advection module

class PARCtorch.differentiator.advection.Advection(finite_difference_method)

Bases: Module

Computes the advection term based on the state variable and velocity field.

Parameters:

finite_difference_method (nn.Module) – Numerical method to calculate spatial deriviatives

forward(state_variable, velocity_field)

Forward pass to compute advection.

Parameters:

  • state_variable (torch.Tensor) – Tensor of shape [N, C, H, W]

  • velocity_field (torch.Tensor) – Tensor of shape [N, 2C, H, W]

Returns:

Advection term of shape [N, 1, H, W]

Return type:

advect (torch.Tensor)

PARCtorch.differentiator.differentiator module

class PARCtorch.differentiator.differentiator.Differentiator(n_state_var, n_fe_features, list_adv_idx, list_dif_idx, feature_extraction, padding_mode, finite_difference_method, spade_random_noise, **kwarg)

Bases: Module

forward(current)

Forward of differentiator. Advection and diffusion will be calculated per channel for those necessary and combined with dynamic features. Those that do not have explicit advection and diffusion calculation will have zero has output. This design choice was made because of certain integrtors (e.g. those in torchdiffeq) requires differentiator to have the same output and input shape.

Parameters:

current (4-d tensor of Float with shape (batch_size, channels, y, x), the current state and velocity variables)

Returns:

t_dot

Return type:

4-d tensor of Float with the same shape as `current`, the predicted time deriviatives on current state and velocity variables

PARCtorch.differentiator.diffusion module

class PARCtorch.differentiator.diffusion.Diffusion(finite_difference_method)

Bases: Module

Computes the Laplacian of the state variable using finite difference filters.

Parameters:

finite_difference_method (nn.Module) – Numerical method to calculate spatial deriviatives

forward(state_variable)

Forward pass to compute the Laplacian of the state variable.

Parameters:

state_variable (torch.Tensor) – Tensor of shape [N, C, H, W]

Returns:

Laplacian of shape [N, C, H, W]

Return type:

laplacian (torch.Tensor)

PARCtorch.differentiator.finitedifference module

class PARCtorch.differentiator.finitedifference.FiniteDifference(channel_size=1, filter_1d=array([-1., 1.]), padding_mode='replicate', device='cuda')

Bases: Module

Computes spatial derivatives using finite difference filters.

Parameters:

  • channel_size (int) – Number of input channels.

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

  • padding_mode (str) – Padding mode. Default value set to reproduce torchmetrics.functional.image.image_gradients.

forward(x)

Forward pass to compute spatial derivatives.

Parameters:

x (torch.Tensor) – Input tensor of shape [N, C, H, W]

Returns:

Derivative along y-axis, shape [N, C, H, W] dx (torch.Tensor): Derivative along x-axis, shape [N, C, H, W]

Return type:

dy (torch.Tensor)

PARCtorch.differentiator.mappingandrecon module

class PARCtorch.differentiator.mappingandrecon.MappingAndRecon(n_base_features=128, n_mask_channel=1, output_channel=1, padding_mode='zeros', add_noise=True)

Bases: Module

forward(dynamic_feature, advec_diff)

Forward pass of the MappingAndRecon.

Parameters:

  • dynamic_feature (torch.Tensor) – Tensor of shape [N, C, H, W], dynamic features from the feature extraction network

  • advec_diff (torch.Tensor) – Tensor of shape [N, M, H, W], channel-concatenated advection and diffusion

Returns:

Output tensor of shape [N, output_channel, H, W]

Return type:

torch.Tensor

Module contents