Skip to content

extended redshift range to 1.0 with new models #447

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

Draft
djones1040 wants to merge 11 commits into
base: main
Choose a base branch
from
Draft

extended redshift range to 1.0 with new models #447

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
45f405e
extended redshift range to 1.0 with new models
Jul 9, 2026
cfe677d
same model for local and global; in future planning to refactor this …
Jul 9, 2026
6f01720
debug line removed
Jul 9, 2026
b12434b
reverted previous change, but anticipating global/local models will b…
Jul 9, 2026
2e54a05
adding standard deviation of S/N to our S/N files and fixing longstan…
Jul 9, 2026
b44f65d
Merge remote-tracking branch 'origin/main' into djones-frankenblast
Jul 12, 2026
96ac013
updated to low-z and high-z SBI++ models. Mag offsets for z < 0.015 …
Jul 15, 2026
ba3dcf1
fine-tuning parameters
Jul 15, 2026
eda1970
apply mag offset to low-z transients based on software version
Jul 15, 2026
e6691ed
updated S/N after fixing vega/AB bug affecting 2MASS, WISE
Jul 18, 2026
4a8336f
fixed duplicate mag values
Jul 18, 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
96 changes: 48 additions & 48 deletions app/host/SBI/run_sbi_blast.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,10 @@
}

sbi_params = {
"anpe_fname_global": f"{settings.SBIPP_ROOT}/SBI_model_global.pt", # trained sbi model
"train_fname_global": f"{settings.SBIPP_PHOT_ROOT}/sbi_phot_global.h5", # training set
"anpe_fname_local": f"{settings.SBIPP_ROOT}/SBI_model_local.pt", # trained sbi model
"train_fname_local": f"{settings.SBIPP_PHOT_ROOT}/sbi_phot_local.h5", # training set
"anpe_fname_standard": f"{settings.SBIPP_ROOT}/SBI_model_standard.pt", # trained sbi model
"train_fname_standard": f"{settings.SBIPP_PHOT_ROOT}/sbi_phot_standard.h5", # training set
"anpe_fname_lowz": f"{settings.SBIPP_ROOT}/SBI_model_lowz.pt", # trained sbi model
"train_fname_lowz": f"{settings.SBIPP_PHOT_ROOT}/sbi_phot_lowz.h5", # training set
"nhidden": 500, # architecture of the trained density estimator
"nblocks": 15, # architecture of the trained density estimator
}
Expand Down Expand Up @@ -76,7 +76,7 @@

# training set
def run_training_set():
for _fit_type in ["global", "local"]:
for _fit_type in ["standard", "lowz"]:
data = h5py.File(sbi_params[f"train_fname_{_fit_type}"], "r")
x_train = np.array(data["theta"]) # physical parameters
y_train = np.array(data["phot"]) # fluxes & uncertainties
Expand Down Expand Up @@ -109,72 +109,72 @@ def run_training_set():
)
)
anpe._x_shape = Ut.x_shape_from_simulation(y_tensor)
if _fit_type == "global":
hatp_x_y_global = anpe.build_posterior(
if _fit_type == "standard":
hatp_x_y_standard = anpe.build_posterior(
p_x_y_estimator, sample_with="rejection"
)
y_train_global = y_train[:]
x_train_global = x_train[:]
elif _fit_type == "local":
hatp_x_y_local = anpe.build_posterior(
y_train_standard = y_train[:]
x_train_standard = x_train[:]
elif _fit_type == "lowz":
hatp_x_y_lowz = anpe.build_posterior(
p_x_y_estimator, sample_with="rejection"
)
y_train_local = y_train[:]
x_train_local = x_train[:]
y_train_lowz = y_train[:]
x_train_lowz = x_train[:]

print("""Storing training sets as data files...""")
with open(
os.path.join(settings.SBI_TRAINING_ROOT, "hatp_x_y_global.pkl"), "wb"
os.path.join(settings.SBI_TRAINING_ROOT, "hatp_x_y_standard.pkl"), "wb"
) as handle:
pickle.dump(hatp_x_y_global, handle)
pickle.dump(hatp_x_y_standard, handle)
with open(
os.path.join(settings.SBI_TRAINING_ROOT, "y_train_global.pkl"), "wb"
os.path.join(settings.SBI_TRAINING_ROOT, "y_train_standard.pkl"), "wb"
) as handle:
pickle.dump(y_train_global, handle)
pickle.dump(y_train_standard, handle)
with open(
os.path.join(settings.SBI_TRAINING_ROOT, "x_train_global.pkl"), "wb"
os.path.join(settings.SBI_TRAINING_ROOT, "x_train_standard.pkl"), "wb"
) as handle:
pickle.dump(x_train_global, handle)
pickle.dump(x_train_standard, handle)
with open(
os.path.join(settings.SBI_TRAINING_ROOT, "hatp_x_y_local.pkl"), "wb"
os.path.join(settings.SBI_TRAINING_ROOT, "hatp_x_y_lowz.pkl"), "wb"
) as handle:
pickle.dump(hatp_x_y_local, handle)
pickle.dump(hatp_x_y_lowz, handle)
with open(
os.path.join(settings.SBI_TRAINING_ROOT, "y_train_local.pkl"), "wb"
os.path.join(settings.SBI_TRAINING_ROOT, "y_train_lowz.pkl"), "wb"
) as handle:
pickle.dump(y_train_local, handle)
pickle.dump(y_train_lowz, handle)
with open(
os.path.join(settings.SBI_TRAINING_ROOT, "x_train_local.pkl"), "wb"
os.path.join(settings.SBI_TRAINING_ROOT, "x_train_lowz.pkl"), "wb"
) as handle:
pickle.dump(x_train_local, handle)
pickle.dump(x_train_lowz, handle)


try:
print("""Loading training sets from data files...""")
with open(
os.path.join(settings.SBI_TRAINING_ROOT, "hatp_x_y_global.pkl"), "rb"
os.path.join(settings.SBI_TRAINING_ROOT, "hatp_x_y_standard.pkl"), "rb"
) as handle:
hatp_x_y_global = pickle.load(handle)
hatp_x_y_standard = pickle.load(handle)
with open(
os.path.join(settings.SBI_TRAINING_ROOT, "y_train_global.pkl"), "rb"
os.path.join(settings.SBI_TRAINING_ROOT, "y_train_standard.pkl"), "rb"
) as handle:
y_train_global = pickle.load(handle)
y_train_standard = pickle.load(handle)
with open(
os.path.join(settings.SBI_TRAINING_ROOT, "x_train_global.pkl"), "rb"
os.path.join(settings.SBI_TRAINING_ROOT, "x_train_standard.pkl"), "rb"
) as handle:
x_train_global = pickle.load(handle)
x_train_standard = pickle.load(handle)
with open(
os.path.join(settings.SBI_TRAINING_ROOT, "hatp_x_y_local.pkl"), "rb"
os.path.join(settings.SBI_TRAINING_ROOT, "hatp_x_y_lowz.pkl"), "rb"
) as handle:
hatp_x_y_local = pickle.load(handle)
hatp_x_y_lowz = pickle.load(handle)
with open(
os.path.join(settings.SBI_TRAINING_ROOT, "y_train_local.pkl"), "rb"
os.path.join(settings.SBI_TRAINING_ROOT, "y_train_lowz.pkl"), "rb"
) as handle:
y_train_local = pickle.load(handle)
y_train_lowz = pickle.load(handle)
with open(
os.path.join(settings.SBI_TRAINING_ROOT, "x_train_local.pkl"), "rb"
os.path.join(settings.SBI_TRAINING_ROOT, "x_train_lowz.pkl"), "rb"
) as handle:
x_train_local = pickle.load(handle)
x_train_lowz = pickle.load(handle)
print("""Training sets loaded.""")
except Exception as err:
print(f"""Error loading training sets: {err}. Regenerating...""")
Expand All @@ -189,14 +189,14 @@ def maggies_to_asinh(x):
return -a * math.asinh((x / 2.0) * np.exp(mu / a)) + mu


def fit_sbi_pp(observations, n_filt_cuts=True, fit_type="global"):
def fit_sbi_pp(observations, n_filt_cuts=True, fit_type="standard"):
np.random.seed(100) # make results reproducible

# toy noise model
meds_sigs, stds_sigs = [], []

for f in all_filters:
toy_noise_x, toy_noise_y = np.loadtxt(
toy_noise_x, toy_noise_y, toy_noise_std = np.loadtxt(
f"host/SBI/snrfiles/{f.name}_magvsnr.txt", dtype=float, unpack=True
)
meds_sigs += [
Expand All @@ -211,7 +211,7 @@ def fit_sbi_pp(observations, n_filt_cuts=True, fit_type="global"):
stds_sigs += [
interp1d(
toy_noise_x,
1.0857 * 1 / toy_noise_y,
(1.0857 * 1 / toy_noise_y)*(toy_noise_std/toy_noise_y),
kind="slinear",
fill_value="extrapolate", # (0.01,1.0),
bounds_error=False,
Expand Down Expand Up @@ -268,14 +268,14 @@ def fit_sbi_pp(observations, n_filt_cuts=True, fit_type="global"):
return {}, 1

# prepare to pass the reconstructed model to sbi_pp
if fit_type == "global":
sbi_params["y_train"] = y_train_global
sbi_params["theta_train"] = x_train_global
sbi_params["hatp_x_y"] = hatp_x_y_global
elif fit_type == "local":
sbi_params["y_train"] = y_train_local
sbi_params["hatp_x_y"] = hatp_x_y_local
sbi_params["theta_train"] = x_train_local
if fit_type == "standard":
sbi_params["y_train"] = y_train_standard
sbi_params["theta_train"] = x_train_standard
sbi_params["hatp_x_y"] = hatp_x_y_standard
elif fit_type == "lowz":
sbi_params["y_train"] = y_train_lowz
sbi_params["hatp_x_y"] = hatp_x_y_lowz
sbi_params["theta_train"] = x_train_lowz

# Run SBI++
chain, obs, flags = sbi_pp.sbi_pp(
Expand Down
82 changes: 46 additions & 36 deletions app/host/SBI/sbi_pp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,12 +118,13 @@ def chidof(mags, obsphot, obsphot_unc, individual=False):
return chi / np.sum(np.isfinite(obsphot))


def gauss_approx_missingband(obs, run_params, sbi_params, max_neighbors=200):
def gauss_approx_missingband(obs, run_params, sbi_params, max_neighbors=10):
"""nearest neighbor approximation of missing bands;
see sec. 4.1.2 for details
"""
use_res = False
y_train = sbi_params["y_train"]
redshift_idx = np.where(abs(obs["redshift"]-sbi_params["y_train"][:,-1]) < 0.05)[0]
y_train = sbi_params["y_train"][redshift_idx]

y_obs = np.copy(obs["mags_sbi"])
sig_obs = np.copy(obs["mags_unc_sbi"])
Expand Down Expand Up @@ -190,9 +191,10 @@ def sbi_missingband(obs, run_params, sbi_params, seconditer=False):
print("sbi missing bands")
ave_theta = []

max_neighbors = 200
max_neighbors = 10
redshift_idx = np.where(abs(obs["redshift"]-sbi_params["y_train"][:,-1]) < 0.05)[0]
hatp_x_y = sbi_params["hatp_x_y"]
y_train = sbi_params["y_train"]
y_train = sbi_params["y_train"][redshift_idx]
y_obs = np.copy(obs["mags_sbi"])
sig_obs = np.copy(obs["mags_unc_sbi"])
invalid_mask = np.copy(obs["missing_mask"])
Expand Down Expand Up @@ -288,20 +290,21 @@ def sbi_missingband(obs, run_params, sbi_params, seconditer=False):
x[22:][not_valid_idx[j]] = y_train[idx_chi2_selected][
np.random.choice(range(len(guess_ndata.T[j])))
][22:][not_valid_idx[j]]
# x[not_valid_idx_unc[j]] = toy_noise(
# flux=x[not_valid_idx[j]],
# meds_sigs=sbi_params["toynoise_meds_sigs"][idx],
# stds_sigs=sbi_params["toynoise_stds_sigs"][idx],
# verbose=run_params["verbose"],
# )[1]
# D. Jones - 7/11/26 -- re-trying toy noise
#x[not_valid_idx_unc[j]] = toy_noise(
# flux=x[not_valid_idx[j]],
# meds_sigs=sbi_params["toynoise_meds_sigs"][idx],
# stds_sigs=sbi_params["toynoise_stds_sigs"][idx],
# verbose=run_params["verbose"],
#)[1]

# the noise model in the training isn't quite right
# Pan-STARRS in particular seems a little off
# we'll have to re-train at some point, but for now just pull
# uncertainties from the training sample
for idx, fname in zip(valid_idx, obs["filternames"][valid_idx]):
chc = np.random.choice(range(len(y_train[idx_chi2_selected])))
x[22:-1][idx] = y_train[idx_chi2_selected][chc][22:-1][idx]
#for idx, fname in zip(valid_idx, obs["filternames"][valid_idx]):
# chc = np.random.choice(range(len(y_train[idx_chi2_selected])))
# x[22:-1][idx] = y_train[idx_chi2_selected][chc][22:-1][idx]

all_x.append(x)

Expand Down Expand Up @@ -371,7 +374,8 @@ def lim_of_noisy_guass(obs, run_params, sbi_params):

use_res = 1

y_train = sbi_params["y_train"]
redshift_idx = np.where(abs(obs["redshift"]-sbi_params["y_train"][:,-1]) < 0.05)[0]
y_train = sbi_params["y_train"][redshift_idx]

y_obs = np.copy(obs["mags_sbi"])
sig_obs = np.copy(obs["mags_unc_sbi"])
Expand Down Expand Up @@ -420,7 +424,7 @@ def lim_of_noisy_guass(obs, run_params, sbi_params):
return lims, use_res


def sbi_mcnoise(obs, run_params, sbi_params, max_neighbors=200):
def sbi_mcnoise(obs, run_params, sbi_params, max_neighbors=10):
"""used when observations have out-of-distribution uncertainties;
see sec. 4.1.1 for details
"""
Expand All @@ -431,8 +435,9 @@ def sbi_mcnoise(obs, run_params, sbi_params, max_neighbors=200):

ave_theta = []

redshift_idx = np.where(abs(obs["redshift"]-sbi_params["y_train"][:,-1]) < 0.05)[0]
hatp_x_y = sbi_params["hatp_x_y"]
y_train = sbi_params["y_train"]
y_train = sbi_params["y_train"][redshift_idx]

y_obs = np.copy(obs["mags_sbi"])
sig_obs = np.copy(obs["mags_unc_sbi"])
Expand Down Expand Up @@ -498,11 +503,11 @@ def sbi_mcnoise(obs, run_params, sbi_params, max_neighbors=200):
)[1]
# signal.alarm(run_params["tmax_per_obj"])

for idx, fname in enumerate(obs["filternames"]):
chc = np.random.choice(range(len(y_train[idx_chi2_selected])))
samp_y_guess[22:-1][idx] = y_train[idx_chi2_selected][chc][22:-1][
idx
]
#for idx, fname in enumerate(obs["filternames"]):
# chc = np.random.choice(range(len(y_train[idx_chi2_selected])))
# samp_y_guess[22:-1][idx] = y_train[idx_chi2_selected][chc][22:-1][
# idx
# ]

do_continue = False
for tmax, npost in zip(
Expand Down Expand Up @@ -564,8 +569,9 @@ def sbi_missing_and_noisy(obs, run_params, sbi_params):

ave_theta = []

redshift_idx = np.where(abs(obs["redshift"]-sbi_params["y_train"][:,-1]) < 0.05)[0]
hatp_x_y = sbi_params["hatp_x_y"]
y_train = sbi_params["y_train"]
y_train = sbi_params["y_train"][redshift_idx]

y_obs = np.copy(obs["mags_sbi"])
sig_obs = np.copy(obs["mags_unc_sbi"])
Expand Down Expand Up @@ -642,10 +648,10 @@ def sbi_missing_and_noisy(obs, run_params, sbi_params):
# Pan-STARRS in particular seems a little off
# we'll have to re-train at some point, but for now just pull
# uncertainties from the training sample
for idx, fname in zip(valid_idx, obs["filternames"][valid_idx]):
#for idx, fname in zip(valid_idx, obs["filternames"][valid_idx]):
# if 'PanSTARRS' in fname or '2MASS' in fname or 'SDSS' in fname or 'DES' in fname:
chc = np.random.choice(range(len(y_train[idx_chi2_selected])))
samp_y_guess[22:][idx] = y_train[idx_chi2_selected][chc][22:][idx]
# chc = np.random.choice(range(len(y_train[idx_chi2_selected])))
# samp_y_guess[22:][idx] = y_train[idx_chi2_selected][chc][22:][idx]

# if we can't get one posterior sample in one second, we should move along
# to the next MC sample
Expand Down Expand Up @@ -696,7 +702,7 @@ def sbi_missing_and_noisy(obs, run_params, sbi_params):
return ave_theta, use_res, timeout_flag, cnt


def sbi_baseline(obs, run_params, sbi_params, max_neighbors=200):
def sbi_baseline(obs, run_params, sbi_params, max_neighbors=10):
signal.signal(signal.SIGALRM, timeout_handler)

if run_params["verbose"]:
Expand All @@ -714,8 +720,9 @@ def sbi_baseline(obs, run_params, sbi_params, max_neighbors=200):
"nsamp_noisy": -99, # number of MC samples drawn
}

redshift_idx = np.where(abs(obs["redshift"]-sbi_params["y_train"][:,-1]) < 0.05)[0]
hatp_x_y = sbi_params["hatp_x_y"]
y_train = sbi_params["y_train"]
y_train = sbi_params["y_train"][redshift_idx]

y_obs = np.copy(obs["mags"])
sig_obs = np.copy(obs["mags_unc"])
Expand Down Expand Up @@ -754,9 +761,9 @@ def sbi_baseline(obs, run_params, sbi_params, max_neighbors=200):

x = np.concatenate([y_obs, sig_obs, [obs["redshift"]]])

for idx, fname in enumerate(obs["filternames"]):
chc = np.random.choice(range(len(y_train[idx_chi2_selected])))
x[22:-1][idx] = y_train[idx_chi2_selected][chc][22:-1][idx]
#for idx, fname in enumerate(obs["filternames"]):
# chc = np.random.choice(range(len(y_train[idx_chi2_selected])))
# x[22:-1][idx] = y_train[idx_chi2_selected][chc][22:-1][idx]

try:
ave_theta = hatp_x_y.sample(
Expand All @@ -777,7 +784,7 @@ def sbi_baseline(obs, run_params, sbi_params, max_neighbors=200):
return ave_theta, obs, flags


def sbi_pp(obs, run_params, sbi_params, max_neighbors=200):
def sbi_pp(obs, run_params, sbi_params, max_neighbors=10):
"""wrapper for sbi++; this should be the only function needed to be called outside this scipt under normal circumstances

obs: a dictionary at least containing
Expand All @@ -801,11 +808,13 @@ def sbi_pp(obs, run_params, sbi_params, max_neighbors=200):
"nsamp_noisy": -99, # number of MC samples drawn
}

redshift_idx = np.where(abs(obs["redshift"]-sbi_params["y_train"][:,-1]) < 0.05)[0]
hatp_x_y = sbi_params["hatp_x_y"]
y_train = sbi_params["y_train"]
y_train = sbi_params["y_train"][redshift_idx]

y_obs = np.copy(obs["mags"])
sig_obs = np.copy(obs["mags_unc"])

# copy the observed data to be used by sbi
# missing data, if any, will be filled in later
obs["mags_sbi"] = y_obs
Expand All @@ -825,9 +834,10 @@ def sbi_pp(obs, run_params, sbi_params, max_neighbors=200):
stds_sigs=sbi_params["toynoise_stds_sigs"][j],
verbose=run_params["verbose"],
)
noisy_mask[j] = (sig_obs[j] - _toynoise[1]) / _toynoise[2] >= run_params[
noisy_mask[j] = abs((sig_obs[j] - _toynoise[1]) / _toynoise[2]) >= run_params[
"noisy_sig"
]

# if noisy_mask[j]:
# import pdb; pdb.set_trace()
noisy_mask &= np.isfinite(y_obs) # idx of noisy bands
Expand Down Expand Up @@ -868,9 +878,9 @@ def sbi_pp(obs, run_params, sbi_params, max_neighbors=200):

x = np.concatenate([y_obs, sig_obs, [obs["redshift"]]])

for idx, fname in enumerate(obs["filternames"]):
chc = np.random.choice(range(len(y_train[idx_chi2_selected])))
x[22:-1][idx] = y_train[idx_chi2_selected][chc][22:-1][idx]
#for idx, fname in enumerate(obs["filternames"]):
# chc = np.random.choice(range(len(y_train[idx_chi2_selected])))
# x[22:-1][idx] = y_train[idx_chi2_selected][chc][22:-1][idx]

try:
ave_theta = hatp_x_y.sample(
Expand Down
Loading