Replace repo contents

README.md

@@ -1,3 +0,0 @@
----
-license: mit
----

README.txt

@@ -0,0 +1,124 @@
+Astro Emulators Toolkit Bundle
+
+Summary:
+  model: transformer_payne
+  release: maja-fe-intensity-tpayne-2ch@0.1.0 (released)
+  bundle_format_version: 1
+  config_schema_version: 1
+  spec_version: 1
+  weights_layout: params_plus_model_state_v1
+  model_family_id: transformer_payne_v1
+  fingerprint_evaluation: present
+  task: regression
+  fit_method: gradient
+  solver_params: not provided
+  solver_diagnostics: not provided
+  solver_design_matrix: not provided
+  role_paths: {'output_leaf': 'outputs/flux', 'parameter_leaf': 'inputs/parameters', 'wavelength_leaf': 'inputs/wavelengths'}
+
+Domain:
+  input_domain: {'kind': 'box_v1', 'max_tree': {'parameters': [7500.0, 5.0, 0.5, 5.0, 0.5, 0.30000001192092896, 0.4000000059604645, 0.4000000059604645, 0.5, 0.5, 1.0], 'wavelengths': 3.7007037171450192}, 'min_tree': {'parameters': [3800.0, 0.0, -2.5, 1.0, -0.20000000298023224, -0.30000001192092896, -0.20000000298023224, -0.20000000298023224, -0.20000000298023224, -0.20000000298023224, 0.1246189996600151], 'wavelengths': 3.6989700043360187}, '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': [7500.0, 5.0, 0.5, 5.0, 0.5, 0.30000001192092896, 0.4000000059604645, 0.4000000059604645, 0.5, 0.5, 1.0], 'wavelengths': 3.7007037171450192}, 'min_tree': {'parameters': [3800.0, 0.0, -2.5, 1.0, -0.20000000298023224, -0.30000001192092896, -0.20000000298023224, -0.20000000298023224, -0.20000000298023224, -0.20000000298023224, 0.1246189996600151], 'wavelengths': 3.6989700043360187}, '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': {'flux': [7.18308162689209, 7.183245658874512]}, 'min_tree': {'flux': [3.442589521408081, 4.903715133666992]}, '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: ['evaluation_scaled_log10_wavelength', 'notes', 'parameter_source_units', 'preprocessing_recipe', 'source_log10_wavelength', 'source_scaled_log10_wavelength', 'source_store', 'source_wavelength', 'training']
+
+Provenance:
+  toolkit_version: 0.1.0
+  created_at: 2026-04-30T14:19:30.339267+00:00
+  python_version: 3.12.13
+  git_commit: ef327b5fc0384a358c53b16364c7c0d688bca7ab
+
+spec:
+input_domain:
+  kind: box_v1
+  max_tree:
+    parameters:
+      array(shape=(11,), dtype=float64)
+    wavelengths: 3.7007037171450192
+  min_tree:
+    parameters:
+      array(shape=(11,), dtype=float64)
+    wavelengths: 3.6989700043360187
+  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:
+      array(shape=(11,), dtype=<U28)
+    wavelengths: None
+  channel_names_tree:
+    parameters:
+      array(shape=(11,), dtype=<U6)
+    wavelengths: None
+  channel_units_tree:
+    parameters:
+      array(shape=(11,), dtype=<U13)
+    wavelengths: None
+  leaf_meanings_tree:
+    parameters: min-max scaled source-grid parameters; see reference_scaling_inputs for raw parameter bounds
+    wavelengths: min-max scaled log10 wavelength; the user applies log10 before reference_scaling_inputs
+  leaf_units_tree:
+    parameters: None
+    wavelengths: dimensionless
+  structure_tree:
+    parameters: None
+    wavelengths: None
+outputs:
+  channel_meanings_tree:
+    flux:
+      - min-max scaled log10 line intensity from source array flux
+      - min-max scaled log10 continuum intensity from source array continuum
+  channel_names_tree:
+    flux:
+      - log_flux_lines
+      - log_flux_continuum
+  channel_units_tree:
+    flux:
+      - dimensionless
+      - dimensionless
+  leaf_meanings_tree:
+    flux: two min-max scaled log10 intensity channels from the Maja archive arrays flux and continuum
+  leaf_units_tree:
+    flux: dimensionless
+  structure_tree:
+    flux: None
+reference_scaling_inputs:
+  applies_to: inputs
+  kind: affine_minmax_v1
+  max_tree:
+    parameters:
+      array(shape=(11,), dtype=float64)
+    wavelengths: 3.7007037171450192
+  min_tree:
+    parameters:
+      array(shape=(11,), dtype=float64)
+    wavelengths: 3.6989700043360187
+  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:
+    flux:
+      - 7.18308162689209
+      - 7.183245658874512
+  min_tree:
+    flux:
+      - 3.442589521408081
+      - 4.903715133666992
+  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,59 @@
+{
+  "algorithm": "sha256",
+  "bundle_id": "sha256:d513b173bef6a04c84ec026d3d519a86318035c8533e5d362609754d0786a687",
+  "integrity_format_version": 1,
+  "tree": [
+    {
+      "path": "README.txt",
+      "sha256": "c0daa7e50513d862171fe0a34ac010ad76953da3fa4a2e1b10aceb23f3d32dfa"
+    },
+    {
+      "path": "config.json",
+      "sha256": "11108861f43c0b23a073534795b59c6f82670da6a45f356ecacc3cbaf7657797"
+    },
+    {
+      "path": "extras/evaluation_scaled_log10_wavelength.safetensors",
+      "sha256": "c00b8e662e97dcd62553c1f13ebf4371cac05933a77fe7165a2600e2e534553d"
+    },
+    {
+      "path": "extras/source_log10_wavelength.safetensors",
+      "sha256": "87630d66ddba8749b3262440edcfa486d86ca1bd044070622a7f85a039095bb3"
+    },
+    {
+      "path": "extras/source_scaled_log10_wavelength.safetensors",
+      "sha256": "f4dcbe6a3d890917ea442b6e775073b042624b3467ea899260c9883142d58b33"
+    },
+    {
+      "path": "extras/source_wavelength.safetensors",
+      "sha256": "32292f9e38853319a5b8e95ff31daa26c1a8e06844542ae436e7cb5155787135"
+    },
+    {
+      "path": "fingerprint_evaluation/inputs.safetensors",
+      "sha256": "2a9da53ff7573f8ef84299c6bbb0990fb49346e06ef17acfcf493977228456bc"
+    },
+    {
+      "path": "fingerprint_evaluation/outputs.safetensors",
+      "sha256": "b47a7750eea84f2cb71c76576f3768c5cd02c5ea5e19e6ee14149fc7f83a93df"
+    },
+    {
+      "path": "input_domain.safetensors",
+      "sha256": "3c641caa7e563613dc40b54822feedc4d0cc16ad90137d6c351129ee96b1b2ca"
+    },
+    {
+      "path": "metadata.json",
+      "sha256": "fa05e398dda0d162fe865c87ee6e67b1087323cea884a69924cedece0d7bb05a"
+    },
+    {
+      "path": "reference_scaling_inputs.safetensors",
+      "sha256": "3c641caa7e563613dc40b54822feedc4d0cc16ad90137d6c351129ee96b1b2ca"
+    },
+    {
+      "path": "reference_scaling_outputs.safetensors",
+      "sha256": "9eaa1947e73b51a122e1dafa8af5a4a07837ee447cb46f791c7c3cc45ddd4530"
+    },
+    {
+      "path": "weights/weights.safetensors",
+      "sha256": "33218ec929ee6f6d7957f1afeca3b3fdd804940f98e5119f9e033814ffaa6cca"
+    }
+  ]
+}

config.json

@@ -0,0 +1,176 @@
+{
+  "bundle": {
+    "bundle_subdir": "bundle"
+  },
+  "hub": {
+    "repo_id": null,
+    "revision": null
+  },
+  "io": {
+    "input_domain": null,
+    "inputs": {
+      "channel_meanings_tree": {
+        "parameters": [
+          "effective temperature",
+          "surface gravity",
+          "metallicity [Fe/H]",
+          "microturbulence velocity",
+          "source-grid parameter [a/Fe]",
+          "source-grid parameter [C/Fe]",
+          "source-grid parameter [N/Fe]",
+          "source-grid parameter [O/Fe]",
+          "source-grid parameter [r/Fe]",
+          "source-grid parameter [s/Fe]",
+          "cosine of viewing angle"
+        ],
+        "wavelengths": null
+      },
+      "channel_names_tree": {
+        "parameters": [
+          "teff",
+          "logg",
+          "[Fe/H]",
+          "vmicro",
+          "[a/Fe]",
+          "[C/Fe]",
+          "[N/Fe]",
+          "[O/Fe]",
+          "[r/Fe]",
+          "[s/Fe]",
+          "mu"
+        ],
+        "wavelengths": null
+      },
+      "channel_units_tree": {
+        "parameters": [
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless"
+        ],
+        "wavelengths": null
+      },
+      "leaf_meanings_tree": {
+        "parameters": "min-max scaled source-grid parameters; see reference_scaling_inputs for raw parameter bounds",
+        "wavelengths": "min-max scaled log10 wavelength; the user applies log10 before reference_scaling_inputs"
+      },
+      "leaf_units_tree": {
+        "parameters": null,
+        "wavelengths": "dimensionless"
+      },
+      "structure_tree": {
+        "parameters": null,
+        "wavelengths": null
+      }
+    },
+    "outputs": {
+      "channel_meanings_tree": {
+        "flux": [
+          "min-max scaled log10 line intensity from source array flux",
+          "min-max scaled log10 continuum intensity from source array continuum"
+        ]
+      },
+      "channel_names_tree": {
+        "flux": [
+          "log_flux_lines",
+          "log_flux_continuum"
+        ]
+      },
+      "channel_units_tree": {
+        "flux": [
+          "dimensionless",
+          "dimensionless"
+        ]
+      },
+      "leaf_meanings_tree": {
+        "flux": "two min-max scaled log10 intensity channels from the Maja archive arrays flux and continuum"
+      },
+      "leaf_units_tree": {
+        "flux": "dimensionless"
+      },
+      "structure_tree": {
+        "flux": null
+      }
+    },
+    "reference_scaling_inputs": null,
+    "reference_scaling_outputs": null
+  },
+  "model": {
+    "init_hints": {
+      "parameter_dim": 11
+    },
+    "name": "transformer_payne",
+    "params": {
+      "channels": 2,
+      "dim": 64,
+      "dim_ff_multiplier": 4,
+      "dim_head": 32,
+      "dtype": "float32",
+      "max_period": 1.0,
+      "min_period": 0.0001,
+      "no_layers": 4,
+      "no_tokens": 8
+    }
+  },
+  "optim": {
+    "b1": 0.9,
+    "b2": 0.999,
+    "decay_steps": 0,
+    "eps": 1e-08,
+    "grad_clip": 0.0,
+    "lr": 0.003,
+    "lr_scaling": null,
+    "name": "soap",
+    "precondition_1d": false,
+    "precondition_frequency": 10,
+    "scale_embedding_lr": 1.0,
+    "schedule": "cosine",
+    "warmup_steps": 1000,
+    "weight_decay": 1e-05
+  },
+  "schema_version": 1,
+  "seed": 0,
+  "solver": {
+    "name": "auto",
+    "params": {}
+  },
+  "task": {
+    "name": "regression",
+    "params": {
+      "loss": "mse",
+      "metric_axes": {
+        "channel": [
+          0
+        ],
+        "global": "all"
+      },
+      "metrics": [
+        "mse",
+        "mae"
+      ]
+    }
+  },
+  "training": {
+    "batch_size": 128,
+    "checkpoint_interval_steps": 1000,
+    "checkpoint_steps": null,
+    "evaluation_interval_steps": 500,
+    "evaluation_steps": null,
+    "logging_interval_steps": 100,
+    "logging_steps": null,
+    "max_saved_checkpoints": 5,
+    "num_steps": 10000,
+    "shuffle": true,
+    "shuffle_seed": 0,
+    "steps_per_epoch": null,
+    "val_fraction": 0.0,
+    "workdir": "./runs/from_bundle"
+  }
+}

extras/evaluation_scaled_log10_wavelength.safetensors

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

extras/source_log10_wavelength.safetensors

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

extras/source_scaled_log10_wavelength.safetensors

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

extras/source_wavelength.safetensors

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

fingerprint_evaluation/inputs.safetensors

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

fingerprint_evaluation/outputs.safetensors

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

input_domain.safetensors

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

metadata.json

@@ -0,0 +1,400 @@
+{
+  "bundle_format_version": 1,
+  "config_schema_version": 1,
+  "extras": {
+    "evaluation_scaled_log10_wavelength": {
+      "__aet_sidecar__": {
+        "format": "safetensors_v1",
+        "layout": "single_array_v1",
+        "path": "extras/evaluation_scaled_log10_wavelength.safetensors"
+      }
+    },
+    "notes": "The source array named 'flux' is treated as line intensity. The source array 'continuum' is treated as continuum intensity. The bundle records min-max scaling metadata; log10 transforms for wavelengths and outputs are explicit user-side preprocessing.",
+    "parameter_source_units": {
+      "[C/Fe]": "",
+      "[Fe/H]": "dex",
+      "[N/Fe]": "",
+      "[O/Fe]": "",
+      "[a/Fe]": "",
+      "[r/Fe]": "",
+      "[s/Fe]": "",
+      "logg": "dex",
+      "mu": "dimensionless",
+      "teff": "K",
+      "vmicro": "km/s"
+    },
+    "preprocessing_recipe": {
+      "outputs": "The model predicts min-max scaled log10 flux and min-max scaled log10 continuum. Invert the min-max transform first; apply 10**y outside the bundle if physical intensities are needed.",
+      "parameters": "Scale raw parameter vector with (x - parameter_min) / (parameter_max - parameter_min).",
+      "wavelengths": "Apply log10 to physical wavelength first, then min-max scale with the log10_wavelength bounds."
+    },
+    "source_log10_wavelength": {
+      "__aet_sidecar__": {
+        "format": "safetensors_v1",
+        "layout": "single_array_v1",
+        "path": "extras/source_log10_wavelength.safetensors"
+      }
+    },
+    "source_scaled_log10_wavelength": {
+      "__aet_sidecar__": {
+        "format": "safetensors_v1",
+        "layout": "single_array_v1",
+        "path": "extras/source_scaled_log10_wavelength.safetensors"
+      }
+    },
+    "source_store": "/Users/tr/data/turbo_spectrum_maja_grids/fe-intensity-20000.zarr",
+    "source_wavelength": {
+      "__aet_sidecar__": {
+        "format": "safetensors_v1",
+        "layout": "single_array_v1",
+        "path": "extras/source_wavelength.safetensors"
+      }
+    },
+    "training": {
+      "dataset_output_keys": [
+        "lines",
+        "continuum"
+      ],
+      "model_output_channel_names": [
+        "log_flux_lines",
+        "log_flux_continuum"
+      ],
+      "n_training_samples": 18000,
+      "n_validation_samples": 2000,
+      "n_wavelength_per_step": 512,
+      "validation_policy": "first 10.0% of spectra by source index"
+    }
+  },
+  "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",
+  "model_family_id": "transformer_payne_v1",
+  "model_init": {
+    "hints": {
+      "parameter_dim": 11
+    },
+    "representation": "model-local init hints only"
+  },
+  "provenance": {
+    "created_at": "2026-04-30T14:19:30.339267+00:00",
+    "dependencies": {
+      "flax": "0.12.6",
+      "jax": "0.9.2",
+      "numpy": "2.4.4",
+      "optax": "0.2.8"
+    },
+    "git_commit": "ef327b5fc0384a358c53b16364c7c0d688bca7ab",
+    "platform": "macOS-26.4.1-arm64-arm-64bit",
+    "python_version": "3.12.13",
+    "toolkit": "astro_emulators_toolkit",
+    "toolkit_version": "0.1.0"
+  },
+  "release": {
+    "name": "maja-fe-intensity-tpayne-2ch",
+    "status": "released",
+    "version": "0.1.0"
+  },
+  "resolved": {
+    "model": {
+      "name": "transformer_payne",
+      "params": {
+        "activation": "gelu",
+        "alpha_att": 1.0,
+        "alpha_emb": 1.0,
+        "bias_attention": false,
+        "bias_dense": false,
+        "bias_feed_forward": null,
+        "bias_output_head": null,
+        "bias_parameter_embedding": null,
+        "channels": 2,
+        "dim": 64,
+        "dim_ff_multiplier": 4,
+        "dim_head": 32,
+        "dtype": "float32",
+        "emb_init": "si",
+        "ff_init": "si",
+        "head_init": "si",
+        "init_att_o": "si",
+        "init_att_q": "si",
+        "max_period": 1.0,
+        "min_period": 0.0001,
+        "no_layers": 4,
+        "no_tokens": 8,
+        "output_activation": "linear",
+        "reference_depth": null,
+        "reference_width": null,
+        "sigma": 1.0
+      }
+    },
+    "solver": {
+      "name": "gradient",
+      "params": {}
+    },
+    "task": {
+      "name": "regression",
+      "params": {
+        "loss": "mse",
+        "loss_weights": null,
+        "metric_axes": {
+          "channel": [
+            0
+          ],
+          "global": "all"
+        },
+        "metrics": [
+          "mse",
+          "mae"
+        ]
+      }
+    }
+  },
+  "runtime_contract": {
+    "affine_leaf_specs": {
+      "inputs/parameters": {
+        "last_axis": 11,
+        "mode": "scalar_or_last_axis"
+      },
+      "inputs/wavelengths": {
+        "mode": "scalar_only"
+      },
+      "outputs/flux": {
+        "last_axis": 2,
+        "mode": "scalar_or_last_axis"
+      }
+    },
+    "role_paths": {
+      "output_leaf": "outputs/flux",
+      "parameter_leaf": "inputs/parameters",
+      "wavelength_leaf": "inputs/wavelengths"
+    },
+    "surface": "canonical_dict_trees_v1",
+    "transformer_payne_channels": [
+      {
+        "dataset_key": "lines",
+        "name": "log_flux_lines"
+      },
+      {
+        "dataset_key": "continuum",
+        "name": "log_flux_continuum"
+      }
+    ]
+  },
+  "spec": {
+    "input_domain": {
+      "kind": "box_v1",
+      "max_tree": {
+        "parameters": [
+          7500.0,
+          5.0,
+          0.5,
+          5.0,
+          0.5,
+          0.30000001192092896,
+          0.4000000059604645,
+          0.4000000059604645,
+          0.5,
+          0.5,
+          1.0
+        ],
+        "wavelengths": 3.7007037171450192
+      },
+      "min_tree": {
+        "parameters": [
+          3800.0,
+          0.0,
+          -2.5,
+          1.0,
+          -0.20000000298023224,
+          -0.30000001192092896,
+          -0.20000000298023224,
+          -0.20000000298023224,
+          -0.20000000298023224,
+          -0.20000000298023224,
+          0.1246189996600151
+        ],
+        "wavelengths": 3.6989700043360187
+      },
+      "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": [
+          "effective temperature",
+          "surface gravity",
+          "metallicity [Fe/H]",
+          "microturbulence velocity",
+          "source-grid parameter [a/Fe]",
+          "source-grid parameter [C/Fe]",
+          "source-grid parameter [N/Fe]",
+          "source-grid parameter [O/Fe]",
+          "source-grid parameter [r/Fe]",
+          "source-grid parameter [s/Fe]",
+          "cosine of viewing angle"
+        ],
+        "wavelengths": null
+      },
+      "channel_names_tree": {
+        "parameters": [
+          "teff",
+          "logg",
+          "[Fe/H]",
+          "vmicro",
+          "[a/Fe]",
+          "[C/Fe]",
+          "[N/Fe]",
+          "[O/Fe]",
+          "[r/Fe]",
+          "[s/Fe]",
+          "mu"
+        ],
+        "wavelengths": null
+      },
+      "channel_units_tree": {
+        "parameters": [
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless",
+          "dimensionless"
+        ],
+        "wavelengths": null
+      },
+      "leaf_meanings_tree": {
+        "parameters": "min-max scaled source-grid parameters; see reference_scaling_inputs for raw parameter bounds",
+        "wavelengths": "min-max scaled log10 wavelength; the user applies log10 before reference_scaling_inputs"
+      },
+      "leaf_units_tree": {
+        "parameters": null,
+        "wavelengths": "dimensionless"
+      },
+      "structure_tree": {
+        "parameters": null,
+        "wavelengths": null
+      }
+    },
+    "outputs": {
+      "channel_meanings_tree": {
+        "flux": [
+          "min-max scaled log10 line intensity from source array flux",
+          "min-max scaled log10 continuum intensity from source array continuum"
+        ]
+      },
+      "channel_names_tree": {
+        "flux": [
+          "log_flux_lines",
+          "log_flux_continuum"
+        ]
+      },
+      "channel_units_tree": {
+        "flux": [
+          "dimensionless",
+          "dimensionless"
+        ]
+      },
+      "leaf_meanings_tree": {
+        "flux": "two min-max scaled log10 intensity channels from the Maja archive arrays flux and continuum"
+      },
+      "leaf_units_tree": {
+        "flux": "dimensionless"
+      },
+      "structure_tree": {
+        "flux": null
+      }
+    },
+    "reference_scaling_inputs": {
+      "applies_to": "inputs",
+      "kind": "affine_minmax_v1",
+      "max_tree": {
+        "parameters": [
+          7500.0,
+          5.0,
+          0.5,
+          5.0,
+          0.5,
+          0.30000001192092896,
+          0.4000000059604645,
+          0.4000000059604645,
+          0.5,
+          0.5,
+          1.0
+        ],
+        "wavelengths": 3.7007037171450192
+      },
+      "min_tree": {
+        "parameters": [
+          3800.0,
+          0.0,
+          -2.5,
+          1.0,
+          -0.20000000298023224,
+          -0.30000001192092896,
+          -0.20000000298023224,
+          -0.20000000298023224,
+          -0.20000000298023224,
+          -0.20000000298023224,
+          0.1246189996600151
+        ],
+        "wavelengths": 3.6989700043360187
+      },
+      "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": {
+        "flux": [
+          7.18308162689209,
+          7.183245658874512
+        ]
+      },
+      "min_tree": {
+        "flux": [
+          3.442589521408081,
+          4.903715133666992
+        ]
+      },
+      "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
+  },
+  "weights_layout": "params_plus_model_state_v1"
+}

reference_likelihood.py

@@ -0,0 +1,279 @@
+"""Quickstart: Gaussian likelihood for Maja log10 intensity spectra.
+
+This compact reference script is meant to travel with the released bundle. It
+uses only bundle metadata, weights, and extras; it does not require the original
+Zarr training archive.
+
+The pattern is:
+
+1. load the local bundle with ``Emulator.from_bundle(...)``;
+2. freeze a JAX callable with ``make_frozen_apply(jit=False)``;
+3. explicitly transform physical parameters and log10 wavelengths into the
+   bundle's canonical min-max space;
+4. explicitly denormalize canonical outputs back to log10 intensities;
+5. evaluate a simple diagonal Gaussian log likelihood in one outer jit.
+
+Inputs are the 11 source-grid parameters recorded in bundle metadata and a
+physical wavelength vector. Wavelengths are converted to log10 before the
+bundle's input min-max normalization, matching the training script.
+
+The output leaf is ``flux`` with two channels:
+
+- ``log_flux_lines``: log10 line intensity from the source array named ``flux``;
+- ``log_flux_continuum``: log10 continuum intensity.
+
+Replace the synthetic ``observed_log10`` vector below with your measured or
+simulated log10 intensity data and uncertainties.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+os.environ.setdefault("JAX_ENABLE_X64", "1")
+os.environ.setdefault("MPLCONFIGDIR", "/private/tmp/matplotlib")
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from astro_emulators_toolkit import Emulator, denormalize_tree, normalize_tree
+
+BUNDLE_DIR = Path(__file__).resolve().parent
+OUTPUT_LEAF = "flux"
+OUTPUT_CHANNELS = ("log_flux_lines", "log_flux_continuum")
+N_LIKELIHOOD_WAVELENGTHS = 1_000
+SIGMA_LOG10_INTENSITY = 0.02
+NOISE_SEED = 0
+
+
+def _as_float_array(value) -> np.ndarray:
+    return np.asarray(value, dtype=np.float32)
+
+
+def _parameter_names(emu: Emulator) -> tuple[str, ...]:
+    inputs = emu.input_spec or {}
+    channel_tree = inputs.get("channel_names_tree")
+    if isinstance(channel_tree, dict):
+        names = channel_tree.get("parameters")
+        if isinstance(names, (list, tuple)):
+            return tuple(str(name) for name in names)
+    ref_inputs = emu.reference_scaling_inputs
+    if ref_inputs is None:
+        return ()
+    n_params = int(np.asarray(ref_inputs["min_tree"]["parameters"]).shape[-1])
+    return tuple(f"parameter_{i}" for i in range(n_params))
+
+
+def _midpoint_physical_parameters(emu: Emulator) -> np.ndarray:
+    ref_inputs = emu.reference_scaling_inputs
+    if ref_inputs is None:
+        raise ValueError("Bundle is missing reference_scaling_inputs metadata.")
+    lo = _as_float_array(ref_inputs["min_tree"]["parameters"])
+    hi = _as_float_array(ref_inputs["max_tree"]["parameters"])
+    return ((lo + hi) * 0.5).astype(np.float32)
+
+
+def _select_wavelengths(emu: Emulator, *, n_wavelength: int) -> np.ndarray:
+    if n_wavelength <= 0:
+        raise ValueError(f"n_wavelength must be positive, got {n_wavelength}")
+
+    extras = emu.bundle_extras
+    source = extras.get("source_wavelength")
+    if source is None:
+        ref_inputs = emu.reference_scaling_inputs
+        if ref_inputs is None:
+            raise ValueError(
+                "Bundle is missing both extras['source_wavelength'] and "
+                "reference_scaling_inputs metadata."
+            )
+        log_lo = float(ref_inputs["min_tree"]["wavelengths"])
+        log_hi = float(ref_inputs["max_tree"]["wavelengths"])
+        return np.power(
+            10.0,
+            np.linspace(log_lo, log_hi, n_wavelength, dtype=np.float32),
+        ).astype(np.float32)
+
+    wave = _as_float_array(source)
+    if wave.ndim != 1:
+        raise ValueError(f"extras['source_wavelength'] must be 1D, got {wave.shape}")
+    n = min(n_wavelength, int(wave.shape[0]))
+    indices = np.linspace(0, wave.shape[0] - 1, num=n, dtype=np.int32)
+    return wave[indices].astype(np.float32)
+
+
+def _show_plot(
+    *,
+    wavelength: np.ndarray,
+    model_log10: np.ndarray,
+    observed_log10: np.ndarray,
+) -> None:
+    import matplotlib.pyplot as plt
+
+    fig, axes = plt.subplots(2, 2, figsize=(13, 7), sharex=True)
+    ax_lines, ax_cont = axes[0]
+    ax_lines_resid, ax_cont_resid = axes[1]
+
+    for channel, name, ax, ax_resid in (
+        (0, OUTPUT_CHANNELS[0], ax_lines, ax_lines_resid),
+        (1, OUTPUT_CHANNELS[1], ax_cont, ax_cont_resid),
+    ):
+        ax.plot(wavelength, model_log10[:, channel], label="model", lw=1.2)
+        ax.scatter(
+            wavelength,
+            observed_log10[:, channel],
+            label="observed + Gaussian noise",
+            s=4,
+            alpha=0.45,
+            linewidths=0,
+        )
+        ax.set_title(name)
+        ax.set_ylabel("log10 intensity")
+        ax.grid(alpha=0.25)
+        ax.legend()
+
+        ax_resid.scatter(
+            wavelength,
+            observed_log10[:, channel] - model_log10[:, channel],
+            s=4,
+            alpha=0.55,
+            linewidths=0,
+        )
+        ax_resid.axhline(0.0, color="0.45", lw=0.8, ls="--")
+        ax_resid.set_xlabel("Wavelength")
+        ax_resid.set_ylabel("observed - model")
+        ax_resid.grid(alpha=0.25)
+
+    fig.suptitle("Maja intensity reference likelihood")
+    fig.tight_layout()
+    plt.show()
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        "--show",
+        action="store_true",
+        help="Show an interactive plot of model, noisy observation, and residuals.",
+    )
+    parser.add_argument(
+        "--n-wavelength",
+        type=int,
+        default=N_LIKELIHOOD_WAVELENGTHS,
+        help="Number of wavelength points to evaluate and plot.",
+    )
+    parser.add_argument(
+        "--sigma-log10",
+        type=float,
+        default=SIGMA_LOG10_INTENSITY,
+        help="Gaussian noise standard deviation in log10 intensity units.",
+    )
+    parser.add_argument(
+        "--noise-seed",
+        type=int,
+        default=NOISE_SEED,
+        help="Random seed for the synthetic Gaussian observation noise.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    emu = Emulator.from_bundle(BUNDLE_DIR, verbose=True)
+    apply_intensity = 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."
+        )
+
+    theta0 = _midpoint_physical_parameters(emu)
+    wavelength = _select_wavelengths(emu, n_wavelength=int(args.n_wavelength))
+
+    def predict_log10_intensity(theta, wavelength_physical):
+        """Predict log10 intensity channels; jit the outer objective."""
+        log10_wavelength = jnp.log10(wavelength_physical)
+        x_physical = {
+            "parameters": theta[None, :],
+            "wavelengths": log10_wavelength[None, :],
+        }
+        x_scaled = normalize_tree(
+            x_physical,
+            ref_inputs["min_tree"],
+            ref_inputs["max_tree"],
+        )
+        y_scaled = apply_intensity(x_scaled)
+        y_log10 = denormalize_tree(
+            y_scaled,
+            ref_outputs["min_tree"],
+            ref_outputs["max_tree"],
+        )
+        return y_log10[OUTPUT_LEAF][0]
+
+    theta_ref = jnp.asarray(theta0, dtype=jnp.float32)
+    wave_ref = jnp.asarray(wavelength, dtype=jnp.float32)
+    model_ref = predict_log10_intensity(theta_ref, wave_ref)
+
+    noise_key = jax.random.key(int(args.noise_seed))
+    sigma_value = float(args.sigma_log10)
+    if sigma_value <= 0.0:
+        raise ValueError(f"sigma-log10 must be positive, got {sigma_value}")
+    noise = jax.random.normal(
+        noise_key,
+        shape=model_ref.shape,
+        dtype=model_ref.dtype,
+    )
+    observed_log10 = jax.lax.stop_gradient(model_ref + sigma_value * noise)
+    sigma_log10 = jnp.full_like(observed_log10, sigma_value)
+
+    @jax.jit
+    def evaluate_likelihood(theta):
+        y_model = predict_log10_intensity(theta, wave_ref)
+        resid = (observed_log10 - y_model) / sigma_log10
+        log_norm = jnp.sum(jnp.log(2.0 * jnp.pi * sigma_log10**2))
+        log_likelihood = -0.5 * (jnp.sum(resid**2) + log_norm)
+        return y_model, log_likelihood
+
+    model_log10_jax, logp_jax = evaluate_likelihood(theta_ref)
+    model_log10 = np.asarray(jax.block_until_ready(model_log10_jax))
+    observed_np = np.asarray(jax.block_until_ready(observed_log10))
+    logp = float(jax.block_until_ready(logp_jax))
+
+    print("bundle:", BUNDLE_DIR)
+    print("parameter vector:")
+    for name, value in zip(_parameter_names(emu), theta0, strict=True):
+        print(f"  {name}: {float(value):.6g}")
+    print(
+        "wavelength range:",
+        f"{float(wavelength[0]):.6g} .. {float(wavelength[-1]):.6g}",
+        f"({wavelength.shape[0]} points)",
+    )
+    print("output channels:", OUTPUT_CHANNELS)
+    print("model log10 output shape:", model_log10.shape)
+    for i, name in enumerate(OUTPUT_CHANNELS):
+        resid = observed_np[:, i] - model_log10[:, i]
+        print(
+            f"  {name}: model median={float(np.median(model_log10[:, i])):.6f}, "
+            f"residual rms={float(np.sqrt(np.mean(resid**2))):.6f}"
+        )
+    print("sigma_log10:", f"{sigma_value:.6f}")
+    print("noise_seed:", int(args.noise_seed))
+    print("log_likelihood:", f"{logp:.6f}")
+
+    if args.show:
+        _show_plot(
+            wavelength=wavelength,
+            model_log10=model_log10,
+            observed_log10=observed_np,
+        )
+
+
+if __name__ == "__main__":
+    main()

reference_scaling_inputs.safetensors

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

reference_scaling_outputs.safetensors

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

weights/weights.safetensors

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