update module names, cleanup some docs warnings

anvil/benchmarks.py

@@ -40,8 +40,8 @@ def free_scalar_theory(couplings, lattice_length):
 def fourier_transform(configs, training_geometry):
     """Takes the Fourier transform of a sample of field configurations.
 
-    Inputs
-    ------
+    Parameters
+    ----------
     configs: torch.tensor
         A (hopefully decorrelated) sample of field configurations in the
         split representation. Shape: (sample_size, lattice_size)

anvil/distributions.py

@@ -14,8 +14,8 @@ class Gaussian:
     """
     Class which handles the generation of a sample of latent Gaussian variables.
 
-    Inputs:
-    -------
+    Parameters
+    ----------
     lattice_size: int
         Number of nodes on the lattice.
     sigma: float
@@ -59,6 +59,18 @@ class PhiFourScalar:
     The forward pass returns the corresponding log density (unnormalised) which
     is equal to -S
 
+    The parameters required differ depending on the parameterisation you're
+    using:
+
+    ================  =============
+    parameterisation  couplings
+    ================  =============
+    standard          m_sq, g
+    albergo2019       m_sq, lam
+    nicoli2020        kappa, lam
+    bosetti2015       beta, lam
+    ================  =============
+
     Parameters
     ----------
     geometry:
@@ -70,14 +82,6 @@ class PhiFourScalar:
         dictionary with two entries that are the couplings of the theory.
         See below.
 
-
-        parameterisation            couplings
-        -------------------------------------
-        standard                    m_sq, g
-        albergo2019                 m_sq, lam
-        nicoli2020                  kappa, lam
-        bosetti2015                 beta, lam
-
     Notes
     -----
     The general form of the action is

anvil/free_scalar.py

@@ -127,8 +127,8 @@ def gen_complex_normal(n_sample, sigma, real=False):
         """Returns a stack of complex arrays where real and imaginary components
         are drawn from a Gaussian distribution with the same width.
 
-        Inputs:
-        -------
+        Parameters
+        ----------
         n_sample: int
             sample size
         sigma: numpy.ndarray
@@ -137,8 +137,8 @@ def gen_complex_normal(n_sample, sigma, real=False):
             (optional) flag. If True, the imaginary component is set to
             zero, but a complex array is still returned.
 
-        Returns:
-        --------
+        Returns
+        -------
         out: numpy.ndarray
             complex array of shape (n_sample, *sigma.shape)
         """
@@ -157,13 +157,13 @@ def gen_eigenmodes(self, n_sample):
         Gaussian distributions with variances given by the eigenvalues of the
         kinetic operator - see _variance() method above.
 
-        Inputs:
-        -------
+        Parameters
+        ----------
         n_sample: int
             sample size
 
-        Returns:
-        --------
+        Returns
+        -------
         eigenmodes: numpy.ndarray
             complex array of eigenmodes with shape (n_sample, L, L)
             where L is the side length of the square lattice.
@@ -206,13 +206,13 @@ def gen_eigenmodes(self, n_sample):
     def gen_real_space_fields(self, n_sample):
         """Returns the inverse fourier transform of a sample of eigenmodes.
 
-        Inputs:
-        -------
+        Parameters
+        ----------
         n_sample: int
             sample size
 
-        Returns:
-        --------
+        Returns
+        -------
         fields: numpy.ndarray
             real array of real-space fields, with shape (n_sample, L, L),
             where L is the side-length of the square lattice.

anvil/layers.py

@@ -33,7 +33,7 @@
 
 .. math::
 
-        f(z) = g_{N_layers}( \ldots ( g_2( g_1( z ) ) ) \ldots )
+        f(z) = g_{N_{\rm layers}}( \ldots ( g_2( g_1( z ) ) ) \ldots )
 
 As a matter of convenience we provide a subclass of
 :py:class:`torch.nn.Sequential`, which is initialised by passing multiple layers
@@ -47,7 +47,7 @@
 import torch
 import torch.nn as nn
 
-from anvil.core import FullyConnectedNeuralNetwork
+from anvil.neural_network import DenseNeuralNetwork
 
 
 class CouplingLayer(nn.Module):
@@ -127,7 +127,7 @@ def __init__(
     ):
         super().__init__(size_half, even_sites)
 
-        self.t_network = FullyConnectedNeuralNetwork(
+        self.t_network = DenseNeuralNetwork(
             size_in=size_half,
             size_out=size_half,
             hidden_shape=hidden_shape,
@@ -175,14 +175,14 @@ def __init__(
     ):
         super().__init__(size_half, even_sites)
 
-        self.s_network = FullyConnectedNeuralNetwork(
+        self.s_network = DenseNeuralNetwork(
             size_in=size_half,
             size_out=size_half,
             hidden_shape=hidden_shape,
             activation=activation,
             bias=not z2_equivar,
         )
-        self.t_network = FullyConnectedNeuralNetwork(
+        self.t_network = DenseNeuralNetwork(
             size_in=size_half,
             size_out=size_half,
             hidden_shape=hidden_shape,
@@ -250,7 +250,7 @@ def __init__(
         self.size_half = size_half
         self.n_segments = n_segments
 
-        self.network = FullyConnectedNeuralNetwork(
+        self.network = DenseNeuralNetwork(
             size_in=size_half,
             size_out=size_half * (3 * n_segments - 1),
             hidden_shape=hidden_shape,
@@ -447,10 +447,11 @@ def forward(self, v_in, log_density, *unused):
 class Sequential(nn.Sequential):
     """Similar to :py:class:`torch.nn.Sequential` except conforms to our
     ``forward`` convention.
-
     """
-
     def forward(self, v, log_density, *args):
+        """overrides the base class ``forward`` method to conform to our
+        conventioned for expected inputs/outputs of ``forward`` methods.
+        """
         for module in self:
             v, log_density = module(v, log_density, *args)
         return v, log_density

anvil/models.py

@@ -18,7 +18,7 @@
 def _coupling_pair(coupling_layer, **kwargs):
     """Helper function which wraps a pair of coupling layers from
     :py:mod:`anvil.layers` in the module container
-    :py:class`layers.Sequential`. The first transformation layer acts on
+    :py:class`anvil.layers.Sequential`. The first transformation layer acts on
     the even sites and the second transformation acts on the odd sites, so one
     of these blocks ensures all sites are transformed as part of an
     active partition.
@@ -41,7 +41,7 @@ def real_nvp(
     r"""Action which returns a sequence of ``n_blocks`` pairs of
     :py:class:`anvil.layers.AffineLayer` s, followed by a single
     :py:class:`anvil.layers.GlobalRescaling` all wrapped in the module container
-    :py:class`layers.Sequential`.
+    :py:class`anvil.layers.Sequential`.
 
     The first ``n_blocks`` elements of the outer ``Sequential``
     are ``Sequential`` s containing a pair of ``AffineLayer`` s which
@@ -71,13 +71,13 @@ def real_nvp(
 
     Returns
     -------
-    real_nvp: layers.Sequential
+    real_nvp: anvil.layers.Sequential
         A sequence of affine transformations, which we refer to as a real NVP
         (Non-volume preserving) flow.
 
     See Also
     --------
-    :py:mod:`anvil.core` contains the fully connected neural network class
+    :py:mod:`anvil.neural_network` contains the fully connected neural network class
     as well as valid choices for activation functions.
 
     """
@@ -102,8 +102,8 @@ def nice(
     z2_equivar=True,
 ):
     r"""Similar to :py:func:`real_nvp`, excepts instead wraps pairs of
-    :py:class:`layers.AdditiveLayer` s followed by a single
-    :py:class:`layers.GlobalRescaling`. The pairs of ``AdditiveLayer`` s
+    :py:class:`anvil.layers.AdditiveLayer` s followed by a single
+    :py:class:`anvil.layers.GlobalRescaling`. The pairs of ``AdditiveLayer`` s
     act on the even and odd sites respectively.
 
     Parameters
@@ -127,7 +127,7 @@ def nice(
 
     Returns
     -------
-    nice: layers.Sequential
+    nice: anvil.layers.Sequential
         A sequence of additive transformations, which we refer to as a
         nice flow.
 
@@ -155,8 +155,8 @@ def rational_quadratic_spline(
     z2_equivar=False,
 ):
     """Similar to :py:func:`real_nvp`, excepts instead wraps pairs of
-    :py:class:`layers.RationalQuadraticSplineLayer` s followed by a single
-    :py:class:`layers.GlobalRescaling`. The pairs of RQS's
+    :py:class:`anvil.layers.RationalQuadraticSplineLayer` s followed by a single
+    :py:class:`anvil.layers.GlobalRescaling`. The pairs of RQS's
     act on the even and odd sites respectively.
 
     Parameters
@@ -206,17 +206,51 @@ def rational_quadratic_spline(
 
 def model_to_load(_normalising_flow):
     """action which wraps a list of layers in
-    :py:class:`layers.Sequential`. This allows the user to specify an
+    :py:class:`anvil.layers.Sequential`. This allows the user to specify an
     arbitrary combination of layers as the model.
 
-    For more information
-    on valid choices for layers, see :py:var:`LAYER_OPTIONS` or the various
+    For more information on valid choices for layers, see
+    ``anvil.models.LAYER_OPTIONS`` or the various
     functions in :py:mod:`anvil.models` which produce sequences of the layers
     found in :py:mod:`anvil.layers`.
 
+    At present, available transformations are:
+
+        - ``nice``
+        - ``real_nvp``
+        - ``rational_quadratic_spline``
+
+    You can see their dependencies using the ``anvil`` provider help, e.g.
+    for ``real_nvp``:
+
+    .. code::
+
+        $ anvil-sample --help real_nvp
+        ...
+        < action docstring - poorly formatted>
+        ...
+        The following resources are read from the configuration:
+
+            lattice_length(int):
+        [Used by lattice_size]
+
+            lattice_dimension(int): Parse lattice dimension from runcard
+        [Used by lattice_size]
+
+        The following additionl arguments can be used to control the
+        behaviour. They are set by default to sensible values:
+
+        n_blocks
+        hidden_shape
+        activation = tanh
+        z2_equivar = True
+
+    ``anvil-train`` will also provide the same information.
+
     """
     return layers.Sequential(*_normalising_flow)
 
+# Update docstring above if you add to this!
 LAYER_OPTIONS = {
     "nice": nice,
     "real_nvp": real_nvp,

anvil/core.py → anvil/neural_network.py

@@ -1,9 +1,10 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 # Copywrite © 2021 anvil Michael Wilson, Joe Marsh Rossney, Luigi Del Debbio
 """
-core.py
+neural_network.py
 
-Module containing project specific extensions to pytorch base classes.
+Module containing neural networks which are used as part of transformation
+layers, found in :py:mod:`anvil.layers`.
 
 """
 import torch
@@ -17,8 +18,8 @@
 }
 
 
-class FullyConnectedNeuralNetwork(nn.Module):
-    """Generic class for neural networks used in coupling layers.
+class DenseNeuralNetwork(nn.Module):
+    """Dense neural networks used in coupling layers.
 
     Parameters
     ----------
@@ -30,11 +31,11 @@ class FullyConnectedNeuralNetwork(nn.Module):
         List specifying the number of nodes in the intermediate layers
     activation: (str, None)
         Key representing the activation function used for each layer
-        except the final one.
-    no_final_activation: bool
-        If True, leave the network output unconstrained.
-    bias: bool
+        except the final one. Valid options can be found in
+        ``ACTIVATION_LAYERS``.
+    bias: bool, default=True
         Whether to use biases in networks.
+
     """
 
     def __init__(

docs/sphinx/source/get-started/basic-usage.rst

@@ -118,7 +118,7 @@ We supply some basic
 machinery to build your own normalising flow models. The relevant modules for
 this purpose are
 
- - :py:mod:`anvil.core`: containing some basic extensions to ``pytorch`` modules relevant for this project.
+ - :py:mod:`anvil.neural_network`: Generic neural networks.
  - :py:mod:`anvil.layers`: a collection of transformation layer classes
  - :py:mod:`anvil.geometry`: classes which transform the output of the transformations into meaningful geometries. These dictate which sites in your lattice get alternated between active and passive partitions.
  - :py:mod:`anvil.distributions`: a collection of distributions which can be used as base distributions (for latent variables) or target distributions.