From e6343d768af081e8d73d0ddab54fbd71d612a4ea Mon Sep 17 00:00:00 2001 From: Cristina Garcia Cardona Date: Fri, 20 Feb 2026 14:34:28 -0700 Subject: [PATCH 1/7] Add more patterns to .gitignore Committer: Cristina Garcia Cardona --- .gitignore | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) 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_*/ From c307485acea4f794c00245ba1b37582b673b0d4b Mon Sep 17 00:00:00 2001 From: crstngc Date: Fri, 20 Feb 2026 15:04:49 -0700 Subject: [PATCH 2/7] Add default base distribution in LFIS and miscellaneous helpers --- rmc/modules/lfis.py | 9 ++++++++- rmc/utils/helpers.py | 21 +++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) create mode 100644 rmc/utils/helpers.py diff --git a/rmc/modules/lfis.py b/rmc/modules/lfis.py index 2ec3a25..cfd09ae 100644 --- a/rmc/modules/lfis.py +++ b/rmc/modules/lfis.py @@ -16,7 +16,9 @@ from rmc.flax.nn_config_dict import NNConfigDict from rmc.flax.trainer import load_model, save_model, train from rmc.utils.density import LogDensityPath, LogDensityPosterior +from rmc.utils.helpers import default from rmc.utils.math_utils import divergence +from rmc.utils.packed_distributions import PackedMultivariateNormal class NN_LiouvilleFlow(nnx.Module): @@ -99,7 +101,12 @@ def __init__( # Use prior density to sample from initial distribution self.distribution0 = densitycl.prior.rvs else: - raise NotImplementedError + # Use provided multivariate normal + # If not provided use multivariate normal with zero mean and identity covariance + d = config["dim"] + mean_base = default(config["mean_base"], jnp.zeros(d).reshape((1, d))) + cov_base = default(config["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..9c57274 --- /dev/null +++ b/rmc/utils/helpers.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- + +"""Miscellaneous helper functions.""" + +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 From 233cb49bc5fe09823894c3251396618d19cf515d Mon Sep 17 00:00:00 2001 From: crstngc Date: Fri, 20 Feb 2026 16:49:45 -0700 Subject: [PATCH 3/7] Fix default distribution0 in LFIS --- examples/simple2D_LFIS.py | 4 +--- rmc/modules/lfis.py | 16 +++++++++------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/examples/simple2D_LFIS.py b/examples/simple2D_LFIS.py index f516b17..bca568e 100644 --- a/examples/simple2D_LFIS.py +++ b/examples/simple2D_LFIS.py @@ -80,8 +80,6 @@ "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, @@ -90,7 +88,7 @@ "max_loss": 5e-1, "max_subiter": 4, # 2, #11, #1, #10, "has_aux": True, - "root_path": "./results-s2D/", + "root_path": "./results_s2D/", } print(f"Flow-based sampling configured --> parameters: {nn_conf}") diff --git a/rmc/modules/lfis.py b/rmc/modules/lfis.py index cfd09ae..a98a0cf 100644 --- a/rmc/modules/lfis.py +++ b/rmc/modules/lfis.py @@ -16,7 +16,6 @@ from rmc.flax.nn_config_dict import NNConfigDict from rmc.flax.trainer import load_model, save_model, train from rmc.utils.density import LogDensityPath, LogDensityPosterior -from rmc.utils.helpers import default from rmc.utils.math_utils import divergence from rmc.utils.packed_distributions import PackedMultivariateNormal @@ -101,12 +100,15 @@ def __init__( # Use prior density to sample from initial distribution self.distribution0 = densitycl.prior.rvs else: - # Use provided multivariate normal - # If not provided use multivariate normal with zero mean and identity covariance - d = config["dim"] - mean_base = default(config["mean_base"], jnp.zeros(d).reshape((1, d))) - cov_base = default(config["cov_base"], jnp.eye(d).reshape((1, d, d))) - self.distribution0 = PackedMultivariateNormal(mean_base, cov_base).rvs + # 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 From cf8f86a0718827a2d369ebfb16f23f1567fac30f Mon Sep 17 00:00:00 2001 From: crstngc Date: Fri, 20 Feb 2026 17:29:33 -0700 Subject: [PATCH 4/7] Clean clutter in simple2D LFIS example --- examples/simple2D_LFIS.py | 33 ++++++++++++++++----------------- 1 file changed, 16 insertions(+), 17 deletions(-) diff --git a/examples/simple2D_LFIS.py b/examples/simple2D_LFIS.py index bca568e..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,38 +56,36 @@ """ 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, - "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/", } From 60ecccc5c1f73a3cc1d7d195c2d88c5f9fdd2ae1 Mon Sep 17 00:00:00 2001 From: crstngc Date: Fri, 20 Feb 2026 17:41:06 -0700 Subject: [PATCH 5/7] Find path for loading data in examples --- examples/elephant_skeleton2D_HMC.py | 9 ++++++++- examples/ionosphere_SMC.py | 9 ++++++++- examples/sonar_SMC.py | 9 ++++++++- 3 files changed, 24 insertions(+), 3 deletions(-) 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/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()}") From 4b10f254de310e1e4bbde44ab5a5f5da3fb491b0 Mon Sep 17 00:00:00 2001 From: crstngc Date: Fri, 20 Feb 2026 18:28:48 -0700 Subject: [PATCH 6/7] Fix import from functools in trainer --- rmc/flax/trainer.py | 1 + 1 file changed, 1 insertion(+) 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 From fe810e99bd3374bada88f9bd6f32f6158d0ea984 Mon Sep 17 00:00:00 2001 From: crstngc Date: Fri, 20 Feb 2026 18:31:53 -0700 Subject: [PATCH 7/7] Fix import from inspect in helpers --- rmc/utils/helpers.py | 1 + 1 file changed, 1 insertion(+) diff --git a/rmc/utils/helpers.py b/rmc/utils/helpers.py index 9c57274..09237a7 100644 --- a/rmc/utils/helpers.py +++ b/rmc/utils/helpers.py @@ -2,6 +2,7 @@ """Miscellaneous helper functions.""" +from inspect import isfunction from typing import Any