From 0729a2d365b7ceb010397401a0b9ac5a05fb5c3e Mon Sep 17 00:00:00 2001 From: Georgios Zacharegkas Date: Fri, 15 May 2026 12:05:47 -0500 Subject: [PATCH 1/3] updates to diffmahnet --- diffhalos/mah/diffmahnet_utils.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/diffhalos/mah/diffmahnet_utils.py b/diffhalos/mah/diffmahnet_utils.py index 2ecb239..a3001a2 100644 --- a/diffhalos/mah/diffmahnet_utils.py +++ b/diffhalos/mah/diffmahnet_utils.py @@ -23,7 +23,7 @@ ) from . import diffmahnet - +from .utils import rescale_mah_parameters DEFAULT_MAH_UPARAMS = get_unbounded_mah_params(DEFAULT_MAH_PARAMS) @@ -172,6 +172,7 @@ def get_mah_from_unbounded_params( mah_params_unbound, logt0, t_grid, + m_vals, ): """ Helper function to generate the MAH from @@ -190,6 +191,9 @@ def get_mah_from_unbounded_params( t_grid: ndarray of shape (n_t, ) cosmic time grid at which to compute the MAH + m_vals: ndarray of shape (n_halo, ) + correct mass values for rescaling the MAH parameters + Returns ------- log_mah: ndarray of shape (n_halo, n_t) @@ -204,9 +208,13 @@ def get_mah_from_unbounded_params( ) mah_params_uncorrected = DEFAULT_MAH_PARAMS._make(mah_params_bound) - _, log_mah = mah_halopop(mah_params_uncorrected, t_grid, logt0) + mah_params_corrected = rescale_mah_parameters( + mah_params_uncorrected, m_vals, log_mah[:, -1] + ) + _, log_mah = mah_halopop(mah_params_corrected, t_grid, logt0) + return log_mah From 3bd8800979e5d91072f39a791eb42db4541017f0 Mon Sep 17 00:00:00 2001 From: Georgios Zacharegkas Date: Sun, 5 Jul 2026 13:45:37 -0500 Subject: [PATCH 2/3] adding diffmahnet log_prob and fixing circ import bug --- diffhalos/mah/diffmahnet/diffmahnet.py | 41 +++++++++++++++++++ .../mah/diffmahnet/tests/test_diffmahflow.py | 20 +++++++++ diffhalos/mah/diffmahnet_utils.py | 9 ++-- 3 files changed, 66 insertions(+), 4 deletions(-) diff --git a/diffhalos/mah/diffmahnet/diffmahnet.py b/diffhalos/mah/diffmahnet/diffmahnet.py index a8c82a2..feef531 100644 --- a/diffhalos/mah/diffmahnet/diffmahnet.py +++ b/diffhalos/mah/diffmahnet/diffmahnet.py @@ -206,6 +206,47 @@ def sample( else: return uparam_array + def make_logprob_diffmahnet(self): + @jax.jit + def _logprob_fn(flow_params, lgm_obs, t_obs, uparam): + condition = jnp.array([lgm_obs, t_obs]).T + return self.log_prob_uparams(condition, uparam, flow_params=flow_params) + + return _logprob_fn + + def log_prob_uparams(self, condition, uparam_array, flow_params=None): + """ + Log density of unbounded Diffmah parameters under the conditional flow. + + Parameters + ---------- + condition: ndarray, shape (n, 2) + Columns are ``(lgm_obs, t_obs)`` + + uparam_array: ndarray, shape (n, 5) + Unbounded Diffmah parameters, before applying the flow scaler + + flow_params: ndarray, optional + Flat flow parameters for functional use + + Returns + ------- + ndarray, shape (n,) + Log probabilities ``log p(u_params | condition)`` + """ + flow = ( + self._flow_from_flat_params(flow_params) + if flow_params is not None + else self.flow + ) + + condition_scaled = scaler_transform(condition, self.scaler.u_scaler) + x_scaled = scaler_transform(uparam_array, self.scaler.x_scaler) + + return jax.vmap(lambda xs, cs: flow.log_prob(xs, condition=cs))( + x_scaled, condition_scaled + ) + def get_params(self): param_tree = self._partition()[0] return self.flatten(param_tree) diff --git a/diffhalos/mah/diffmahnet/tests/test_diffmahflow.py b/diffhalos/mah/diffmahnet/tests/test_diffmahflow.py index 82bafd6..c510448 100644 --- a/diffhalos/mah/diffmahnet/tests/test_diffmahflow.py +++ b/diffhalos/mah/diffmahnet/tests/test_diffmahflow.py @@ -1,5 +1,6 @@ import jax import jax.numpy as jnp +import numpy as np from .. import diffmahnet @@ -30,3 +31,22 @@ def test_diffmahflow(): # Make sure asparams=True gives tuple output mahparams_prediction = flow.sample(fake_conditions, keys[2], asparams=True) assert isinstance(mahparams_prediction, tuple) + + +def test_log_prob_uparams_diffmahnet(): + + randkey = jax.random.key(0) + keys = jax.random.split(randkey, 6) + + ndata = 1000 + + # m_obs and t_obs + fake_conditions = jax.random.normal(keys[0], (ndata, 2)) + 1.5 + fake_mah_uparams = jax.random.normal(keys[1], (ndata, 5)) * 0.2 - 4.0 + scaler = diffmahnet.Scaler.compute(fake_mah_uparams, fake_conditions) + + flow = diffmahnet.DiffMahFlow(scaler) + + log_prob_mah = flow.log_prob_uparams(fake_conditions, fake_mah_uparams) + + assert np.all(np.isfinite(log_prob_mah)) diff --git a/diffhalos/mah/diffmahnet_utils.py b/diffhalos/mah/diffmahnet_utils.py index a3001a2..27b4581 100644 --- a/diffhalos/mah/diffmahnet_utils.py +++ b/diffhalos/mah/diffmahnet_utils.py @@ -23,7 +23,6 @@ ) from . import diffmahnet -from .utils import rescale_mah_parameters DEFAULT_MAH_UPARAMS = get_unbounded_mah_params(DEFAULT_MAH_PARAMS) @@ -210,9 +209,11 @@ def get_mah_from_unbounded_params( mah_params_uncorrected = DEFAULT_MAH_PARAMS._make(mah_params_bound) _, log_mah = mah_halopop(mah_params_uncorrected, t_grid, logt0) - mah_params_corrected = rescale_mah_parameters( - mah_params_uncorrected, m_vals, log_mah[:, -1] - ) + # rescale mah parameters + delta_logm_obs = log_mah[:, -1] - m_vals + logm0_rescaled = mah_params_uncorrected.logm0 - delta_logm_obs + mah_params_corrected = mah_params_uncorrected._replace(logm0=logm0_rescaled) + _, log_mah = mah_halopop(mah_params_corrected, t_grid, logt0) return log_mah From 148aa42f37bc06d6f7ce66caf4ab779e8b5419e6 Mon Sep 17 00:00:00 2001 From: Georgios Zacharegkas Date: Tue, 7 Jul 2026 13:46:09 -0500 Subject: [PATCH 3/3] removing new log_prob diffmahnet from PR --- diffhalos/mah/diffmahnet/diffmahnet.py | 41 ------------------- .../mah/diffmahnet/tests/test_diffmahflow.py | 20 --------- 2 files changed, 61 deletions(-) diff --git a/diffhalos/mah/diffmahnet/diffmahnet.py b/diffhalos/mah/diffmahnet/diffmahnet.py index feef531..a8c82a2 100644 --- a/diffhalos/mah/diffmahnet/diffmahnet.py +++ b/diffhalos/mah/diffmahnet/diffmahnet.py @@ -206,47 +206,6 @@ def sample( else: return uparam_array - def make_logprob_diffmahnet(self): - @jax.jit - def _logprob_fn(flow_params, lgm_obs, t_obs, uparam): - condition = jnp.array([lgm_obs, t_obs]).T - return self.log_prob_uparams(condition, uparam, flow_params=flow_params) - - return _logprob_fn - - def log_prob_uparams(self, condition, uparam_array, flow_params=None): - """ - Log density of unbounded Diffmah parameters under the conditional flow. - - Parameters - ---------- - condition: ndarray, shape (n, 2) - Columns are ``(lgm_obs, t_obs)`` - - uparam_array: ndarray, shape (n, 5) - Unbounded Diffmah parameters, before applying the flow scaler - - flow_params: ndarray, optional - Flat flow parameters for functional use - - Returns - ------- - ndarray, shape (n,) - Log probabilities ``log p(u_params | condition)`` - """ - flow = ( - self._flow_from_flat_params(flow_params) - if flow_params is not None - else self.flow - ) - - condition_scaled = scaler_transform(condition, self.scaler.u_scaler) - x_scaled = scaler_transform(uparam_array, self.scaler.x_scaler) - - return jax.vmap(lambda xs, cs: flow.log_prob(xs, condition=cs))( - x_scaled, condition_scaled - ) - def get_params(self): param_tree = self._partition()[0] return self.flatten(param_tree) diff --git a/diffhalos/mah/diffmahnet/tests/test_diffmahflow.py b/diffhalos/mah/diffmahnet/tests/test_diffmahflow.py index c510448..82bafd6 100644 --- a/diffhalos/mah/diffmahnet/tests/test_diffmahflow.py +++ b/diffhalos/mah/diffmahnet/tests/test_diffmahflow.py @@ -1,6 +1,5 @@ import jax import jax.numpy as jnp -import numpy as np from .. import diffmahnet @@ -31,22 +30,3 @@ def test_diffmahflow(): # Make sure asparams=True gives tuple output mahparams_prediction = flow.sample(fake_conditions, keys[2], asparams=True) assert isinstance(mahparams_prediction, tuple) - - -def test_log_prob_uparams_diffmahnet(): - - randkey = jax.random.key(0) - keys = jax.random.split(randkey, 6) - - ndata = 1000 - - # m_obs and t_obs - fake_conditions = jax.random.normal(keys[0], (ndata, 2)) + 1.5 - fake_mah_uparams = jax.random.normal(keys[1], (ndata, 5)) * 0.2 - 4.0 - scaler = diffmahnet.Scaler.compute(fake_mah_uparams, fake_conditions) - - flow = diffmahnet.DiffMahFlow(scaler) - - log_prob_mah = flow.log_prob_uparams(fake_conditions, fake_mah_uparams) - - assert np.all(np.isfinite(log_prob_mah))