Skip to content

Fixing satquench model #116

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
aphearin merged 2 commits into from
May 5, 2026
Merged

Fixing satquench model #116

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cebc58b
Fixing behaviour of satquench model, which was currently not having a…
alexalar May 5, 2026
24ca072
Updating test_gradients to correctly test the satquench params
alexalar May 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions diffstar/diffstarpop/kernels/diffstarpop_mgash.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,14 +96,21 @@ def _diffstarpop_means_covs(
means_covs = _sfh_pdf_scalar_kernel(sfh_pdf_cens_params, logmp0, tpeak)

# Modify frac_q for satellites
frac_q = means_covs[0]
(
frac_quench_cen,
frac_quench_sat,
) = means_covs[:2]
satquench_params = DEFAULT_SATQUENCHPOP_PARAMS._make(
[getattr(diffstarpop_params, x) for x in DEFAULT_SATQUENCHPOP_PARAMS._fields]
)
frac_q = get_qprob_sat(
satquench_params, lgmu_infall, logmhost_infall, gyr_since_infall, frac_q
frac_quench_sat_updated = get_qprob_sat(
satquench_params,
lgmu_infall,
logmhost_infall,
gyr_since_infall,
frac_quench_sat,
)
means_covs = (frac_q, *means_covs[1:])
means_covs = (frac_quench_cen, frac_quench_sat_updated, *means_covs[2:])
return means_covs


Expand Down
16 changes: 8 additions & 8 deletions diffstar/diffstarpop/tests/test_gradients.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,15 +202,15 @@ def test_gradients_of_diffstarpop_pdf_satquench_params_are_nonzero():
gyr_since_infall,
)
_res = _diffstarpop_means_covs(*args)
frac_quench = _res[0]
(
frac_quench_cen,
frac_quench_sat,
mu_mseq,
mu_qseq,
cov_mseq_ms_block,
cov_qseq_ms_block,
cov_qseq_q_block,
) = _res[1:]
) = _res

# Generate an alternate galpop at some other point in param space
ran_params_key, ran_key = jran.split(ran_key, 2)
Expand All @@ -225,22 +225,22 @@ def test_gradients_of_diffstarpop_pdf_satquench_params_are_nonzero():
gyr_since_infall,
)
_res = _diffstarpop_means_covs(*args)
frac_quench2 = _res[0]
(
frac_quench_cen2,
frac_quench_sat2,
mu_mseq2,
mu_qseq2,
cov_mseq_ms_block2,
cov_qseq_ms_block2,
cov_qseq_q_block2,
) = _res[1:]
) = _res

assert not np.allclose(mu_mseq2, mu_mseq)
assert not np.allclose(cov_qseq_ms_block2, cov_qseq_ms_block)
assert not np.allclose(cov_qseq_q_block2, cov_qseq_q_block)
assert not np.allclose(frac_quench2, frac_quench)
assert not np.allclose(frac_quench_sat, frac_quench_sat2)

frac_q_target = np.copy(frac_quench)
frac_q_target_sat = np.copy(frac_quench_sat)

@jjit
def _loss(u_params):
Expand All @@ -254,8 +254,8 @@ def _loss(u_params):
gyr_since_infall,
)
_res = _diffstarpop_means_covs(*args)
frac_q_pred = _res[0]
return _mse(frac_q_pred, frac_q_target)
frac_q_pred_sat = _res[1]
return _mse(frac_q_pred_sat, frac_q_target_sat)

frac_q_loss, frac_q_grads = value_and_grad(_loss)(alt_dpp_u_params)
assert np.isfinite(frac_q_loss)
Expand Down
Loading