Replace repo contents

README.md

@@ -1,127 +0,0 @@
----
-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()
-```

README.txt

@@ -0,0 +1,129 @@
+Astro Emulators Toolkit Bundle
+
+Summary:
+  model: mlp
+  release: mist-isochrone-modern-mlp-10m-128x3@0.1.0-collab.1 (released)
+  bundle_format_version: 1
+  config_schema_version: 1
+  spec_version: 1
+  weights_layout: params_plus_model_state_v1
+  model_family_id: mlp_v1
+  fingerprint_evaluation: present
+  task: regression
+  fit_method: gradient
+  solver_params: not provided
+  solver_diagnostics: not provided
+  solver_design_matrix: not provided
+  role_paths: {'input_leaf': 'inputs/parameters', 'output_leaf': 'outputs/magnitudes'}
+
+Domain:
+  input_domain: {'kind': 'box_v1', 'max_tree': {'parameters': [10.299357414245605, 454.0, 0.5952290296554565]}, 'min_tree': {'parameters': [5.861983299255371, 202.0, -0.8797748684883118]}, 'storage': {'filename': 'input_domain.safetensors', 'format': 'safetensors_v1', 'layout': 'split_minmax_tree_v1'}, 'value_space': 'physical_input_dict_tree_v1'}
+  reference_scaling_inputs: {'applies_to': 'inputs', 'kind': 'affine_minmax_v1', 'max_tree': {'parameters': [10.299357414245605, 454.0, 0.5952290296554565]}, 'min_tree': {'parameters': [5.861983299255371, 202.0, -0.8797748684883118]}, 'source_space': 'physical_input_dict_tree_v1', 'storage': {'filename': 'reference_scaling_inputs.safetensors', 'format': 'safetensors_v1', 'layout': 'split_minmax_tree_v1'}, 'target_space': 'canonical_input_dict_tree_v1'}
+  reference_scaling_outputs: {'applies_to': 'outputs', 'kind': 'affine_minmax_v1', 'max_tree': {'magnitudes': [15.017570495605469, 18.439416885375977, 13.620195388793945]}, 'min_tree': {'magnitudes': [-2.3778717517852783, -2.4398915767669678, -2.2926206588745117]}, 'source_space': 'canonical_output_dict_tree_v1', 'storage': {'filename': 'reference_scaling_outputs.safetensors', 'format': 'safetensors_v1', 'layout': 'split_minmax_tree_v1'}, 'target_space': 'physical_output_dict_tree_v1'}
+  extras: not provided
+
+Provenance:
+  toolkit_version: 0.1.0
+  created_at: 2026-04-21T18:10:12.535436+00:00
+  python_version: 3.12.13
+  git_commit: b3415cfd04a48359232624dba9a1a746cf91313f
+
+spec:
+input_domain:
+  kind: box_v1
+  max_tree:
+    parameters:
+      - 10.299357414245605
+      - 454.0
+      - 0.5952290296554565
+  min_tree:
+    parameters:
+      - 5.861983299255371
+      - 202.0
+      - -0.8797748684883118
+  storage:
+    filename: input_domain.safetensors
+    format: safetensors_v1
+    layout: split_minmax_tree_v1
+  value_space: physical_input_dict_tree_v1
+inputs:
+  channel_meanings_tree:
+    parameters:
+      - log10 stellar age in years
+      - equivalent evolutionary phase
+      - metallicity relative to solar
+  channel_names_tree:
+    parameters:
+      - age
+      - eep
+      - feh
+  channel_units_tree:
+    parameters:
+      - log10(age [yr])
+      - 
+      - [Fe/H]
+  leaf_meanings_tree: None
+  leaf_units_tree: None
+  structure_tree:
+    parameters: None
+outputs:
+  channel_meanings_tree:
+    magnitudes:
+      - absolute Gaia G magnitude
+      - absolute Gaia BP magnitude
+      - absolute Gaia RP magnitude
+  channel_names_tree:
+    magnitudes:
+      - G_mag
+      - BP_mag
+      - RP_mag
+  channel_units_tree:
+    magnitudes:
+      - abs_mag
+      - abs_mag
+      - abs_mag
+  leaf_meanings_tree: None
+  leaf_units_tree: None
+  structure_tree:
+    magnitudes: None
+reference_scaling_inputs:
+  applies_to: inputs
+  kind: affine_minmax_v1
+  max_tree:
+    parameters:
+      - 10.299357414245605
+      - 454.0
+      - 0.5952290296554565
+  min_tree:
+    parameters:
+      - 5.861983299255371
+      - 202.0
+      - -0.8797748684883118
+  source_space: physical_input_dict_tree_v1
+  storage:
+    filename: reference_scaling_inputs.safetensors
+    format: safetensors_v1
+    layout: split_minmax_tree_v1
+  target_space: canonical_input_dict_tree_v1
+reference_scaling_outputs:
+  applies_to: outputs
+  kind: affine_minmax_v1
+  max_tree:
+    magnitudes:
+      - 15.017570495605469
+      - 18.439416885375977
+      - 13.620195388793945
+  min_tree:
+    magnitudes:
+      - -2.3778717517852783
+      - -2.4398915767669678
+      - -2.2926206588745117
+  source_space: canonical_output_dict_tree_v1
+  storage:
+    filename: reference_scaling_outputs.safetensors
+    format: safetensors_v1
+    layout: split_minmax_tree_v1
+  target_space: physical_output_dict_tree_v1
+spec_version: 1
+
+Note: this bundle is the canonical emulator artifact. Physical-space composition is external.

bundle_integrity.json

@@ -0,0 +1,43 @@
+{
+  "algorithm": "sha256",
+  "bundle_id": "sha256:f1d5c3f5792c57ac841a07dfa844f1637a357d402fa7e01a2139989d559462ab",
+  "integrity_format_version": 1,
+  "tree": [
+    {
+      "path": "README.txt",
+      "sha256": "5ecf592e929e06c29b11da5d2286047e58107465971c928276562130a5829cc7"
+    },
+    {
+      "path": "config.json",
+      "sha256": "8beec961e9bac4a9c13f087cb8a70af1fd91caf7e09733f2a681a133a02e1b7d"
+    },
+    {
+      "path": "fingerprint_evaluation/inputs.safetensors",
+      "sha256": "e411d571388dd13d2c5c9f020386c14378f2a63e63a8a28182b3ac2a45f60e60"
+    },
+    {
+      "path": "fingerprint_evaluation/outputs.safetensors",
+      "sha256": "5b168542ec3041851b3ac0bf8dd278c0ccb3f53623dc7c12f6172305e6ed25a2"
+    },
+    {
+      "path": "input_domain.safetensors",
+      "sha256": "5d8acd2ed09e2a417ff6f83a4482d0aa88be08ad949144969dc5f2e3b9329c66"
+    },
+    {
+      "path": "metadata.json",
+      "sha256": "3d6032c9918be3c88751164653632090eb22d1f578e1fb93c23faf60f83ec03a"
+    },
+    {
+      "path": "reference_scaling_inputs.safetensors",
+      "sha256": "5d8acd2ed09e2a417ff6f83a4482d0aa88be08ad949144969dc5f2e3b9329c66"
+    },
+    {
+      "path": "reference_scaling_outputs.safetensors",
+      "sha256": "d3420c76743566a2093846f876d0472c0da282e67ff09b8dab72b7b16c7bd5d1"
+    },
+    {
+      "path": "weights/weights.safetensors",
+      "sha256": "d407b20c90e930cd335ee6b3ef287114a1ae635c12323c126d3e51e043abda7b"
+    }
+  ]
+}

config.json

@@ -2,41 +2,129 @@
   "bundle": {
     "bundle_subdir": "bundle"
   },
-  "data": {
-    "columns": null,
-    "dtype": "float32",
-    "inputs": [],
-    "memmap": true,
-    "path": "",
-    "targets": []
-  },
   "hub": {
     "repo_id": null,
     "revision": null
   },
   "io": {
-    "x_dim": 3,
-    "x_names": [
-      "age",
-      "eep",
-      "feh"
-    ],
-    "y_dim": 3,
-    "y_names": [
-      "G_mag",
-      "BP_mag",
-      "RP_mag"
-    ]
+    "input_domain": {
+      "max_tree": {
+        "parameters": [
+          10.299357414245605,
+          454.0,
+          0.5952290296554565
+        ]
+      },
+      "min_tree": {
+        "parameters": [
+          5.861983299255371,
+          202.0,
+          -0.8797748684883118
+        ]
+      }
+    },
+    "inputs": {
+      "channel_meanings_tree": {
+        "parameters": [
+          "log10 stellar age in years",
+          "equivalent evolutionary phase",
+          "metallicity relative to solar"
+        ]
+      },
+      "channel_names_tree": {
+        "parameters": [
+          "age",
+          "eep",
+          "feh"
+        ]
+      },
+      "channel_units_tree": {
+        "parameters": [
+          "log10(age [yr])",
+          "",
+          "[Fe/H]"
+        ]
+      },
+      "leaf_meanings_tree": null,
+      "leaf_units_tree": null,
+      "structure_tree": {
+        "parameters": null
+      }
+    },
+    "outputs": {
+      "channel_meanings_tree": {
+        "magnitudes": [
+          "absolute Gaia G magnitude",
+          "absolute Gaia BP magnitude",
+          "absolute Gaia RP magnitude"
+        ]
+      },
+      "channel_names_tree": {
+        "magnitudes": [
+          "G_mag",
+          "BP_mag",
+          "RP_mag"
+        ]
+      },
+      "channel_units_tree": {
+        "magnitudes": [
+          "abs_mag",
+          "abs_mag",
+          "abs_mag"
+        ]
+      },
+      "leaf_meanings_tree": null,
+      "leaf_units_tree": null,
+      "structure_tree": {
+        "magnitudes": null
+      }
+    },
+    "reference_scaling_inputs": {
+      "max_tree": {
+        "parameters": [
+          10.299357414245605,
+          454.0,
+          0.5952290296554565
+        ]
+      },
+      "min_tree": {
+        "parameters": [
+          5.861983299255371,
+          202.0,
+          -0.8797748684883118
+        ]
+      }
+    },
+    "reference_scaling_outputs": {
+      "max_tree": {
+        "magnitudes": [
+          15.017570495605469,
+          18.439416885375977,
+          13.620195388793945
+        ]
+      },
+      "min_tree": {
+        "magnitudes": [
+          -2.3778717517852783,
+          -2.4398915767669678,
+          -2.2926206588745117
+        ]
+      }
+    }
   },
   "model": {
+    "init_hints": {
+      "input_last_axis": 3,
+      "output_last_axis": 3
+    },
     "name": "mlp",
     "params": {
       "activation": "gelu",
       "dtype": "float32",
       "hidden_sizes": [
-        64,
-        64,
-        64
+        128,
+        128,
+        128
       ]
     }
   },
@@ -45,16 +133,20 @@
     "b2": 0.999,
     "decay_steps": 0,
     "eps": 1e-08,
-    "lr": 0.003,
+    "lr": 0.001,
     "name": "soap",
     "precondition_1d": false,
-    "precondition_frequency": 10,
+    "precondition_frequency": 20,
     "schedule": "cosine",
-    "warmup_steps": 10000,
+    "warmup_steps": 1000000,
     "weight_decay": 1e-05
   },
   "schema_version": 1,
   "seed": 0,
+  "solver": {
+    "name": "auto",
+    "params": {}
+  },
   "task": {
     "name": "regression",
     "params": {
@@ -71,15 +163,18 @@
   },
   "training": {
     "batch_size": 2048,
-    "checkpoint_every_steps": 150000,
-    "eval_every_steps": 5000,
-    "log_every_steps": 1000,
-    "max_checkpoints": 3,
-    "num_steps": 100000,
+    "checkpoint_interval_steps": null,
+    "checkpoint_steps": null,
+    "evaluation_interval_steps": 50000,
+    "evaluation_steps": null,
+    "logging_interval_steps": 10000,
+    "logging_steps": null,
+    "max_saved_checkpoints": 0,
+    "num_steps": 10000000,
     "shuffle": true,
     "shuffle_seed": 0,
     "steps_per_epoch": null,
     "val_fraction": 0.1,
-    "workdir": "/Users/tr/repos/astro-emulators-toolkit/large_local_examples/isochrones/run2"
+    "workdir": "./runs/from_bundle"
   }
 }

weights.safetensors → fingerprint_evaluation/inputs.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7916af20b55930166ba2f40866892ecf13cec5c9354d827ac44055d6797bf8e8
-size 35876
+oid sha256:e411d571388dd13d2c5c9f020386c14378f2a63e63a8a28182b3ac2a45f60e60
+size 92

fingerprint_evaluation/outputs.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b168542ec3041851b3ac0bf8dd278c0ccb3f53623dc7c12f6172305e6ed25a2
+size 92

input_domain.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d8acd2ed09e2a417ff6f83a4482d0aa88be08ad949144969dc5f2e3b9329c66
+size 192

metadata.json

@@ -1,23 +1,235 @@
 {
+  "bundle_format_version": 1,
+  "config_schema_version": 1,
+  "extras": {},
+  "fingerprint_evaluation": {
+    "atol": 1e-07,
+    "inputs": {
+      "filename": "fingerprint_evaluation/inputs.safetensors",
+      "format": "safetensors_v1",
+      "layout": "numeric_dict_tree_v1",
+      "space": "canonical_input_dict_trees_v1"
+    },
+    "kind": "canonical_inputs_outputs_v1",
+    "outputs": {
+      "filename": "fingerprint_evaluation/outputs.safetensors",
+      "format": "safetensors_v1",
+      "layout": "numeric_dict_tree_v1",
+      "space": "canonical_output_dict_trees_v1"
+    },
+    "rtol": 1e-05,
+    "selection_strategy": "midpoint_from_input_domain_then_reference_scaling_inputs_v1"
+  },
   "fit_method": "gradient",
-  "schema_version": 1,
-  "tool": "astro_emulators_toolkit",
-  "x_schema": {
-    "dim": 3,
-    "names": [
-      "age",
-      "eep",
-      "feh"
-    ],
-    "representation": "model-space inputs emitted directly by the dataset"
+  "model_family_id": "mlp_v1",
+  "model_init": {
+    "hints": {
+      "input_last_axis": 3,
+      "output_last_axis": 3
+    },
+    "representation": "model-local init hints only"
+  },
+  "provenance": {
+    "created_at": "2026-04-21T18:10:12.535436+00:00",
+    "dependencies": {
+      "flax": "0.12.6",
+      "jax": "0.9.2",
+      "numpy": "2.4.4",
+      "optax": "0.2.8"
+    },
+    "git_commit": "b3415cfd04a48359232624dba9a1a746cf91313f",
+    "platform": "macOS-26.4.1-arm64-arm-64bit",
+    "python_version": "3.12.13",
+    "toolkit": "astro_emulators_toolkit",
+    "toolkit_version": "0.1.0"
+  },
+  "release": {
+    "name": "mist-isochrone-modern-mlp-10m-128x3",
+    "status": "released",
+    "version": "0.1.0-collab.1"
+  },
+  "resolved": {
+    "model": {
+      "name": "mlp",
+      "params": {
+        "activation": "gelu",
+        "dtype": "float32",
+        "hidden_sizes": [
+          128,
+          128,
+          128
+        ],
+        "use_bias": true
+      }
+    },
+    "solver": {
+      "name": "gradient",
+      "params": {}
+    },
+    "task": {
+      "name": "regression",
+      "params": {
+        "loss": "mse",
+        "loss_weights": null,
+        "metric_axes": {
+          "global": "all",
+          "per_dim": []
+        },
+        "metrics": [
+          "mse",
+          "mae"
+        ]
+      }
+    }
+  },
+  "runtime_contract": {
+    "affine_leaf_specs": {
+      "inputs/parameters": {
+        "last_axis": 3,
+        "mode": "scalar_or_last_axis"
+      },
+      "outputs/magnitudes": {
+        "last_axis": 3,
+        "mode": "scalar_or_last_axis"
+      }
+    },
+    "role_paths": {
+      "input_leaf": "inputs/parameters",
+      "output_leaf": "outputs/magnitudes"
+    },
+    "surface": "canonical_dict_trees_v1"
+  },
+  "spec": {
+    "input_domain": {
+      "kind": "box_v1",
+      "max_tree": {
+        "parameters": [
+          10.299357414245605,
+          454.0,
+          0.5952290296554565
+        ]
+      },
+      "min_tree": {
+        "parameters": [
+          5.861983299255371,
+          202.0,
+          -0.8797748684883118
+        ]
+      },
+      "storage": {
+        "filename": "input_domain.safetensors",
+        "format": "safetensors_v1",
+        "layout": "split_minmax_tree_v1"
+      },
+      "value_space": "physical_input_dict_tree_v1"
+    },
+    "inputs": {
+      "channel_meanings_tree": {
+        "parameters": [
+          "log10 stellar age in years",
+          "equivalent evolutionary phase",
+          "metallicity relative to solar"
+        ]
+      },
+      "channel_names_tree": {
+        "parameters": [
+          "age",
+          "eep",
+          "feh"
+        ]
+      },
+      "channel_units_tree": {
+        "parameters": [
+          "log10(age [yr])",
+          "",
+          "[Fe/H]"
+        ]
+      },
+      "leaf_meanings_tree": null,
+      "leaf_units_tree": null,
+      "structure_tree": {
+        "parameters": null
+      }
+    },
+    "outputs": {
+      "channel_meanings_tree": {
+        "magnitudes": [
+          "absolute Gaia G magnitude",
+          "absolute Gaia BP magnitude",
+          "absolute Gaia RP magnitude"
+        ]
+      },
+      "channel_names_tree": {
+        "magnitudes": [
+          "G_mag",
+          "BP_mag",
+          "RP_mag"
+        ]
+      },
+      "channel_units_tree": {
+        "magnitudes": [
+          "abs_mag",
+          "abs_mag",
+          "abs_mag"
+        ]
+      },
+      "leaf_meanings_tree": null,
+      "leaf_units_tree": null,
+      "structure_tree": {
+        "magnitudes": null
+      }
+    },
+    "reference_scaling_inputs": {
+      "applies_to": "inputs",
+      "kind": "affine_minmax_v1",
+      "max_tree": {
+        "parameters": [
+          10.299357414245605,
+          454.0,
+          0.5952290296554565
+        ]
+      },
+      "min_tree": {
+        "parameters": [
+          5.861983299255371,
+          202.0,
+          -0.8797748684883118
+        ]
+      },
+      "source_space": "physical_input_dict_tree_v1",
+      "storage": {
+        "filename": "reference_scaling_inputs.safetensors",
+        "format": "safetensors_v1",
+        "layout": "split_minmax_tree_v1"
+      },
+      "target_space": "canonical_input_dict_tree_v1"
+    },
+    "reference_scaling_outputs": {
+      "applies_to": "outputs",
+      "kind": "affine_minmax_v1",
+      "max_tree": {
+        "magnitudes": [
+          15.017570495605469,
+          18.439416885375977,
+          13.620195388793945
+        ]
+      },
+      "min_tree": {
+        "magnitudes": [
+          -2.3778717517852783,
+          -2.4398915767669678,
+          -2.2926206588745117
+        ]
+      },
+      "source_space": "canonical_output_dict_tree_v1",
+      "storage": {
+        "filename": "reference_scaling_outputs.safetensors",
+        "format": "safetensors_v1",
+        "layout": "split_minmax_tree_v1"
+      },
+      "target_space": "physical_output_dict_tree_v1"
+    },
+    "spec_version": 1
   },
-  "y_schema": {
-    "dim": 3,
-    "names": [
-      "G_mag",
-      "BP_mag",
-      "RP_mag"
-    ],
-    "representation": "model-space targets emitted directly by the dataset"
-  }
+  "weights_layout": "params_plus_model_state_v1"
 }

reference_likelihood.py

@@ -0,0 +1,109 @@
+"""Quickstart: Gaussian likelihood for one Gaia photometric measurement.
+
+This is a compact reference script for collaborators loading the released
+isochrone emulator from Hugging Face. It follows the same pattern as the main
+README examples:
+
+1. download a bundle with ``Emulator.from_pretrained(...)``;
+2. freeze a JAX callable with ``make_frozen_apply(jit=False)``;
+3. explicitly normalize physical inputs into the bundle's canonical space;
+4. explicitly denormalize canonical outputs back to physical magnitudes;
+5. evaluate a simple diagonal Gaussian log likelihood in one outer jit.
+
+The input vector is ``[log10_age_yr, eep, feh]``. The measurement vector is
+absolute ``[G_mag, BP_mag, RP_mag]``. If your data are apparent magnitudes,
+include distance modulus, extinction, calibration offsets, or other nuisance
+terms in your own likelihood around the emulator prediction.
+
+Related examples in the source repository:
+
+- examples/basic/02_load_bundle_predict.py
+- examples/basic/04_use_bundle_in_map_fit.py
+"""
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from astro_emulators_toolkit import Emulator, denormalize_tree, normalize_tree
+
+REPO_ID = "RozanskiT/isochrones-mlp"
+REVISION = None
+CACHE_DIR = ".emuspec_cache"
+
+OUTPUT_LEAF = "magnitudes"
+OUTPUT_CHANNELS = ("G_mag", "BP_mag", "RP_mag")
+
+# One trial isochrone point: log10(age [yr]), EEP, [Fe/H].
+THETA_PHYSICAL = np.asarray([9.4, 300.0, 0.0], dtype=np.float32)
+
+# One example absolute photometric measurement in G, BP, RP.
+# Replace these with your own absolute magnitudes and uncertainties.
+OBSERVED_MAGNITUDES = np.asarray([6.94, 7.52, 6.21], dtype=np.float32)
+OBSERVED_SIGMA_MAG = np.asarray([0.03, 0.03, 0.03], dtype=np.float32)
+
+
+def main() -> None:
+    emu = Emulator.from_pretrained(
+        REPO_ID,
+        revision=REVISION,
+        cache_dir=CACHE_DIR,
+        verbose=True,
+    )
+    apply_magnitudes = emu.make_frozen_apply(jit=False)
+
+    ref_inputs = emu.reference_scaling_inputs
+    ref_outputs = emu.reference_scaling_outputs
+    if ref_inputs is None or ref_outputs is None:
+        raise ValueError(
+            "This likelihood example requires reference_scaling_inputs and "
+            "reference_scaling_outputs in the bundle metadata."
+        )
+
+    y_obs = jnp.asarray(OBSERVED_MAGNITUDES, dtype=jnp.float32)
+    y_err = jnp.asarray(OBSERVED_SIGMA_MAG, dtype=jnp.float32)
+
+    def predict_magnitudes(theta):
+        """Predict physical magnitudes; jit the outer objective, not this helper."""
+        x_physical = {"parameters": theta[None, :]}
+        x_scaled = normalize_tree(
+            x_physical,
+            ref_inputs["min_tree"],
+            ref_inputs["max_tree"],
+        )
+        y_scaled = apply_magnitudes(x_scaled)
+        y_physical = denormalize_tree(
+            y_scaled,
+            ref_outputs["min_tree"],
+            ref_outputs["max_tree"],
+        )
+        return y_physical[OUTPUT_LEAF][0]
+
+    @jax.jit
+    def evaluate_likelihood(theta):
+        y_model = predict_magnitudes(theta)
+        resid = (y_obs - y_model) / y_err
+        log_norm = jnp.sum(jnp.log(2.0 * jnp.pi * y_err**2))
+        log_likelihood = -0.5 * (jnp.sum(resid**2) + log_norm)
+        return y_model, log_likelihood
+
+    theta = jnp.asarray(THETA_PHYSICAL, dtype=jnp.float32)
+    model_magnitudes_jax, logp_jax = evaluate_likelihood(theta)
+    model_magnitudes = np.asarray(jax.block_until_ready(model_magnitudes_jax))
+    logp = float(jax.block_until_ready(logp_jax))
+
+    print("theta_physical [age, eep, feh]:", THETA_PHYSICAL.tolist())
+    print("model absolute magnitudes:")
+    for name, value in zip(OUTPUT_CHANNELS, model_magnitudes, strict=True):
+        print(f"  {name}: {value:.6f}")
+    print("observed absolute magnitudes:")
+    for name, value in zip(OUTPUT_CHANNELS, OBSERVED_MAGNITUDES, strict=True):
+        print(f"  {name}: {value:.6f}")
+    print("sigma_mag:", OBSERVED_SIGMA_MAG.tolist())
+    print("log_likelihood:", f"{logp:.6f}")
+
+
+if __name__ == "__main__":
+    main()

reference_scaling_inputs.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d8acd2ed09e2a417ff6f83a4482d0aa88be08ad949144969dc5f2e3b9329c66
+size 192

reference_scaling_outputs.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3420c76743566a2093846f876d0472c0da282e67ff09b8dab72b7b16c7bd5d1
+size 192

weights/weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d407b20c90e930cd335ee6b3ef287114a1ae635c12323c126d3e51e043abda7b
+size 136588