Update README.md

README.md

@@ -1,3 +1,127 @@
----
-license: mit
----
+---
+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()
+```