|
|
--- |
|
|
license: mit |
|
|
--- |
|
|
``` |
|
|
from pathlib import Path |
|
|
|
|
|
import numpy as np |
|
|
import matplotlib.pyplot as plt |
|
|
import jax |
|
|
import jax.numpy as jnp |
|
|
import time |
|
|
|
|
|
from astro_emulators_toolkit import Emulator |
|
|
|
|
|
script_dir = Path(__file__).parent.resolve() |
|
|
|
|
|
# ------------------------------------------------------------------------------ |
|
|
# Model description and data scaling info for physical prediction |
|
|
# ------------------------------------------------------------------------------ |
|
|
|
|
|
DEFAULT_INPUTS = ("age", "eep", "feh") |
|
|
DEFAULT_TARGETS = ("G_mag", "BP_mag", "RP_mag") |
|
|
|
|
|
MIN_VAL = np.array([5.8619833, 202.0, -0.87977487, -2.3778718, -2.4398916, -2.2926207], dtype=np.float32) |
|
|
MAX_VAL = np.array([1.02993574e01, 4.54000000e02, 5.95229030e-01, 1.50175705e01, 1.84394169e01, 1.36201954e01], dtype=np.float32) |
|
|
|
|
|
|
|
|
# ------------------------------------------------------------------------------ |
|
|
# Load pretrained emulator bundle from Hugging Face and build a physical predictor |
|
|
# ------------------------------------------------------------------------------ |
|
|
|
|
|
print("Attempting to load pretrained emulator bundle from Hugging Face...") |
|
|
repo_id = "RozanskiT/isochrones-mlp" |
|
|
try: |
|
|
emu = Emulator.from_pretrained( |
|
|
repo_id, |
|
|
cache_dir=script_dir / ".emuspec_cache", |
|
|
) |
|
|
print(f"Loaded pretrained emulator from Hugging Face: {repo_id}") |
|
|
except Exception as exc: |
|
|
print(f"Hugging Face load failed ({exc}).") |
|
|
|
|
|
|
|
|
# ------------------------------------------------------------------------------ |
|
|
# Build a physical predictor that scales inputs and applies the frozen model |
|
|
# ------------------------------------------------------------------------------ |
|
|
|
|
|
def build_physical_predictor(emu: Emulator): |
|
|
"""Return a jitted predictor that scales physical inputs then applies frozen model.""" |
|
|
|
|
|
frozen_apply = emu.make_frozen_apply_fn(postprocess=True, jit=False) |
|
|
x_min = jax.device_put(MIN_VAL[:3]) |
|
|
x_scale = jax.device_put(MAX_VAL[:3] - MIN_VAL[:3]) |
|
|
y_min = jax.device_put(MIN_VAL[3:]) |
|
|
y_scale = jax.device_put(MAX_VAL[3:] - MIN_VAL[3:]) |
|
|
|
|
|
@jax.jit |
|
|
def predict_physical(x_physical): |
|
|
x_norm = (x_physical - x_min) / x_scale |
|
|
y_norm = frozen_apply(x_norm) |
|
|
return y_norm * y_scale + y_min |
|
|
|
|
|
return predict_physical |
|
|
|
|
|
predict_physical = build_physical_predictor(emu) |
|
|
|
|
|
|
|
|
# ------------------------------------------------------------------------------ |
|
|
# Make some physical inputs |
|
|
# ------------------------------------------------------------------------------ |
|
|
|
|
|
no_points = 1000 |
|
|
batch_of_predictions = np.zeros((no_points, 3)) # dummy batch of 10 input points with 3 features (age, eep, feh) |
|
|
batch_of_predictions[:,0] = 9.4 # age |
|
|
batch_of_predictions[:,1] = np.linspace(202, 454, no_points) # eep |
|
|
batch_of_predictions[:,2] = 0.0 # feh |
|
|
|
|
|
# simplified check of domain: |
|
|
assert np.all(batch_of_predictions[:, 0] >= MIN_VAL[0]) and np.all(batch_of_predictions[:, 0] <= MAX_VAL[0]), "Age out of domain" |
|
|
assert np.all(batch_of_predictions[:, 1] >= MIN_VAL[1]) and np.all(batch_of_predictions[:, 1] <= MAX_VAL[1]), "EEP out of domain" |
|
|
assert np.all(batch_of_predictions[:, 2] >= MIN_VAL[2]) and np.all(batch_of_predictions[:, 2] <= MAX_VAL[2]), "FeH out of domain" |
|
|
|
|
|
# move to jax (eg. GPU when availible) |
|
|
batch_of_predictions = jnp.array(batch_of_predictions) |
|
|
|
|
|
# ------------------------------------------------------------------------------ |
|
|
# Create a predictor that uses the frozen model but scales physical inputs, then predict on the batch and time it |
|
|
# This could be extended to include a distance, extinction, by analitical model |
|
|
# ------------------------------------------------------------------------------ |
|
|
|
|
|
# A bit of timing info to see how fast the predictions are after the initial compilation. |
|
|
t0 = time.perf_counter() |
|
|
y_pred_first = predict_physical(batch_of_predictions) |
|
|
y_pred_first = np.asarray(jax.block_until_ready(y_pred_first)) |
|
|
t1 = time.perf_counter() |
|
|
|
|
|
y_pred_second = predict_physical(batch_of_predictions) |
|
|
y_pred_second = np.asarray(jax.block_until_ready(y_pred_second)) |
|
|
t2 = time.perf_counter() |
|
|
|
|
|
# Summarize timings and prediction shape |
|
|
print(f"First call (compile + run): {t1 - t0:.6f} s") |
|
|
print(f"Second call (run only): {t2 - t1:.6f} s") |
|
|
print(f"Predictions size: {y_pred_second.shape}") |
|
|
|
|
|
|
|
|
# color-magnitude at the left and magntude vs step at the right |
|
|
fig, axs = plt.subplots(1, 2, figsize=(12, 4)) |
|
|
ax_cmd = axs[0] |
|
|
ax_cmd.scatter(y_pred_second[:, 1] - y_pred_second[:, 2], y_pred_second[:, 0], s=18, alpha=0.8, color="tab:orange") |
|
|
ax_cmd.set_xlabel("BP - RP") |
|
|
ax_cmd.set_ylabel("G") |
|
|
ax_cmd.set_title(f"CMD (batch of {no_points} predictions)") |
|
|
ax_cmd.grid(alpha=0.25) |
|
|
ax_cmd.invert_yaxis() # Magnitudes are brighter when smaller, so invert y-axis for CMD |
|
|
|
|
|
ax_step = axs[1] |
|
|
for i in range(y_pred_second.shape[1]): |
|
|
ax_step.plot(y_pred_second[:, i], "-", color="tab:orange", alpha=0.9, label=f"Pred {DEFAULT_TARGETS[i]}") |
|
|
ax_step.set_xlabel("Batch Index") |
|
|
ax_step.set_ylabel("Magnitude") |
|
|
ax_step.set_title(f"Predicted {DEFAULT_TARGETS[i]} vs Batch Index") |
|
|
ax_step.legend() |
|
|
ax_step.grid(alpha=0.25) |
|
|
plt.tight_layout() |
|
|
plt.show() |
|
|
``` |
