diff --git a/.gitignore b/.gitignore index 757c01c..258ff9b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,27 @@ +# Python bytecode __pycache__/ -*.pyc +*.py[cod] + +# Docs compilation docs/build/ -docs/source/_autosummary/ \ No newline at end of file +docs/source/_autosummary/ + +# Notebooks +.ipynb_checkpoints + +# Packaging +dist/ +*.egg-info/ + +# Unit test +.coverage +.pytest_cache/ + +# OS +*.DS_Store + +# VS Code +.vscode/ + +# Runs +results_*/ diff --git a/examples/elephant_skeleton2D_HMC.py b/examples/elephant_skeleton2D_HMC.py index cdb09f1..a5950e0 100644 --- a/examples/elephant_skeleton2D_HMC.py +++ b/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 diff --git a/examples/ionosphere_SMC.py b/examples/ionosphere_SMC.py index 7216b3d..a524a2b 100644 --- a/examples/ionosphere_SMC.py +++ b/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()}") diff --git a/examples/simple2D_LFIS.py b/examples/simple2D_LFIS.py index f516b17..474e51a 100644 --- a/examples/simple2D_LFIS.py +++ b/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}") diff --git a/examples/sonar_SMC.py b/examples/sonar_SMC.py index 4179deb..7bb3d14 100644 --- a/examples/sonar_SMC.py +++ b/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()}") diff --git a/rmc/flax/trainer.py b/rmc/flax/trainer.py index c901ea1..ff571eb 100644 --- a/rmc/flax/trainer.py +++ b/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 diff --git a/rmc/modules/lfis.py b/rmc/modules/lfis.py index 2ec3a25..a98a0cf 100644 --- a/rmc/modules/lfis.py +++ b/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 diff --git a/rmc/utils/helpers.py b/rmc/utils/helpers.py new file mode 100644 index 0000000..09237a7 --- /dev/null +++ b/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