Skip to content

Add ThermoResult object for streamlined thermo-abundance workflow#19

Open
smsharma wants to merge 1 commit intomainfrom
feature/thermo-result
Open

Add ThermoResult object for streamlined thermo-abundance workflow#19
smsharma wants to merge 1 commit intomainfrom
feature/thermo-result

Conversation

@smsharma
Copy link
Collaborator

Summary

Closes #14

BackgroundModel now returns a ThermoResult NamedTuple that AbundanceModel can accept directly, with automatic temperature range detection.

Changes

  • Add ThermoResult NamedTuple to linx/thermo.py containing all thermo outputs plus T_start/T_end
  • BackgroundModel returns ThermoResult instead of plain tuple
  • AbundanceModel accepts either:
    • ThermoResult object (auto-extracts all values including T_start/T_end)
    • Individual arrays (backward compatible)

Usage

# New streamlined API
thermo_result = BackgroundModel()(0.)
abundances = AbundanceModel(nuclear_net)(thermo_result)
# T_start/T_end auto-detected from thermo_result

# Backward compatible usage
t, a, rho_g, rho_nu, rho_extra, P_extra, Neff = thermo_model(0.)[:7]
abundances = abundance_model(rho_g, rho_nu, rho_extra, P_extra, t_vec=t, a_vec=a)

Test plan

  • Existing tests pass
  • Tuple unpacking backward compatibility verified

🤖 Generated with Claude Code

Closes #14

- Add ThermoResult NamedTuple to linx/thermo.py containing all thermo
  outputs plus T_start/T_end used for integration
- BackgroundModel now returns ThermoResult instead of plain tuple
- AbundanceModel accepts ThermoResult directly with automatic
  temperature range detection from the thermo result
- Maintains backward compatibility via tuple unpacking (use [:7] slice)

Usage:
  # New streamlined API
  thermo_result = BackgroundModel()(0.)
  abundances = AbundanceModel(nuclear_net)(thermo_result)

  # Backward compatible
  t, a, rho_g, rho_nu, rho_extra, P_extra, Neff = thermo_model(0.)[:7]

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@cgiovanetti
Copy link
Owner

Thanks for putting this together--I'm going to hold off on this one because it requires us to also update the example notebooks, and I want to play with it a bit in some different contexts before pulling the trigger on this switch. I'm a little backed up trying to get the stuff on cara/me-hardcode integrated in and waiting for a response on jax-ml/jax#33666, once I clear that away I'll come back to this

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Abundances T_start and T_end could track thermo by default

2 participants