basic reportengine loop over layers

anvil/benchmark_config/free_scalar_sample.yml

@@ -1,5 +1,6 @@
 training_output: /tmp/del_me_anvil_benchmark
 cp_id: -1
+layer_id: -1
 
 sample_size: 100000
 thermalization: 10
@@ -16,9 +17,6 @@ template_text: |
  # Eigenvalues of kinetic operator
  {@plot_kinetic_eigenvalues@}
  {@table_kinetic_eigenvalues@}
- # Layer-wise breakdown
- {@plot_layerwise_histograms@}
- {@plot_layerwise_weights@}
 
 actions_:
   - report(main=True)

anvil/config.py

@@ -13,6 +13,7 @@
 from anvil.checkpoint import TrainingOutput
 from anvil.models import MODEL_OPTIONS
 from anvil.distributions import BASE_OPTIONS, TARGET_OPTIONS
+import anvil.sample as sample
 
 from random import randint
 from sys import maxsize
@@ -121,6 +122,17 @@ def produce_checkpoint(self, cp_id=None, training_output=None):
         # get index from training_output class
         return training_output.checkpoints[training_output.cp_ids.index(cp_id)]
 
+    @element_of("layer_ids")
+    def parse_layer_id(self, layer_id: int = -1):
+        return layer_id
+
+    @explicit_node
+    def produce_configs(self, layer_id):
+        if layer_id == -1:
+            return sample.configs_from_metropolis
+        else:
+            return sample.configs_from_model
+
     def produce_training_context(self, training_output):
         """Given a training output produce the context of that training"""
         # NOTE: This seems a bit hacky, exposing the entire training configuration
@@ -183,12 +195,11 @@ def parse_bootstrap_sample_size(self, n_boot: int):
         log.warning(f"Using user specified bootstrap sample size: {n_boot}")
         return n_boot
 
-    def produce_bootstrap_seed(
-        self, manual_bootstrap_seed: (int, type(None)) = None):
+    def produce_bootstrap_seed(self, manual_bootstrap_seed: (int, type(None)) = None):
         if manual_bootstrap_seed is None:
             return randint(0, maxsize)
         # numpy is actually this strict but let's keep it sensible.
-        if (manual_bootstrap_seed < 0) or (manual_bootstrap_seed > 2**32):
+        if (manual_bootstrap_seed < 0) or (manual_bootstrap_seed > 2 ** 32):
             raise ConfigError("Seed is outside of appropriate range: [0, 2 ** 32]")
         return manual_bootstrap_seed
 
@@ -213,4 +224,4 @@ def produce_use_multiprocessing(self):
         """Don't use mp on MacOS"""
         if platform.system() == "Darwin":
             return False
-        return True
+        return True

anvil/observables.py

@@ -12,6 +12,7 @@
 
 log = logging.getLogger(__name__)
 
+
 def cosh_shift(x, xi, A, c):
     return A * np.cosh(-x / xi) + c
 
@@ -42,8 +43,11 @@ def fit_zero_momentum_correlator(zero_momentum_correlator, training_geometry):
 
 
 def correlation_length_from_fit(fit_zero_momentum_correlator):
-    popt, pcov, _ = fit_zero_momentum_correlator
-    return popt[0], np.sqrt(pcov[0, 0])
+    if fit_zero_momentum_correlator is not None:
+        popt, pcov, _ = fit_zero_momentum_correlator
+        return popt[0], np.sqrt(pcov[0, 0])
+    else:
+        return None, None
 
 
 def autocorrelation(chain):

anvil/plot.py

@@ -32,19 +32,28 @@ def plot_layerwise_weights(plot_layer_weights):
     yield from plot_layer_weights
 
 
-def plot_layer_histogram(layerwise_configs):
-    for v in layerwise_configs:
-        v = v.numpy()
-        v_pos = v[v.sum(axis=1) > 0].flatten()
-        v_neg = v[v.sum(axis=1) < 0].flatten()
-        fig, ax = plt.subplots()
-        ax.hist([v_pos, v_neg], bins=50, density=True, histtype="step")
-        yield fig
+@figure
+def plot_layer_histogram(configs):
+    v = configs.numpy()
+    v_pos = v[v.sum(axis=1) > 0].flatten()
+    v_neg = v[v.sum(axis=1) < 0].flatten()
+    fig, ax = plt.subplots()
+    ax.hist([v_pos, v_neg], bins=50, density=True, histtype="step")
+    return fig
 
 
-@figuregen
-def plot_layerwise_histograms(plot_layer_histogram):
-    yield from plot_layer_histogram
+@figure
+def plot_correlation_length(table_correlation_length):
+    fig, ax = plt.subplots()
+    ax.errorbar(
+        x=table_correlation_length.index,
+        y=table_correlation_length.value,
+        yerr=table_correlation_length.error,
+        linestyle="",
+        marker="o",
+    )
+    ax.set_xticklabels(table_correlation_length.index, rotation=45)
+    return fig
 
 
 @figure
@@ -63,14 +72,14 @@ def plot_zero_momentum_correlator(
 
     if fit_zero_momentum_correlator is not None:
         popt, pcov, t0 = fit_zero_momentum_correlator
-        shift = popt[2]
+        xi, A, shift = popt
 
         t = np.linspace(t0, T - t0, 100)
         ax.plot(
             t,
-            cosh_shift(t - T // 2, *popt) - popt[2],
+            cosh_shift(t - T // 2, *popt) - shift,
             "r--",
-            label=r"fit $A \cosh(-(t - T/2) / \xi) + c$",
+            label=r"fit $A \cosh(-(t - T/2) / \xi) + c$" + "\n" + fr"$\xi = ${xi:.2f}",
         )
     ax.errorbar(
         x=np.arange(T),

anvil/sample.py

@@ -172,7 +172,7 @@ def metropolis_hastings(
 _metropolis_hastings = collect("metropolis_hastings", ("training_context",))
 
 
-def configs(_metropolis_hastings):
+def configs_from_metropolis(_metropolis_hastings):
     return _metropolis_hastings[0][0]
 
 
@@ -186,24 +186,26 @@ def acceptance(_metropolis_hastings):
 
 # TODO: figure out how to name each coupling block
 @torch.no_grad()
-def yield_configs_layerwise(loaded_model, base_dist, metropolis_hastings):
-    v, _ = base_dist(BATCH_SIZE)
-    yield v
+def yield_configs_layerwise(loaded_model, base_dist, sample_size, layer_id):
+    v, _ = base_dist(sample_size)
+    if layer_id == 0:
+        return v
 
     negative_mag = (v.sum(dim=1).sign() < 0).nonzero().squeeze()
 
+    i = 1
     for block in loaded_model:
         v, _ = block(v, 0, negative_mag)
         # only want coupling layers
         if len([tensor for tensor in block.state_dict().values()]) > 1:
-            yield v
-
-    v = metropolis_hastings[0]
-    yield v
+            if i == layer_id:
+                return v
+            else:
+                i += 1
 
 
-_layerwise_configs = collect("yield_configs_layerwise", ("training_context",))
+_configs_from_model = collect("yield_configs_layerwise", ("training_context",))
 
 
-def layerwise_configs(_layerwise_configs):
-    return _layerwise_configs[0]
+def configs_from_model(_configs_from_model):
+    return _configs_from_model[0]

anvil/table.py

@@ -42,6 +42,9 @@ def table_fit(fit_zero_momentum_correlator, training_geometry):
             index=["xi_fit", "m_fit"],
         )
         return df
+    else:
+        # TODO should fail better than this
+        return pd.DataFrame([])
 
 
 @table
@@ -100,15 +103,15 @@ def table_correlation_length(
 
     df = pd.DataFrame(
         res,
-        columns=["Mean", "Standard deviation"],
+        columns=["value", "error"],
         index=[
             "Estimate from fit",
             "Estimate using arcosh",
             "Second moment estimate",
             "Low momentum estimate",
         ],
     )
-    df["No. correlation lengths"] = training_geometry.length / df["Mean"]
+    df["No. correlation lengths"] = training_geometry.length / df["value"]
     return df
 
 

examples/runcards/layerwise.md

@@ -0,0 +1,30 @@
+Layer-wise breakdown
+====================
+
+Model weights
+-------------
+{@plot_layerwise_weights@}
+
+Field variables
+---------------
+{@with layer_ids@}
+{@plot_layer_histogram@}
+{@endwith@}
+
+Correlation function
+--------------------
+{@with layer_ids@}
+{@plot_two_point_correlator@}
+{@endwith@}
+
+Correlation function at zero momentum
+-------------------------------------
+{@with layer_ids@}
+{@plot_zero_momentum_correlator@}
+{@endwith@}
+
+Correlation length
+------------------
+{@with layer_ids@}
+{@plot_correlation_length@}
+{@endwith@}

examples/runcards/layerwise.yml

@@ -0,0 +1,19 @@
+training_output: train
+cp_id: -1
+layer_ids: [1, 2, 3, -1]
+
+sample_size: 10000
+thermalization: 1000
+sample_interval: 1
+
+bootstrap_sample_size: 100
+
+meta:
+  author: Author
+  title: Training Report
+  keywords: [example]
+
+template: layerwise.md
+
+actions_:
+    - report(main=True)

examples/runcards/report.md

@@ -10,12 +10,10 @@
 {@plot_zero_momentum_correlator@}
 ## Correlation length
 {@plot_effective_pole_mass@}
+{@plot_correlation_length@}
 {@table_correlation_length@}
 ## Magnetisation
 {@table_magnetization@}
 {@plot_magnetization_series@}
 {@plot_magnetization_autocorr@}
 {@plot_magnetization_integrated_autocorr@}
-## Layer-wise breakdown
-{@plot_layerwise_histograms@}
-{@plot_layerwise_weights@}

examples/runcards/report.yml

@@ -1,11 +1,12 @@
 training_output: train
 cp_id: -1
+layer_id: -1
 
 sample_size: 10000
 thermalization: 10
 sample_interval: 1
 
-bootstrap_sample_size: 100
+bootstrap_sample_size: 1000
 
 meta:
   author: Author

examples/runcards/sample.yml

@@ -1,5 +1,6 @@
 training_output: train
 cp_id: -1
+layer_id: -1
 
 sample_size: 10000
 thermalization: 1000