[GSOC] use previous iteration ion populations as first guess in NLTE solver by Aaryan-Dadu · Pull Request #3454 · tardis-sn/tardis

Aaryan-Dadu · 2026-02-20T22:03:41Z

📝 Description

Type: 🚀 feature

This PR improves the convergence speed of the NLTE solver by using the ion populations and electron densities from the previous iteration as the initial guess for subsequent iterations, rather than resetting to the default "all singly ionized" guess every time.

Changes:

tardis/plasma/properties/nlte.py: Added PreviousIonNumberDensity class to store ion number densities from the previous iteration
tardis/plasma/properties/nlte_rate_equation_solver.py: Modified calculate_first_guess to utilize these previous values if they exist
Added Regression Tests for it

Closes #2485

📌 Resources

Examples, notebooks, and links to useful references.

🚦 Testing

How did you test these changes?

Testing pipeline
Other method (self tested and then added regression tests for the same)
My changes can't be tested (explain why)

☑️ Checklist

I requested two reviewers for this pull request
I updated the documentation according to my changes
I built the documentation by applying the build_docs label

Note: If you are not allowed to perform any of these actions, ping (@) a contributor.

AI Usage

Gemini 2.5 is used for research and understanding parts of code.
Github Copilot is used for better naming of variables and framing docstrings sentences

- Implement PreviousIonNumberDensity - update NLTEPopulationSolverRoot to utilize previous iteration values

tardis-bot · 2026-02-20T22:06:10Z

*beep* *bop*
Hi human,
I ran ruff on the latest commit (bf8d951).
Here are the outputs produced.
Results can also be downloaded as artifacts here.
Summarised output:

Details

91	F405  	[ ] undefined-local-with-import-star-usage
8	W293  	[*] blank-line-with-whitespace
2	W291  	[*] trailing-whitespace
2	G004  	[ ] logging-f-string
2	I001  	[*] unsorted-imports
2	RUF022	[*] unsorted-dunder-all
1	F403  	[ ] undefined-local-with-import-star
Found 108 errors.
[*] 14 fixable with the `--fix` option.

Complete output(might be large):

Details

tardis/plasma/properties/nlte.py:9:11: RUF022 [*] `__all__` is not sorted
tardis/plasma/properties/nlte_rate_equation_solver.py:9:11: RUF022 [*] `__all__` is not sorted
tardis/plasma/properties/nlte_rate_equation_solver.py:316:25: G004 Logging statement uses f-string
tardis/plasma/properties/nlte_rate_equation_solver.py:323:25: G004 Logging statement uses f-string
tardis/plasma/properties/property_collections.py:1:1: F403 `from tardis.plasma.properties import *` used; unable to detect undefined names
tardis/plasma/properties/property_collections.py:1:1: I001 [*] Import block is un-sorted or un-formatted
tardis/plasma/properties/property_collections.py:11:9: F405 `DilutePlanckianRadField` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:12:9: F405 `NumberDensity` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:13:9: F405 `TimeExplosion` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:14:9: F405 `AtomicData` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:15:9: F405 `JBlues` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:16:9: F405 `LinkTRadTElectron` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:17:9: F405 `HeliumTreatment` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:18:9: F405 `ContinuumInteractionSpecies` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:19:9: F405 `NLTEIonizationSpecies` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:20:9: F405 `NLTEExcitationSpecies` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:25:9: F405 `TRadiative` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:26:9: F405 `DilutionFactor` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:27:9: F405 `BetaRadiation` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:28:9: F405 `Levels` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:29:9: F405 `Lines` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:30:9: F405 `PartitionFunction` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:31:9: F405 `GElectron` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:32:9: F405 `IonizationData` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:33:9: F405 `LinesLowerLevelIndex` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:34:9: F405 `LinesUpperLevelIndex` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:35:9: F405 `StimulatedEmissionFactor` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:36:9: F405 `SelectedAtoms` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:37:9: F405 `ElectronTemperature` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:38:9: F405 `ThermalLevelBoltzmannFactorLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:39:9: F405 `ThermalLTEPartitionFunction` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:40:9: F405 `BetaElectron` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:41:9: F405 `ThermalGElectron` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:42:9: F405 `ThermalPhiSahaLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:43:9: F405 `SahaFactor` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:46:55: F405 `PhiSahaLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:47:55: F405 `LevelBoltzmannFactorLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:50:6: F405 `PhiSahaNebular` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:50:22: F405 `ZetaData` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:50:32: F405 `BetaElectron` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:50:46: F405 `RadiationFieldCorrection` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:53:6: F405 `LevelBoltzmannFactorDiluteLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:55:49: F405 `LevelBoltzmannFactorNoNLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:58:9: F405 `LevelBoltzmannFactorNLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:59:9: F405 `NLTEData` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:60:9: F405 `PreviousElectronDensities` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:61:9: F405 `PreviousBetaSobolev` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:65:6: F405 `NLTEIndexHelper` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:65:23: F405 `NLTEPopulationSolverRoot` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:65:49: F405 `PreviousIonNumberDensity` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:68:6: F405 `NLTEIndexHelper` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:68:23: F405 `NLTEPopulationSolverLU` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:72:9: F405 `HeliumNLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:73:9: F405 `RadiationFieldCorrection` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:74:9: F405 `ZetaData` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:75:9: F405 `BetaElectron` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:76:9: F405 `LevelNumberDensityHeNLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:77:9: F405 `IonNumberDensityHeNLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:81:6: F405 `LevelNumberDensity` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:81:26: F405 `IonNumberDensity` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:84:6: F405 `HeliumNumericalNLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:88:9: F405 `PhotoIonRateCoeff` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:89:9: F405 `StimRecombRateFactor` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:90:9: F405 `BfHeatingRateCoeffEstimator` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:91:9: F405 `StimRecombCoolingRateCoeffEstimator` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:92:9: F405 `YgData` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:97:9: F405 `StimRecombRateCoeff` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:98:9: F405 `PhotoIonizationData` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:99:9: F405 `SpontRecombRateCoeff` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:100:9: F405 `ThermalLevelBoltzmannFactorLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:101:9: F405 `ThermalLTEPartitionFunction` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:102:9: F405 `BetaElectron` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:103:9: F405 `ThermalGElectron` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:104:9: F405 `ThermalPhiSahaLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:105:9: F405 `SahaFactor` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:106:9: F405 `CorrPhotoIonRateCoeff` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:107:9: F405 `SpontRecombCoolingRateCoeff` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:108:9: F405 `YgInterpolator` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:109:9: F405 `CollExcRateCoeff` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:110:9: F405 `CollDeexcRateCoeff` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:111:9: F405 `RawCollisionTransProbs` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:112:9: F405 `MarkovChainIndex` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:113:9: F405 `FreeFreeCoolingRate` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:114:9: F405 `FreeBoundCoolingRate` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:115:9: F405 `LevelNumberDensityLTE` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:116:9: F405 `PhotoIonBoltzmannFactor` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:117:9: F405 `FreeBoundEmissionCDF` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:118:9: F405 `LevelIdxs2LineIdx` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:119:9: F405 `LevelIdxs2TransitionIdx` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:120:9: F405 `CollIonRateCoeffSeaton` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:121:9: F405 `CollRecombRateCoeff` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:122:9: F405 `ContinuumInteractionHandler` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:126:58: F405 `AdiabaticCoolingRate` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:129:9: F405 `RawTwoPhotonTransProbs` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:130:9: F405 `TwoPhotonData` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:131:9: F405 `TwoPhotonEmissionCDF` may be undefined, or defined from star imports
tardis/plasma/properties/property_collections.py:132:9: F405 `TwoPhotonFrequencySampler` may be undefined, or defined from star imports
tardis/plasma/tests/test_nlte_solver.py:1:1: I001 [*] Import block is un-sorted or un-formatted
tardis/plasma/tests/test_nlte_solver.py:422:1: W293 [*] Blank line contains whitespace
tardis/plasma/tests/test_nlte_solver.py:424:20: W291 [*] Trailing whitespace
tardis/plasma/tests/test_nlte_solver.py:428:1: W293 [*] Blank line contains whitespace
tardis/plasma/tests/test_nlte_solver.py:437:1: W293 [*] Blank line contains whitespace
tardis/plasma/tests/test_nlte_solver.py:449:1: W293 [*] Blank line contains whitespace
tardis/plasma/tests/test_nlte_solver.py:458:1: W293 [*] Blank line contains whitespace
tardis/plasma/tests/test_nlte_solver.py:470:1: W293 [*] Blank line contains whitespace
tardis/plasma/tests/test_nlte_solver.py:473:15: W291 [*] Trailing whitespace
tardis/plasma/tests/test_nlte_solver.py:477:1: W293 [*] Blank line contains whitespace
tardis/plasma/tests/test_nlte_solver.py:486:1: W293 [*] Blank line contains whitespace
Found 108 errors.
[*] 14 fixable with the `--fix` option.

Aaryan-Dadu · 2026-02-21T05:25:47Z

@andrewfullard please have a look at it

wkerzendorf · 2026-02-23T14:22:34Z

tardis/plasma/properties/nlte_rate_equation_solver.py

+        and previous_electron_densities is not None
+    ):
+        first_guess = pd.Series(0.0, index=rate_matrix_index)
+        ion_indices = rate_matrix_index.drop("n_e")


why are you dropping n_e and what is that?

Answered in #3454 (comment)

Aaryan-Dadu · 2026-02-23T15:07:12Z

tardis/plasma/properties/nlte_rate_equation_solver.py

        )

        # TODO: Don't create the rate_matrix_index with n_e in the first place
        rate_matrix_index = rate_matrix_index.drop("n_e")


I noticed this line so be on the safer side I added the line to drop n_e which is a label for electron density and also ion indices dont need these electron density labels same as how in line 108 for atomic numbers it was dropped because atomic numbers are of element and electrons are not atoms similarly ions should not have n_e as its not an ion state

but for ionization we need electron density, no?

but this is label n_e, we are using the electron_density argument for the calculations

wkerzendorf · 2026-02-23T19:03:54Z

please go through the PR requirements for TARDIS and follow them.

Aaryan-Dadu · 2026-02-23T19:47:57Z

please go through the PR requirements for TARDIS and follow them.

Added the AI Usage statement as per the PR Checklist
Please tell me if there is anything else that needs to change.

atharva-2001 · 2026-02-24T02:32:58Z

Can you check it once again please.

tardis-bot added the GSoC 🌞 label Feb 20, 2026

Aaryan-Dadu added 2 commits February 21, 2026 03:36

use previous iteration ion populations as first guess in NLTE solver

d91a5fb

- Implement PreviousIonNumberDensity - update NLTEPopulationSolverRoot to utilize previous iteration values

tests for changes in NLTE solver

bf8d951

Aaryan-Dadu force-pushed the feat-2485 branch from 1b67707 to bf8d951 Compare February 20, 2026 22:06

wkerzendorf reviewed Feb 23, 2026

View reviewed changes

Aaryan-Dadu commented Feb 23, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Comments

[GSOC] use previous iteration ion populations as first guess in NLTE solver#3454

[GSOC] use previous iteration ion populations as first guess in NLTE solver#3454
Aaryan-Dadu wants to merge 2 commits intotardis-sn:masterfrom
Aaryan-Dadu:feat-2485

Aaryan-Dadu commented Feb 20, 2026 •

edited

Loading

Uh oh!

tardis-bot commented Feb 20, 2026 •

edited

Loading

Uh oh!

Aaryan-Dadu commented Feb 21, 2026

Uh oh!

wkerzendorf Feb 23, 2026

Uh oh!

Aaryan-Dadu Feb 23, 2026

Uh oh!

Aaryan-Dadu Feb 23, 2026 •

edited

Loading

Uh oh!

wkerzendorf Feb 23, 2026

Uh oh!

Aaryan-Dadu Feb 23, 2026

Uh oh!

wkerzendorf commented Feb 23, 2026

Uh oh!

Aaryan-Dadu commented Feb 23, 2026

Uh oh!

atharva-2001 commented Feb 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Comments

Conversation

Aaryan-Dadu commented Feb 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📝 Description

📌 Resources

🚦 Testing

☑️ Checklist

AI Usage

Uh oh!

tardis-bot commented Feb 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Aaryan-Dadu commented Feb 21, 2026

Uh oh!

wkerzendorf Feb 23, 2026

Choose a reason for hiding this comment

Uh oh!

Aaryan-Dadu Feb 23, 2026

Choose a reason for hiding this comment

Uh oh!

Aaryan-Dadu Feb 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wkerzendorf Feb 23, 2026

Choose a reason for hiding this comment

Uh oh!

Aaryan-Dadu Feb 23, 2026

Choose a reason for hiding this comment

Uh oh!

wkerzendorf commented Feb 23, 2026

Uh oh!

Aaryan-Dadu commented Feb 23, 2026

Uh oh!

atharva-2001 commented Feb 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Aaryan-Dadu commented Feb 20, 2026 •

edited

Loading

tardis-bot commented Feb 20, 2026 •

edited

Loading

Aaryan-Dadu Feb 23, 2026 •

edited

Loading