diff --git a/docs/api/arch/conv.md b/docs/api/arch/conv.md new file mode 100644 index 0000000..e541b30 --- /dev/null +++ b/docs/api/arch/conv.md @@ -0,0 +1,7 @@ +# Feed-Forward Conv Net + +::: pdequinox.arch.ConvNet + options: + members: + - __init__ + - __call__ \ No newline at end of file diff --git a/docs/api/arch/dilated.md b/docs/api/arch/dilated.md new file mode 100644 index 0000000..a1e1120 --- /dev/null +++ b/docs/api/arch/dilated.md @@ -0,0 +1,9 @@ +# Dilated ResNet + +A sequential convolutional architecture employing dilated convolutions for long-range interactions. + +::: pdequinox.arch.DilatedResNet + options: + members: + - __init__ + - __call__ \ No newline at end of file diff --git a/docs/api/arch/fno.md b/docs/api/arch/fno.md new file mode 100644 index 0000000..d4849d8 --- /dev/null +++ b/docs/api/arch/fno.md @@ -0,0 +1,9 @@ +# Fourier Neural Operator (FNO) + +A Fourier pseudo-spectral architecture. + +::: pdequinox.arch.ClassicFNO + options: + members: + - __init__ + - __call__ \ No newline at end of file diff --git a/docs/api/arch/mlp.md b/docs/api/arch/mlp.md new file mode 100644 index 0000000..2814f7c --- /dev/null +++ b/docs/api/arch/mlp.md @@ -0,0 +1,7 @@ +# Multi-Layer Perceptron (MLP) + +::: pdequinox.arch.MLP + options: + members: + - __init__ + - __call__ \ No newline at end of file diff --git a/docs/api/arch/resnet.md b/docs/api/arch/resnet.md new file mode 100644 index 0000000..bf7f904 --- /dev/null +++ b/docs/api/arch/resnet.md @@ -0,0 +1,15 @@ +# (Convolutional) ResNet + +::: pdequinox.arch.ClassicResNet + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.arch.ModernResNet + options: + members: + - __init__ + - __call__ \ No newline at end of file diff --git a/docs/api/arch/unet.md b/docs/api/arch/unet.md new file mode 100644 index 0000000..4b5dbd3 --- /dev/null +++ b/docs/api/arch/unet.md @@ -0,0 +1,17 @@ +# UNet + +A hierarchical multi-scale convolutional network. + +::: pdequinox.arch.ClassicUNet + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.arch.ModernUNet + options: + members: + - __init__ + - __call__ \ No newline at end of file diff --git a/docs/api/arch_constructors.md b/docs/api/arch_constructors.md new file mode 100644 index 0000000..844e4d0 --- /dev/null +++ b/docs/api/arch_constructors.md @@ -0,0 +1,15 @@ +# Architectural Constructors + +::: pdequinox.Sequential + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.Hierarchical + options: + members: + - __init__ + - __call__ diff --git a/docs/api/block_factories.md b/docs/api/block_factories.md new file mode 100644 index 0000000..0910d7c --- /dev/null +++ b/docs/api/block_factories.md @@ -0,0 +1,71 @@ +# Block Factories + +::: pdequinox.blocks.ClassicDoubleConvBlockFactory + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.ClassicResBlockFactory + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.ModernResBlockFactory + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.ClassicSpectralBlockFactory + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.DilatedResBlockFactory + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.LinearConvBlockFactory + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.LinearConvDownBlockFactory + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.LinearConvUpBlockFactory + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.BlockFactory + options: + members: + - __init__ + - __call__ \ No newline at end of file diff --git a/docs/api/blocks.md b/docs/api/blocks.md new file mode 100644 index 0000000..1719d20 --- /dev/null +++ b/docs/api/blocks.md @@ -0,0 +1,71 @@ +# Blocks + +::: pdequinox.blocks.ClassicDoubleConvBlock + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.ClassicResBlock + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.ModernResBlock + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.ClassicSpectralBlock + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.DilatedResBlock + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.LinearConvBlock + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.LinearConvDownBlock + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.LinearConvUpBlock + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.blocks.Block + options: + members: + - __init__ + - __call__ \ No newline at end of file diff --git a/docs/api/conv.md b/docs/api/conv.md new file mode 100644 index 0000000..0cea01c --- /dev/null +++ b/docs/api/conv.md @@ -0,0 +1,34 @@ +# Low-Level Convolution Routines + +They wrap the `equinox.conv` module to provide an interface based on `"SAME"` +convolutions with variable boundary modes or implement spectral convolutions. + +::: pdequinox.conv.PhysicsConv + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.conv.PhysicsConvTranspose + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.conv.SpectralConv + options: + members: + - __init__ + - __call__ + +--- + +::: pdequinox.conv.PointwiseLinearConv + options: + members: + - __init__ + - __call__ \ No newline at end of file diff --git a/docs/api/sample_data.md b/docs/api/sample_data.md new file mode 100644 index 0000000..d3f962c --- /dev/null +++ b/docs/api/sample_data.md @@ -0,0 +1,3 @@ +# Sample Data + +::: pdequinox.sample_data \ No newline at end of file diff --git a/docs/api/utils.md b/docs/api/utils.md new file mode 100644 index 0000000..cb52a74 --- /dev/null +++ b/docs/api/utils.md @@ -0,0 +1,15 @@ +# Utilities + +::: pdequinox.dataloader + +--- + +::: pdequinox.cycling_dataloader + +--- + +::: pdequinox.extract_from_ensemble + +--- + +::: pdequinox.combine_to_ensemble \ No newline at end of file diff --git a/docs/index.md b/docs/index.md index 401e4f5..be0fb87 100644 --- a/docs/index.md +++ b/docs/index.md @@ -64,32 +64,6 @@ for epoch in tqdm(range(100)): ) loss_history.append(loss) ``` -## Background - -Neural Emulators are networks learned to efficienty predict physical phenomena, -often associated with PDEs. In the simplest case this can be a linear advection -equation, all the way to more complicated Navier-Stokes cases. If we work on -Uniform Cartesian grids* (which this package assumes), one can borrow plenty of -architectures from image-to-image tasks in computer vision (e.g., for -segmentation). This includes: - -* Standard Feedforward ConvNets -* Convolutional ResNets ([He et al.](https://arxiv.org/abs/1512.03385)) -* U-Nets ([Ronneberger et al.](https://arxiv.org/abs/1505.04597)) -* Dilated ResNets ([Yu et al.](https://arxiv.org/abs/1511.07122), [Stachenfeld et al.](https://arxiv.org/abs/2112.15275)) -* Fourier Neural Operators ([Li et al.](https://arxiv.org/abs/2010.08895)) - -It is interesting to note that most of these architectures resemble classical -numerical methods or at least share similarities with them. For example, -ConvNets (or convolutions in general) are related to finite differences, while -U-Nets resemble multigrid methods. Fourier Neural Operators are related to -spectral methods. The difference is that the emulators' free parameters are -found based on a (data-driven) numerical optimization not a symbolic -manipulation of the differential equations. - -(*) This means that we essentially have a pixel or voxel grid on which space is -discretized. Hence, the space can only be the scaled unit cube $\Omega = (0, -L)^D$ ## Features @@ -106,76 +80,6 @@ L)^D$ * Composability * Tools to count parameters and assess receptive fields -## Boundary Conditions - -This package assumes that the boundary condition is baked into the neural -emulator. Hence, most components allow setting `boundary_mode` which can be -`"dirichlet"`, `"neumann"`, or `"periodic"`. This affects what is considered a -degree of freedom in the grid. - -![](imgs/three_boundary_conditions.svg) - -Dirichlet boundaries fully eliminate degrees of freedom on the boundary. -Periodic boundaries only keep one end of the domain as a degree of freedom (This -package follows the convention that the left boundary is the degree of freedom). Neumann boundaries keep both ends as degrees of freedom. - -## Constructors - -There are two primary architectural constructors for Sequential and Hierarchical -Networks that allow for composability with the `PDEquinox` blocks. - -### Sequential Constructor - -![](imgs/sequential_net.svg) - -The squential network constructor is defined by: -* a lifting block $\mathcal{L}$ -* $N$ blocks $\left \{ \mathcal{B}_i \right\}_{i=1}^N$ -* a projection block $\mathcal{P}$ -* the hidden channels within the sequential processing -* the number of blocks $N$ (one can also supply a list of hidden channels if they shall be different between blocks) - -### Hierarchical Constructor - -![](imgs/hierarchical_net.svg) - -The hierarchical network constructor is defined by: -* a lifting block $\mathcal{L}$ -* The number of levels $D$ (i.e., the number of additional hierarchies). Setting $D = 0$ recovers the sequential processing. -* a list of $D$ blocks $\left \{ \mathcal{D}_i \right\}_{i=1}^D$ for - downsampling, i.e. mapping downwards to the lower hierarchy (oftentimes this - is that they halve the spatial axes while keeping the number of channels) -* a list of $D$ blocks $\left \{ \mathcal{B}_i^l \right\}_{i=1}^D$ for - processing in the left arc (oftentimes this changes the number of channels, - e.g. doubles it such that the combination of downsampling and left processing - halves the spatial resolution and doubles the feature count) -* a list of $D$ blocks $\left \{ \mathcal{U}_i \right\}_{i=1}^D$ for upsamping, - i.e., mapping upwards to the higher hierarchy (oftentimes this doubles the - spatial resolution; at the same time it halves the feature count such that we - can concatenate a skip connection) -* a list of $D$ blocks $\left \{ \mathcal{B}_i^r \right\}_{i=1}^D$ for - processing in the right arc (oftentimes this changes the number of channels, - e.g. halves it such that the combination of upsampling and right processing - doubles the spatial resolution and halves the feature count) -* a projection block $\mathcal{P}$ -* the hidden channels within the hierarchical processing (if just an integer is - provided; this is assumed to be the number of hidden channels in the highest - hierarchy.) - -### Beyond Architectural Constructors - -For completion, `pdequinox.arch` also provides a `ConvNet` which is a simple -feed-forward convolutional network. It also provides `MLP` which is a dense -networks which also requires pre-defining the number of resolution points. - -## Related - -Similar packages that provide a collection of emulator architectures are -[PDEBench](https://github.com/pdebench/PDEBench) and -[PDEArena](https://github.com/pdearena/pdearena). With focus on Phyiscs-informed -Neural Networks and Neural Operators, there are also -[DeepXDE](https://github.com/lululxvi/deepxde) and [NVIDIA -Modulus](https://developer.nvidia.com/modulus). ## License diff --git a/docs/more_details.md b/docs/more_details.md new file mode 100644 index 0000000..bd2547d --- /dev/null +++ b/docs/more_details.md @@ -0,0 +1,99 @@ +# More Details + +## Background + +Neural Emulators are networks learned to efficienty predict physical phenomena, +often associated with PDEs. In the simplest case this can be a linear advection +equation, all the way to more complicated Navier-Stokes cases. If we work on +Uniform Cartesian grids* (which this package assumes), one can borrow plenty of +architectures from image-to-image tasks in computer vision (e.g., for +segmentation). This includes: + +* Standard Feedforward ConvNets +* Convolutional ResNets ([He et al.](https://arxiv.org/abs/1512.03385)) +* U-Nets ([Ronneberger et al.](https://arxiv.org/abs/1505.04597)) +* Dilated ResNets ([Yu et al.](https://arxiv.org/abs/1511.07122), [Stachenfeld et al.](https://arxiv.org/abs/2112.15275)) +* Fourier Neural Operators ([Li et al.](https://arxiv.org/abs/2010.08895)) + +It is interesting to note that most of these architectures resemble classical +numerical methods or at least share similarities with them. For example, +ConvNets (or convolutions in general) are related to finite differences, while +U-Nets resemble multigrid methods. Fourier Neural Operators are related to +spectral methods. The difference is that the emulators' free parameters are +found based on a (data-driven) numerical optimization not a symbolic +manipulation of the differential equations. + +(*) This means that we essentially have a pixel or voxel grid on which space is +discretized. Hence, the space can only be the scaled unit cube $\Omega = (0, +L)^D$ + +## Boundary Conditions + +This package assumes that the boundary condition is baked into the neural +emulator. Hence, most components allow setting `boundary_mode` which can be +`"dirichlet"`, `"neumann"`, or `"periodic"`. This affects what is considered a +degree of freedom in the grid. + +![](imgs/three_boundary_conditions.svg) + +Dirichlet boundaries fully eliminate degrees of freedom on the boundary. +Periodic boundaries only keep one end of the domain as a degree of freedom (This +package follows the convention that the left boundary is the degree of freedom). Neumann boundaries keep both ends as degrees of freedom. + +## Constructors + +There are two primary architectural constructors for Sequential and Hierarchical +Networks that allow for composability with the `PDEquinox` blocks. + +### Sequential Constructor + +![](imgs/sequential_net.svg) + +The squential network constructor is defined by: +* a lifting block $\mathcal{L}$ +* $N$ blocks $\left \{ \mathcal{B}_i \right\}_{i=1}^N$ +* a projection block $\mathcal{P}$ +* the hidden channels within the sequential processing +* the number of blocks $N$ (one can also supply a list of hidden channels if they shall be different between blocks) + +### Hierarchical Constructor + +![](imgs/hierarchical_net.svg) + +The hierarchical network constructor is defined by: +* a lifting block $\mathcal{L}$ +* The number of levels $D$ (i.e., the number of additional hierarchies). Setting $D = 0$ recovers the sequential processing. +* a list of $D$ blocks $\left \{ \mathcal{D}_i \right\}_{i=1}^D$ for + downsampling, i.e. mapping downwards to the lower hierarchy (oftentimes this + is that they halve the spatial axes while keeping the number of channels) +* a list of $D$ blocks $\left \{ \mathcal{B}_i^l \right\}_{i=1}^D$ for + processing in the left arc (oftentimes this changes the number of channels, + e.g. doubles it such that the combination of downsampling and left processing + halves the spatial resolution and doubles the feature count) +* a list of $D$ blocks $\left \{ \mathcal{U}_i \right\}_{i=1}^D$ for upsamping, + i.e., mapping upwards to the higher hierarchy (oftentimes this doubles the + spatial resolution; at the same time it halves the feature count such that we + can concatenate a skip connection) +* a list of $D$ blocks $\left \{ \mathcal{B}_i^r \right\}_{i=1}^D$ for + processing in the right arc (oftentimes this changes the number of channels, + e.g. halves it such that the combination of upsampling and right processing + doubles the spatial resolution and halves the feature count) +* a projection block $\mathcal{P}$ +* the hidden channels within the hierarchical processing (if just an integer is + provided; this is assumed to be the number of hidden channels in the highest + hierarchy.) + +### Beyond Architectural Constructors + +For completion, `pdequinox.arch` also provides a `ConvNet` which is a simple +feed-forward convolutional network. It also provides `MLP` which is a dense +networks which also requires pre-defining the number of resolution points. + +## Related + +Similar packages that provide a collection of emulator architectures are +[PDEBench](https://github.com/pdebench/PDEBench) and +[PDEArena](https://github.com/pdearena/pdearena). With focus on Phyiscs-informed +Neural Networks and Neural Operators, there are also +[DeepXDE](https://github.com/lululxvi/deepxde) and [NVIDIA +Modulus](https://developer.nvidia.com/modulus). \ No newline at end of file diff --git a/mkdocs.yml b/mkdocs.yml index 795c7d8..aef94e9 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -81,4 +81,27 @@ plugins: docstring_style: null nav: - - 'index.md' \ No newline at end of file + - 'index.md' + - 'more_details.md' + - Examples: + - Introductory: + - Train UNet as Poisson Emulator: 'examples/train_unet_as_poisson_solver.ipynb' + - Parameter Counting & Receptive Field: 'examples/parameter_count_and_receptive_field.ipynb' + - Architecture Showcase: 'examples/architecture_showcase.ipynb' + - Advanced: + - Seed Parallel Training: 'examples/seed_parallel_training.ipynb' + - High-Level API: + - Architectures: + - MLP: 'api/arch/mlp.md' + - Conv: 'api/arch/conv.md' + - ResNet: 'api/arch/resnet.md' + - UNet: 'api/arch/unet.md' + - Dilated: 'api/arch/dilated.md' + - FNO: 'api/arch/fno.md' + - Utilities: 'api/utils.md' + - Sample Data: 'api/sample_data.md' + - Low-Level API: + - Arch Constructors: 'api/arch_constructors.md' + - Blocks: 'api/blocks.md' + - Block Factories: 'api/block_factories.md' + - Conv: 'api/conv.md' \ No newline at end of file diff --git a/pdequinox/conv/_physics_conv.py b/pdequinox/conv/_physics_conv.py index 08f31b6..c7617fe 100644 --- a/pdequinox/conv/_physics_conv.py +++ b/pdequinox/conv/_physics_conv.py @@ -189,11 +189,12 @@ def __init__( the input to a convolution. Moreover, to do so with the same "connectivity", i.e. which inputs can affect which outputs. - Relative to an [`PhysicsConv`][] layer, this can be accomplished by - switching the values of `in_channels` and `out_channels`, whilst - keeping `kernel_size`, `stride`, `dilation`, and `groups` the same. + Relative to an [`pdequinox.conv.PhysicsConv`][] layer, this can be + accomplished by switching the values of `in_channels` and + `out_channels`, whilst keeping `kernel_size`, `stride`, `dilation`, + and `groups` the same. - When `stride > 1` then [`PhysicsConv`][] maps multiple input shapes + When `stride > 1` then [`pdequinox.conv.PhysicsConv`][] maps multiple input shapes to the same output shape. `output_padding` is provided to resolve this ambiguity, by adding a little extra padding to just the bottom/right edges of the input.