Miscellaneous Fixes

.gitignore

@@ -1,4 +1,27 @@
+# Python bytecode
 __pycache__/
-*.pyc
+*.py[cod]
+
+# Docs compilation
 docs/build/
-docs/source/_autosummary/
+docs/source/_autosummary/
+
+# Notebooks
+.ipynb_checkpoints
+
+# Packaging
+dist/
+*.egg-info/
+
+# Unit test
+.coverage
+.pytest_cache/
+
+# OS
+*.DS_Store
+
+# VS Code
+.vscode/
+
+# Runs
+results_*/

examples/elephant_skeleton2D_HMC.py

@@ -31,7 +31,14 @@
 Set density: 2D skeleton. This skeleton corresponds to
 the shape of a 2D elephant. The centers are read from a file.
 """
-z = np.load("examples/datasets/elephantz.npy")
+# Determine path to file
+path, dir = os.path.split(os.getcwd())
+if dir == "examples":
+    path2file = "datasets/elephantz.npy"
+else:
+    path2file = "examples/datasets/elephantz.npy"
+
+z = np.load(path2file)
 D = z.shape[0]  # Number of points in skeleton
 d = z.shape[1]  # Dimension of points in skeleton
 

examples/ionosphere_SMC.py

@@ -27,7 +27,14 @@
 """
 Read ionosphere data and pre-process.
 """
-x, y = load_data("examples/datasets/ionosphere_full.pkl")
+# Determine path to file
+path, dir = os.path.split(os.getcwd())
+if dir == "examples":
+    path2file = "datasets/ionosphere_full.pkl"
+else:
+    path2file = "examples/datasets/ionosphere_full.pkl"
+
+x, y = load_data(path2file)
 print(f"Data read shapes: x --> {x.shape}, y --> {y.shape}")
 print(f"Range x: min --> {x.min()}, max --> {x.max()}")
 print(f"Range y: min --> {y.min()}, max --> {y.max()}")

examples/simple2D_LFIS.py

@@ -26,6 +26,7 @@
 from rmc import (
     CosineSchedule,
     LiouvilleFlow,
+    PackedMultivariateNormal,
     plot_quiver,
     plot_samples,
     plot_trajectories,
@@ -55,42 +56,38 @@
 """
 Configure sampling run.
 """
-# define prior
-prior_mean = 0.0  # prior mean
-prior_std = 2.0  # prior standard deviation
-prior_mean_vec = prior_mean * jnp.ones((1, d))
-prior_std_vec = prior_std * jnp.ones((1, d))
+# define base distribution
+mean_base = jnp.zeros(d).reshape((1, d))
+cov_base = 1.5 * jnp.eye(d).reshape((1, d, d))
+distribution0 = PackedMultivariateNormal(mean_base, cov_base)
 
 """
 Construct Louville Flow (LF) Model, a Flax neural network (NN) model,
 specifically a multi-layer perceptron (MLP).
 """
 # NN configuration
-# layer_widths = [64, 64, 64] # number of neurons per layer
-layer_widths = [16, 16, 16]  # number of neurons per layer
+layer_widths = [64, 64, 64]  # number of neurons per layer
 nn_conf: NNConfigDict = {
     "seed": 10,
     "task": "train",
-    "batch_size": 200,  # 500,#20000,
-    "method": "withoutweight",  # "withweight_resample"
-    # "method": "withweight_resample",
+    "batch_size": 500,
+    "method": "withoutweight",
     "dim": d,
     "layer_widths": layer_widths,
     "activation_func": nnx.silu,
     "opt_type": "ADAM",
     "base_lr": 1e-2,
-    "max_epochs": 500,
-    "mu0_mean": prior_mean * jnp.ones((1, d)),
-    "mu0_covariance": jnp.diagflat((prior_std * jnp.ones((d,))) ** 2).reshape((1, d, d)),
-    "dt_max": 2e-1,  # 4e-3,
-    "max_samples": 1000,
-    "nsamples": 1000,  # 1000,#500,#250,
+    "max_epochs": 1000,
+    "dt_max": 2e-1,
+    "dist0": distribution0,
+    "max_samples": 500,
+    "nsamples": 2000,
     "eval_every": 100,
-    "warm_start": False,  # True,
-    "max_loss": 5e-1,
-    "max_subiter": 4,  # 2, #11, #1, #10,
+    "warm_start": False,
+    "max_loss": 5e-2,
+    "max_subiter": 5,
     "has_aux": True,
-    "root_path": "./results-s2D/",
+    "root_path": "./results_s2D/",
 }
 print(f"Flow-based sampling configured --> parameters: {nn_conf}")
 

examples/sonar_SMC.py

@@ -27,7 +27,14 @@
 """
 Read sonar data and pre-process.
 """
-x, y = load_data("examples/datasets/sonar_full.pkl")
+# Determine path to file
+path, dir = os.path.split(os.getcwd())
+if dir == "examples":
+    path2file = "datasets/sonar_full.pkl"
+else:
+    path2file = "examples/datasets/sonar_full.pkl"
+
+x, y = load_data(path2file)
 print(f"Data read shapes: x --> {x.shape}, y --> {y.shape}")
 print(f"Range x: min --> {x.min()}, max --> {x.max()}")
 print(f"Range y: min --> {y.min()}, max --> {y.max()}")

rmc/flax/trainer.py

@@ -3,6 +3,7 @@
 """Definition of functions for training Flax neural network."""
 
 import pickle
+from functools import partial
 from pathlib import Path
 from typing import Any, Callable, Optional
 

rmc/modules/lfis.py

@@ -17,6 +17,7 @@
 from rmc.flax.trainer import load_model, save_model, train
 from rmc.utils.density import LogDensityPath, LogDensityPosterior
 from rmc.utils.math_utils import divergence
+from rmc.utils.packed_distributions import PackedMultivariateNormal
 
 
 class NN_LiouvilleFlow(nnx.Module):
@@ -99,7 +100,15 @@ def __init__(
             # Use prior density to sample from initial distribution
             self.distribution0 = densitycl.prior.rvs
         else:
-            raise NotImplementedError
+            # Use provided distribution
+            # If not provided, use multivariate normal with zero mean and identity covariance
+            if "dist0" in config.keys():
+                self.distribution0 = config["dist0"].rvs
+            else:
+                d = config["dim"]
+                mean_base = jnp.zeros(d).reshape((1, d))
+                cov_base = jnp.eye(d).reshape((1, d, d))
+                self.distribution0 = PackedMultivariateNormal(mean_base, cov_base).rvs
 
         # Store schedule
         self.schedule = schedule

rmc/utils/helpers.py

@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+
+"""Miscellaneous helper functions."""
+
+from inspect import isfunction
+from typing import Any
+
+
+def exists(x: Any):
+    """Determine if x is not none."""
+    return x is not None
+
+
+def default(val: Any, d: Any):
+    """Return default value if given. Otherwise return object d.
+    Args:
+        val: Default value.
+        d: Function or variable to return if no default value provided.
+    """
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d