Explicit model specification

anvil/benchmark_config/free_scalar_train.yml

@@ -12,15 +12,15 @@ couplings:
 # Model
 base: gaussian
 
-#model: rational_quadratic_spline
-#model: real_nvp
-model: nice
-model_params:
-    n_affine: 2
-    n_additive: 2
+model:
+  - layer: nice
+    n_blocks: 2
     hidden_shape: [36]
     activation: tanh
     z2_equivar: True
+  - layer: global_rescaling
+    scale: 1
+    learnable: True
 
 # Training length
 n_batch: 2000

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/config.py

@@ -13,7 +13,7 @@
 
 from anvil.geometry import Geometry2D
 from anvil.checkpoint import TrainingOutput
-from anvil.models import MODEL_OPTIONS
+from anvil.models import LAYER_OPTIONS
 from anvil.distributions import BASE_OPTIONS, TARGET_OPTIONS
 
 from random import randint
@@ -84,12 +84,12 @@ def parse_parameterisation(self, param: str):
         return param
 
     @explicit_node
-    def produce_model_action(self, model: str):
+    def produce_layer_action(self, layer: str):
         """Given a string, return the flow model action indexed by that string."""
         try:
-            return MODEL_OPTIONS[model]
+            return LAYER_OPTIONS[layer]
         except KeyError:
-            raise ConfigError(f"Invalid model {model}", model, MODEL_OPTIONS.keys())
+            raise ConfigError(f"Invalid model {layer}", layer, LAYER_OPTIONS.keys())
 
     def parse_n_batch(self, nb: int):
         """Batch size for training."""

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.