SYMBOL INDEX (412 symbols across 60 files)

FILE: baselines/data.py
  class DarcyFlow (line 10) | class DarcyFlow(Dataset):
    method __init__ (line 11) | def __init__(self,
    method __len__ (line 24) | def __len__(self):
    method __getitem__ (line 27) | def __getitem__(self, item):
  class NSLong (line 32) | class NSLong(object):
    method __init__ (line 33) | def __init__(self,
    method get_boundary_value (line 62) | def get_boundary_value(self, component=0):
    method get_boundary_points (line 82) | def get_boundary_points(self, num_x, num_y, num_t):
    method get_test_xyt (line 88) | def get_test_xyt(self):
  class NSdata (line 106) | class NSdata(object):
    method __init__ (line 107) | def __init__(self, datapath1,
    method get_init_cond (line 151) | def get_init_cond(self):
    method get_boundary_value (line 157) | def get_boundary_value(self, component=0):
    method get_boundary_points (line 177) | def get_boundary_points(self, num_x, num_y, num_t):
    method get_test_xyt (line 225) | def get_test_xyt(self):
    method extract (line 244) | def extract(data):
  class DeepOnetNS (line 270) | class DeepOnetNS(Dataset):
    method __init__ (line 274) | def __init__(self, datapath,
    method __len__ (line 296) | def __len__(self):
    method __getitem__ (line 299) | def __getitem__(self, idx):
  class DeepONetCPNS (line 309) | class DeepONetCPNS(Dataset):
    method __init__ (line 314) | def __init__(self, datapath,
    method __len__ (line 336) | def __len__(self):
    method __getitem__ (line 339) | def __getitem__(self, idx):

FILE: baselines/deepxde_deeponet.py
  function train (line 11) | def train(config):

FILE: baselines/loss.py
  function boundary_loss (line 7) | def boundary_loss(model, npt=100):
  function resf_NS (line 24) | def resf_NS(u, v, p, x, y, t, re=40):

FILE: baselines/model.py
  class DeepONet (line 8) | class DeepONet(nn.Module):
    method __init__ (line 9) | def __init__(self, branch_layer, trunk_layer):
    method forward (line 14) | def forward(self, u0, grid):
  class DeepONetCP (line 22) | class DeepONetCP(nn.Module):
    method __init__ (line 23) | def __init__(self, branch_layer, trunk_layer):
    method forward (line 28) | def forward(self, u0, grid):
  class SAWeight (line 36) | class SAWeight(nn.Module):
    method __init__ (line 37) | def __init__(self, out_dim, num_init: List, num_bd: List, num_collo: L...
    method forward (line 51) | def forward(self, init_cond: List, bd_cond: List, residual: List):

FILE: baselines/pinns_ns_05s.py
  function forcing (line 17) | def forcing(x):
  function pde (line 21) | def pde(x, u):
  function eval (line 54) | def eval(model, dataset,
  function train (line 83) | def train(offset, config, args):

FILE: baselines/pinns_ns_50s.py
  function forcing (line 17) | def forcing(x):
  function pde (line 22) | def pde(x, u):
  function eval (line 55) | def eval(model, dataset,
  function train_longtime (line 102) | def train_longtime(offset, config, args):

FILE: baselines/sapinns-50s.py
  function forcing (line 20) | def forcing(x):
  function pde (line 25) | def pde(x, u):
  function eval (line 59) | def eval(model, dataset,
  function train_sapinn (line 90) | def train_sapinn(offset, config, args):

FILE: baselines/sapinns.py
  function forcing (line 22) | def forcing(x):
  function pde (line 26) | def pde(x, u):
  function eval (line 59) | def eval(model, dataset,
  function train_sapinn (line 90) | def train_sapinn(offset, config, args):

FILE: baselines/test.py
  function test (line 11) | def test(model,
  function test_deeponet_ns (line 42) | def test_deeponet_ns(config):
  function test_deeponet_darcy (line 70) | def test_deeponet_darcy(config):

FILE: baselines/tqd_sapinns.py
  function forcing (line 19) | def forcing(x):
  function bd_model (line 23) | def bd_model(u_model, x, y, t):
  function f_model (line 29) | def f_model(u_model, x, y, t):
  function eval (line 59) | def eval(model, dataset,
  function train_sa (line 88) | def train_sa(offset, config, args):

FILE: baselines/tqd_utils.py
  class PointsIC (line 7) | class PointsIC(BC):
    method __init__ (line 11) | def __init__(self, domain, values, var, n_values=None):
    method create_input (line 29) | def create_input(self):
    method create_target (line 40) | def create_target(self, values):
    method loss (line 53) | def loss(self):

FILE: baselines/train_darcy.py
  function train_deeponet_darcy (line 14) | def train_deeponet_darcy(config):

FILE: baselines/train_ns.py
  function train_deeponet_cp (line 15) | def train_deeponet_cp(config):
  function train_deeponet (line 74) | def train_deeponet(config):

FILE: baselines/unet3d.py
  class BaseModel (line 12) | class BaseModel(nn.Module):
    method __init__ (line 13) | def __init__(self):
    method device (line 18) | def device(self):
    method data_dict_to_input (line 22) | def data_dict_to_input(self, data_dict, **kwargs):
    method loss_dict (line 28) | def loss_dict(self, data_dict, **kwargs):
    method eval_dict (line 35) | def eval_dict(self, data_dict, **kwargs):
  function create_conv (line 41) | def create_conv(
  class SingleConv (line 129) | class SingleConv(nn.Sequential):
    method __init__ (line 148) | def __init__(
  class DoubleConv (line 166) | class DoubleConv(nn.Sequential):
    method __init__ (line 190) | def __init__(
  class Encoder (line 242) | class Encoder(nn.Module):
    method __init__ (line 265) | def __init__(
    method forward (line 306) | def forward(self, x):
  class Decoder (line 313) | class Decoder(nn.Module):
    method __init__ (line 334) | def __init__(
    method forward (line 385) | def forward(self, encoder_features, x):
    method _joining (line 392) | def _joining(encoder_features, x, concat):
  function create_encoders (line 399) | def create_encoders(
  function create_decoders (line 444) | def create_decoders(
  class AbstractUpsampling (line 472) | class AbstractUpsampling(nn.Module):
    method __init__ (line 478) | def __init__(self, upsample):
    method forward (line 482) | def forward(self, encoder_features, x):
  class InterpolateUpsampling (line 489) | class InterpolateUpsampling(AbstractUpsampling):
    method __init__ (line 497) | def __init__(self, mode="nearest"):
    method _interpolate (line 502) | def _interpolate(x, size, mode):
  class TransposeConvUpsampling (line 506) | class TransposeConvUpsampling(AbstractUpsampling):
    method __init__ (line 520) | def __init__(
  class NoUpsampling (line 534) | class NoUpsampling(AbstractUpsampling):
    method __init__ (line 535) | def __init__(self):
    method _no_upsampling (line 539) | def _no_upsampling(x, size):
  function number_of_features_per_level (line 543) | def number_of_features_per_level(init_channel_number, num_levels):
  class AbstractUNet (line 547) | class AbstractUNet(BaseModel):
    method __init__ (line 577) | def __init__(
    method forward (line 645) | def forward(self, x):
  class UNet3D (line 673) | class UNet3D(AbstractUNet):
    method __init__ (line 682) | def __init__(

FILE: baselines/utils.py
  function weighted_mse (line 7) | def weighted_mse(pred, target, weight=None):
  function get_3dboundary_points (line 14) | def get_3dboundary_points(num_x,                # number of points on x ...
  function get_3dboundary (line 57) | def get_3dboundary(value):
  function get_xytgrid (line 78) | def get_xytgrid(S, T, bot=[0, 0, 0], top=[1, 1, 1]):
  function get_2dgird (line 100) | def get_2dgird(num=31):
  function get_3dgrid (line 110) | def get_3dgrid(num=11):
  function get_4dgrid (line 121) | def get_4dgrid(num=11):
  function vel2vor (line 143) | def vel2vor(u, v, x, y):
  function sub_mse (line 150) | def sub_mse(vec):
  function get_sample (line 164) | def get_sample(npt=100):
  function concat (line 179) | def concat(xy, z, t=0.0, offset=0):
  function cal_mixgrad (line 200) | def cal_mixgrad(outputs, inputs):

FILE: cavity_flow.py
  class SpectralConv3d (line 29) | class SpectralConv3d(nn.Module):
    method __init__ (line 30) | def __init__(self, in_channels, out_channels, modes1, modes2, modes3):
    method compl_mul3d (line 58) | def compl_mul3d(self, input, weights):
    method forward (line 62) | def forward(self, x):
  class FNO3d (line 84) | class FNO3d(nn.Module):
    method __init__ (line 85) | def __init__(self, modes1, modes2, modes3, width, padding):
    method forward (line 122) | def forward(self, x):
    method get_grid (line 157) | def get_grid(self, shape, device):
  function PINO_loss_Fourier_f (line 215) | def PINO_loss_Fourier_f(out, Re=500):
  function PINO_loss_FDM_f (line 260) | def PINO_loss_FDM_f(out, Re=500):
  function PINO_loss_ic (line 289) | def PINO_loss_ic(out, y):
  function PINO_loss_bc (line 299) | def PINO_loss_bc(out, y):

FILE: download_data.py
  function download_file (line 23) | def download_file(url, file_path):
  function main (line 33) | def main(args):

FILE: eval_operator.py
  function test_3d (line 14) | def test_3d(config):
  function test_2d (line 48) | def test_2d(config):

FILE: generate_data.py
  function legacy_solver (line 15) | def legacy_solver(args):
  function gen_data (line 58) | def gen_data(args):

FILE: inference.py
  function get_pred (line 20) | def get_pred(args):

FILE: instance_opt.py
  function train_ns (line 25) | def train_ns(model,
  function subprocess (line 122) | def subprocess(args):

FILE: inverse-darcy-foward.py
  class SpectralConv2d (line 24) | class SpectralConv2d(nn.Module):
    method __init__ (line 25) | def __init__(self, in_channels, out_channels, modes1, modes2):
    method compl_mul2d (line 44) | def compl_mul2d(self, input, weights):
    method forward (line 48) | def forward(self, x):
  class FNO2d (line 66) | class FNO2d(nn.Module):
    method __init__ (line 67) | def __init__(self, modes1, modes2, width):
    method forward (line 102) | def forward(self, x):
    method get_grid (line 139) | def get_grid(self, shape, device):
  function FDM_Darcy (line 214) | def FDM_Darcy(u, a, D=1, f=1):
  function PINO_loss (line 248) | def PINO_loss(u, a):
  function darcy_mask1 (line 349) | def darcy_mask1(x):
  function darcy_mask2 (line 352) | def darcy_mask2(x):
  function total_variance (line 359) | def total_variance(x):

FILE: inverse-darcy.py
  class SpectralConv2d (line 22) | class SpectralConv2d(nn.Module):
    method __init__ (line 23) | def __init__(self, in_channels, out_channels, modes1, modes2):
    method compl_mul2d (line 42) | def compl_mul2d(self, input, weights):
    method forward (line 46) | def forward(self, x):
  class FNO2d (line 64) | class FNO2d(nn.Module):
    method __init__ (line 65) | def __init__(self, modes1, modes2, width):
    method forward (line 100) | def forward(self, x):
    method get_grid (line 136) | def get_grid(self, shape, device):
  function FDM_Darcy (line 211) | def FDM_Darcy(u, a, D=1, f=1):
  function PINO_loss (line 245) | def PINO_loss(u, a):
  function darcy_mask1 (line 279) | def darcy_mask1(x):
  function darcy_mask2 (line 282) | def darcy_mask2(x):
  function total_variance (line 289) | def total_variance(x):

FILE: models/FCN.py
  function linear_block (line 4) | def linear_block(in_channel, out_channel):
  class FCNet (line 12) | class FCNet(nn.Module):
    method __init__ (line 18) | def __init__(self, layers=[2, 10, 1]):
    method forward (line 26) | def forward(self, x):
  class DenseNet (line 30) | class DenseNet(nn.Module):
    method __init__ (line 31) | def __init__(self, layers, nonlinearity, out_nonlinearity=None, normal...
    method forward (line 57) | def forward(self, x):

FILE: models/basics.py
  function compl_mul1d (line 8) | def compl_mul1d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
  function compl_mul2d (line 15) | def compl_mul2d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
  function compl_mul3d (line 22) | def compl_mul3d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
  class SpectralConv1d (line 31) | class SpectralConv1d(nn.Module):
    method __init__ (line 32) | def __init__(self, in_channels, out_channels, modes1):
    method forward (line 48) | def forward(self, x):
  class SpectralConv2d (line 66) | class SpectralConv2d(nn.Module):
    method __init__ (line 67) | def __init__(self, in_channels, out_channels, modes1, modes2):
    method forward (line 81) | def forward(self, x):
  class SpectralConv3d (line 101) | class SpectralConv3d(nn.Module):
    method __init__ (line 102) | def __init__(self, in_channels, out_channels, modes1, modes2, modes3):
    method forward (line 116) | def forward(self, x):
  class FourierBlock (line 148) | class FourierBlock(nn.Module):
    method __init__ (line 149) | def __init__(self, in_channels, out_channels, modes1, modes2, modes3, ...
    method forward (line 164) | def forward(self, x):

FILE: models/core.py
  function contract_1D (line 7) | def contract_1D(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
  function contract_2D (line 13) | def contract_2D(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
  function contract_3D (line 19) | def contract_3D(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
  class FactorizedSpectralConv3d (line 24) | class FactorizedSpectralConv3d(nn.Module):
    method __init__ (line 25) | def __init__(self, in_channels, out_channels, modes_height, modes_widt...
    method _get_weight_factorized (line 70) | def _get_weight_factorized(self, layer_index, corner_index):
    method _get_weight_dense (line 77) | def _get_weight_dense(self, layer_index, corner_index):
    method forward (line 84) | def forward(self, x, indices=0):
    method get_conv (line 116) | def get_conv(self, indices):
    method __getitem__ (line 125) | def __getitem__(self, indices):
  class FactorizedSpectralConv2d (line 129) | class FactorizedSpectralConv2d(nn.Module):
    method __init__ (line 130) | def __init__(self, in_channels, out_channels, modes_height, modes_widt...
    method _get_weight_factorized (line 170) | def _get_weight_factorized(self, layer_index, corner_index):
    method _get_weight_dense (line 177) | def _get_weight_dense(self, layer_index, corner_index):
    method forward (line 184) | def forward(self, x, indices=0, super_res=1):
    method get_conv (line 208) | def get_conv(self, indices):
    method __getitem__ (line 217) | def __getitem__(self, indices):
  class SubConv (line 221) | class SubConv(nn.Module):
    method __init__ (line 229) | def __init__(self, main_conv, indices):
    method forward (line 234) | def forward(self, x, **kwargs):
  class FactorizedSpectralConv1d (line 238) | class FactorizedSpectralConv1d(nn.Module):
    method __init__ (line 239) | def __init__(self, in_channels, out_channels, modes, n_layers=1,
    method _get_weight_factorized (line 282) | def _get_weight_factorized(self, layer_index):
    method _get_weight_dense (line 286) | def _get_weight_dense(self, layer_index):
    method forward (line 290) | def forward(self, x, indices=0, s=None):
    method get_conv (line 309) | def get_conv(self, indices):
    method __getitem__ (line 318) | def __getitem__(self, indices):
  class JointFactorizedSpectralConv1d (line 323) | class JointFactorizedSpectralConv1d(nn.Module):
    method __init__ (line 324) | def __init__(self, modes, width, n_layers=1, joint_factorization=True,...
    method forward (line 384) | def forward(self, x, s=None):

FILE: models/fourier1d.py
  class FNO1d (line 6) | class FNO1d(nn.Module):
    method __init__ (line 7) | def __init__(self,
    method forward (line 45) | def forward(self, x):

FILE: models/fourier2d.py
  class FNO2d (line 6) | class FNO2d(nn.Module):
    method __init__ (line 7) | def __init__(self, modes1, modes2,
    method forward (line 53) | def forward(self, x):

FILE: models/fourier3d.py
  class FNO3d (line 6) | class FNO3d(nn.Module):
    method __init__ (line 7) | def __init__(self,
    method forward (line 58) | def forward(self, x):

FILE: models/lowrank2d.py
  class LowRank2d (line 8) | class LowRank2d(nn.Module):
    method __init__ (line 9) | def __init__(self, in_channels, out_channels):
    method get_grid (line 17) | def get_grid(self, S1, S2, batchsize, device):
    method forward (line 24) | def forward(self, x, gridy=None):

FILE: models/tfno.py
  class FactorizedFNO3d (line 6) | class FactorizedFNO3d(nn.Module):
    method __init__ (line 7) | def __init__(self, modes_height, modes_width,  modes_depth, width, fc_...
    method forward (line 65) | def forward(self, x, super_res=1):
  class FactorizedFNO2d (line 98) | class FactorizedFNO2d(nn.Module):
    method __init__ (line 99) | def __init__(self, modes_height, modes_width,  width, fc_channels=256,...
    method forward (line 160) | def forward(self, x, super_res=1):
  class FactorizedFNO1d (line 202) | class FactorizedFNO1d(nn.Module):
    method __init__ (line 203) | def __init__(self, modes, width, in_channels=2, out_channels=1, n_laye...
    method forward (line 231) | def forward(self, x, s=None):

FILE: models/utils.py
  function add_padding (line 4) | def add_padding(x, num_pad):
  function add_padding2 (line 12) | def add_padding2(x, num_pad1, num_pad2):
  function remove_padding (line 20) | def remove_padding(x, num_pad):
  function remove_padding2 (line 28) | def remove_padding2(x, num_pad1, num_pad2):
  function _get_act (line 36) | def _get_act(act):

FILE: prepare_data.py
  function shuffle_data (line 5) | def shuffle_data(datapath):
  function test_data (line 13) | def test_data(datapath):
  function get_slice (line 21) | def get_slice(datapath):
  function plot_test (line 30) | def plot_test(datapath):

FILE: run_pino2d.py
  function train (line 14) | def train(args, config):

FILE: run_pino3d.py
  function run_instance (line 15) | def run_instance(loader, config, data_config):

FILE: run_solver.py
  function solve (line 15) | def solve(a,

FILE: solver/kolmogorov_flow.py
  class KolmogorovFlow2d (line 5) | class KolmogorovFlow2d(object):
    method __init__ (line 7) | def __init__(self, w0, Re, n):
    method vorticity (line 59) | def vorticity(self, stream_f=None, real_space=True):
    method stream_function (line 71) | def stream_function(self, w_h=None, real_space=False):
    method velocity_field (line 86) | def velocity_field(self, stream_f=None, real_space=True):
    method nonlinear_term (line 104) | def nonlinear_term(self, w_h):
    method advance (line 127) | def advance(self, t, delta_t=1e-3):

FILE: solver/legacy_solver.py
  class GaussianRF (line 8) | class GaussianRF(object):
    method __init__ (line 10) | def __init__(self, dim, size, alpha=2, tau=3, sigma=None, boundary="pe...
    method sample (line 54) | def sample(self, N):
  function navier_stokes_2d (line 72) | def navier_stokes_2d(w0, f, visc, T, delta_t=1e-4, record_steps=1):

FILE: solver/periodic.py
  class Poisson2d (line 9) | class Poisson2d(object):
    method __init__ (line 11) | def __init__(self, s1, s2, L1=2*math.pi, L2=2*math.pi, device=None, dt...
    method solve (line 28) | def solve(self, f):
    method __call__ (line 31) | def __call__(self, f):
  class NavierStokes2d (line 38) | class NavierStokes2d(object):
    method __init__ (line 40) | def __init__(self, s1, s2, L1=2*math.pi, L2=2*math.pi, device=None, dt...
    method stream_function (line 81) | def stream_function(self, w_h, real_space=False):
    method velocity_field (line 91) | def velocity_field(self, stream_f, real_space=True):
    method nonlinear_term (line 104) | def nonlinear_term(self, w_h, f_h=None):
    method time_step (line 120) | def time_step(self, q, v, f, Re):
    method advance (line 139) | def advance(self, w, f=None, T=1.0, Re=100, adaptive=True, delta_t=1e-3):
    method __call__ (line 185) | def __call__(self, w, f=None, T=1.0, Re=100, adaptive=True, delta_t=1e...

FILE: solver/random_fields.py
  class GaussianRF (line 8) | class GaussianRF(object):
    method __init__ (line 9) | def __init__(self, dim, size, length=1.0, alpha=2.0, tau=3.0, sigma=No...
    method sample (line 67) | def sample(self, N):
  class GaussianRF2d (line 76) | class GaussianRF2d(object):
    method __init__ (line 78) | def __init__(self, s1, s2, L1=2*math.pi, L2=2*math.pi, alpha=2.0, tau=...
    method sample (line 107) | def sample(self, N, xi=None):

FILE: solver/rfsampler.py
  class GaussianRF (line 5) | class GaussianRF(object):
    method __init__ (line 6) | def __init__(self, dim, size, alpha=2, tau=3, sigma=None, boundary="pe...
    method sample (line 51) | def sample(self, N):

FILE: solver/spectrum.py
  function navier_stokes_2d (line 17) | def navier_stokes_2d(w0, f, visc, T, delta_t=1e-4, record_steps=1):

FILE: train_PINO3d.py
  function subprocess_fn (line 20) | def subprocess_fn(rank, args):

FILE: train_burgers.py
  function run (line 12) | def run(args, config):
  function test (line 50) | def test(config):

FILE: train_darcy.py
  function get_molifier (line 26) | def get_molifier(mesh, device):
  function eval_darcy (line 32) | def eval_darcy(model, val_loader, criterion,
  function train (line 50) | def train(model,
  function subprocess (line 139) | def subprocess(args):

FILE: train_no.py
  function pad_input (line 26) | def pad_input(x, num_pad):
  function train_ns (line 34) | def train_ns(model,
  function eval_ns (line 158) | def eval_ns(model, val_loader, device, config, args):
  function subprocess (line 186) | def subprocess(args):

FILE: train_operator.py
  function train_3d (line 15) | def train_3d(args, config):
  function train_2d (line 78) | def train_2d(args, config):

FILE: train_pino.py
  function eval_ns (line 28) | def eval_ns(model, val_loader, criterion, device):
  function train_ns (line 44) | def train_ns(model,
  function subprocess (line 141) | def subprocess(args):

FILE: train_unet.py
  function eval_ns (line 28) | def eval_ns(model, val_loader, criterion, device):
  function train_ns (line 46) | def train_ns(model,
  function subprocess (line 121) | def subprocess(args):

FILE: train_utils/adam.py
  function adam (line 8) | def adam(params: List[Tensor],
  class Adam (line 54) | class Adam(Optimizer):
    method __init__ (line 79) | def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
    method __setstate__ (line 95) | def __setstate__(self, state):
    method step (line 101) | def step(self, closure=None):

FILE: train_utils/data_utils.py
  function sample_data (line 4) | def sample_data(loader):
  function data_sampler (line 10) | def data_sampler(dataset, shuffle, distributed):

FILE: train_utils/datasets.py
  function online_loader (line 15) | def online_loader(sampler, S, T, time_scale, batchsize=1):
  function sample_data (line 24) | def sample_data(loader):
  class MatReader (line 30) | class MatReader(object):
    method __init__ (line 31) | def __init__(self, file_path, to_torch=True, to_cuda=False, to_float=T...
    method _load_file (line 44) | def _load_file(self):
    method load_file (line 48) | def load_file(self, file_path):
    method read_field (line 52) | def read_field(self, field):
    method set_cuda (line 70) | def set_cuda(self, to_cuda):
    method set_torch (line 73) | def set_torch(self, to_torch):
    method set_float (line 76) | def set_float(self, to_float):
  class BurgersLoader (line 80) | class BurgersLoader(object):
    method __init__ (line 81) | def __init__(self, datapath, nx=2 ** 10, nt=100, sub=8, sub_t=1, new=F...
    method make_loader (line 94) | def make_loader(self, n_sample, batch_size, start=0, train=True):
  class NSLoader (line 117) | class NSLoader(object):
    method __init__ (line 118) | def __init__(self, datapath1,
    method make_loader (line 154) | def make_loader(self, n_sample, batch_size, start=0, train=True):
    method make_dataset (line 169) | def make_dataset(self, n_sample, start=0, train=True):
    method extract (line 187) | def extract(data):
  class NS3DDataset (line 213) | class NS3DDataset(Dataset):
    method __init__ (line 214) | def __init__(self, paths,
    method load (line 231) | def load(self, train=True, sub_x=1, sub_t=1):
    method __getitem__ (line 263) | def __getitem__(self, idx):
    method __len__ (line 266) | def __len__(self, ):
    method extract (line 270) | def extract(data):
  class KFDataset (line 296) | class KFDataset(Dataset):
    method __init__ (line 297) | def __init__(self, paths,
    method load (line 327) | def load(self):
    method partition (line 360) | def partition(self, data):
    method __getitem__ (line 378) | def __getitem__(self, idx):
    method __len__ (line 385) | def __len__(self, ):
  class BurgerData (line 389) | class BurgerData(Dataset):
    method __init__ (line 399) | def __init__(self, datapath):
    method __len__ (line 411) | def __len__(self):
    method __getitem__ (line 414) | def __getitem__(self, idx):
    method get_flatten_data (line 417) | def get_flatten_data(self):
    method get_boundary_data (line 428) | def get_boundary_data(self):
    method sample_xt (line 438) | def sample_xt(self, N=10000):
    method sample_xu (line 448) | def sample_xu(self, N=100):
  class DarcyFlow (line 461) | class DarcyFlow(Dataset):
    method __init__ (line 462) | def __init__(self,
    method __len__ (line 475) | def __len__(self):
    method __getitem__ (line 478) | def __getitem__(self, item):
  class DarcyIC (line 483) | class DarcyIC(Dataset):
    method __init__ (line 484) | def __init__(self,
    method __len__ (line 501) | def __len__(self):
    method __getitem__ (line 504) | def __getitem__(self, item):
  class DarcyCombo (line 509) | class DarcyCombo(Dataset):
    method __init__ (line 510) | def __init__(self,
    method __len__ (line 527) | def __len__(self):
    method __getitem__ (line 530) | def __getitem__(self, item):
  class KFaDataset (line 540) | class KFaDataset(Dataset):
    method __init__ (line 541) | def __init__(self, paths,
    method load (line 561) | def load(self):
    method __getitem__ (line 584) | def __getitem__(self, idx):
    method __len__ (line 591) | def __len__(self, ):

FILE: train_utils/distributed.py
  function setup (line 6) | def setup(rank, world_size):
  function cleanup (line 12) | def cleanup():
  function get_world_size (line 16) | def get_world_size():
  function all_reduce_mean (line 23) | def all_reduce_mean(tensor):
  function reduce_sum (line 37) | def reduce_sum(tensor):
  function reduce_loss_dict (line 52) | def reduce_loss_dict(loss_dict):

FILE: train_utils/eval_2d.py
  function eval_darcy (line 14) | def eval_darcy(model,
  function eval_burgers (line 61) | def eval_burgers(model,

FILE: train_utils/eval_3d.py
  function eval_ns (line 15) | def eval_ns(model,  # model

FILE: train_utils/losses.py
  function FDM_Darcy (line 6) | def FDM_Darcy(u, a, D=1):
  function darcy_loss (line 39) | def darcy_loss(u, a):
  function FDM_NS_vorticity (line 68) | def FDM_NS_vorticity(w, v=1/40, t_interval=1.0):
  function Autograd_Burgers (line 108) | def Autograd_Burgers(u, grid, v=1/100):
  function AD_loss (line 119) | def AD_loss(u, u0, grid, index_ic=None, p=None, q=None):
  class LpLoss (line 152) | class LpLoss(object):
    method __init__ (line 156) | def __init__(self, d=2, p=2, size_average=True, reduction=True):
    method abs (line 167) | def abs(self, x, y):
    method rel (line 183) | def rel(self, x, y):
    method __call__ (line 197) | def __call__(self, x, y):
  function FDM_Burgers (line 201) | def FDM_Burgers(u, v, D=1):
  function PINO_loss (line 224) | def PINO_loss(u, u0, v):
  function PINO_loss3d (line 247) | def PINO_loss3d(u, u0, forcing, v=1/40, t_interval=1.0):
  function PDELoss (line 266) | def PDELoss(model, x, t, nu):
  function get_forcing (line 289) | def get_forcing(S):

FILE: train_utils/negadam.py
  function adam (line 8) | def adam(params: List[Tensor],
  class NAdam (line 54) | class NAdam(Optimizer):
    method __init__ (line 79) | def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
    method __setstate__ (line 95) | def __setstate__(self, state):
    method step (line 101) | def step(self, closure=None):

FILE: train_utils/train_2d.py
  function train_2d_operator (line 13) | def train_2d_operator(model,
  function train_2d_burger (line 119) | def train_2d_burger(model,

FILE: train_utils/train_3d.py
  function train (line 16) | def train(model,
  function mixed_train (line 125) | def mixed_train(model,              # model of neural operator
  function progressive_train (line 258) | def progressive_train(model,

FILE: train_utils/utils.py
  function vor2vel (line 6) | def vor2vel(w, L=2 * np.pi):
  function get_sample (line 45) | def get_sample(N, T, s, p, q):
  function get_grid (line 72) | def get_grid(N, T, s):
  function get_2dgrid (line 81) | def get_2dgrid(S):
  function torch2dgrid (line 97) | def torch2dgrid(num_x, num_y, bot=(0,0), top=(1,1)):
  function get_grid3d (line 107) | def get_grid3d(S, T, time_scale=1.0, device='cpu'):
  function convert_ic (line 117) | def convert_ic(u0, N, S, T, time_scale=1.0):
  function requires_grad (line 126) | def requires_grad(model, flag=True):
  function set_grad (line 131) | def set_grad(tensors, flag=True):
  function zero_grad (line 136) | def zero_grad(params):
  function count_params (line 149) | def count_params(net):
  function save_checkpoint (line 156) | def save_checkpoint(path, name, model, optimizer=None):
  function save_ckpt (line 178) | def save_ckpt(path, model, optimizer=None, scheduler=None):
  function dict2str (line 197) | def dict2str(log_dict):