From 91fb18c81f83b2d6b4744387e6194e2d83a680d7 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 10:35:08 +0000
Subject: [PATCH 001/150] Update inputmodel_misc.py

---
 artistools/inputmodel/inputmodel_misc.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/artistools/inputmodel/inputmodel_misc.py b/artistools/inputmodel/inputmodel_misc.py
index 663f04eff..4a83143ee 100644
--- a/artistools/inputmodel/inputmodel_misc.py
+++ b/artistools/inputmodel/inputmodel_misc.py
@@ -337,7 +337,7 @@ def get_modeldata_polars(
         print(f"{filename} has been modified after {filenameparquet}. Deleting out of date parquet file.")
         filenameparquet.unlink()
 
-    dfmodel: pl.LazyFrame | None | pl.DataFrame
+    dfmodel: pl.LazyFrame | None | pl.DataFrame = None
     if not getheadersonly and filenameparquet.is_file():
         if not printwarningsonly:
             print(f"Reading data table from {filenameparquet}")
@@ -346,22 +346,20 @@ def get_modeldata_polars(
         except pl.exceptions.ComputeError:
             print(f"Problem reading {filenameparquet}. Will regenerate and overwite from text source.")
             dfmodel = None
-    else:
-        dfmodel = None
 
     if dfmodel is not None:
         # already read in the model data, just need to get the metadata
         getheadersonly = True
 
-    dfmodel_in, modelmeta = read_modelfile_text(
+    dfmodel_textfile, modelmeta = read_modelfile_text(
         filename=filename,
         printwarningsonly=printwarningsonly,
         getheadersonly=getheadersonly,
     )
 
     if dfmodel is None:
-        dfmodel = dfmodel_in
-    elif dfmodel.schema != dfmodel_in.schema:
+        dfmodel = dfmodel_textfile
+    elif dfmodel.schema != dfmodel_textfile.schema:
         print(f"ERROR: parquet schema does not match model.txt. Remove {filenameparquet} and try again.")
         raise AssertionError
 

From cdfc4adc1a72e175244cd38948b56dd2b03ed32c Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 10:35:19 +0000
Subject: [PATCH 002/150] Update modelfromhydro.py

---
 artistools/inputmodel/modelfromhydro.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/artistools/inputmodel/modelfromhydro.py b/artistools/inputmodel/modelfromhydro.py
index ff28bbbcd..5e78a211c 100755
--- a/artistools/inputmodel/modelfromhydro.py
+++ b/artistools/inputmodel/modelfromhydro.py
@@ -436,7 +436,6 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         argcomplete.autocomplete(parser)
         args = parser.parse_args(argsraw)
 
-    pd.options.mode.copy_on_write = True
     gridfolderpath = args.gridfolderpath
     if not Path(gridfolderpath, "grid.dat").is_file() or not Path(gridfolderpath, "gridcontributions.txt").is_file():
         print("grid.dat and gridcontributions.txt are required. Run artistools-maptogrid")

From 48d6262dd9cce7a71b3b9a6c137c2cc2bf198c59 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 12:43:44 +0000
Subject: [PATCH 003/150] Update estimators.py

---
 artistools/estimators/estimators.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index fae81ec88..0143c0a44 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -213,7 +213,7 @@ def parse_estimfile(
 
             elif row[0] == "heating:" and get_heatingcooling:
                 for heatingtype, value in zip(row[1::2], row[2::2]):
-                    key = heatingtype if heatingtype.startswith("heating_") else "heating_" + heatingtype
+                    key = heatingtype if heatingtype.startswith("heating_") else f"heating_{heatingtype}"
                     estimblock[key] = float(value)
 
                 if "heating_gamma/gamma_dep" in estimblock and estimblock["heating_gamma/gamma_dep"] > 0:
@@ -223,7 +223,7 @@ def parse_estimfile(
 
             elif row[0] == "cooling:" and get_heatingcooling:
                 for coolingtype, value in zip(row[1::2], row[2::2]):
-                    estimblock["cooling_" + coolingtype] = float(value)
+                    estimblock[f"cooling_{coolingtype}"] = float(value)
 
     # reached the end of file
     if timestep >= 0 and modelgridindex >= 0 and (not skip_emptycells or not estimblock.get("emptycell", True)):

From 1a885984ade7f14c14b4592f6ca054b534f08aa6 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 13:11:12 +0000
Subject: [PATCH 004/150] Update estimators.py

---
 artistools/estimators/estimators.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 0143c0a44..f254dc6af 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -168,6 +168,7 @@ def parse_estimfile(
                 else:
                     atomic_number = int(row[1].split("=")[1])
                     startindex = 2
+                elsymbol = at.get_elsymbol(atomic_number)
 
                 estimblock.setdefault(variablename, {})
 
@@ -184,10 +185,10 @@ def parse_estimfile(
                     try:
                         ion_stage = int(ion_stage_str.rstrip(":"))
                     except ValueError:
-                        if variablename == "populations" and ion_stage_str.startswith(at.get_elsymbol(atomic_number)):
+                        if variablename == "populations" and ion_stage_str.startswith(elsymbol):
                             estimblock[variablename][ion_stage_str.rstrip(":")] = float(value)
                         else:
-                            print(ion_stage_str, at.get_elsymbol(atomic_number))
+                            print(ion_stage_str, elsymbol)
                             print(f"Cannot parse row: {row}")
                         continue
 

From b757d3d5413b45573ba9c98f98a44f8e0592e071 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 13:14:53 +0000
Subject: [PATCH 005/150] Update estimators.py

---
 artistools/estimators/estimators.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index f254dc6af..c742a0867 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -173,12 +173,13 @@ def parse_estimfile(
                 estimblock.setdefault(variablename, {})
 
                 for ion_stage_str, value in zip(row[startindex::2], row[startindex + 1 :: 2]):
-                    if ion_stage_str.strip() == "(or":
+                    ion_stage_str_strip = ion_stage_str.strip()
+                    if ion_stage_str_strip == "(or":
                         continue
 
                     value_thision = float(value.rstrip(","))
 
-                    if ion_stage_str.strip() == "SUM:":
+                    if ion_stage_str_strip == "SUM:":
                         estimblock[variablename][atomic_number] = value_thision
                         continue
 

From 0920c4de964ae6e3d06327a101f2216c1c8ddfb0 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 14:13:09 +0000
Subject: [PATCH 006/150] Flatten estimator dictionary (not nested dicts)

---
 artistools/codecomparison.py            |  8 +--
 artistools/estimators/estimators.py     | 92 ++++++++++---------------
 artistools/estimators/plotestimators.py | 73 ++++++++------------
 artistools/transitions.py               | 17 +++--
 pyproject.toml                          |  6 +-
 5 files changed, 81 insertions(+), 115 deletions(-)

diff --git a/artistools/codecomparison.py b/artistools/codecomparison.py
index 76040c73c..37e150612 100644
--- a/artistools/codecomparison.py
+++ b/artistools/codecomparison.py
@@ -148,8 +148,6 @@ def read_reference_estimators(
                     cur_modelgridindex += 1
 
                     tsmgi = (cur_timestep, cur_modelgridindex)
-                    if "populations" not in estimators[tsmgi]:
-                        estimators[tsmgi]["populations"] = {}
 
                     assert len(row) == nstages + 1
                     assert len(iontuples) == nstages
@@ -158,9 +156,9 @@ def read_reference_estimators(
                             ionfrac = float(strionfrac)
                             ionpop = ionfrac * estimators[tsmgi]["nntot"]
                             if ionpop > 1e-80:
-                                estimators[tsmgi]["populations"][(atomic_number, ion_stage)] = ionpop
-                                estimators[tsmgi]["populations"].setdefault(atomic_number, 0.0)
-                                estimators[tsmgi]["populations"][atomic_number] += ionpop
+                                estimators[tsmgi][f"populations_{atomic_number}_{ion_stage}"] = ionpop
+                                estimators[tsmgi].setdefault(f"populations_{atomic_number}", 0.0)
+                                estimators[tsmgi][f"populations_{atomic_number}"] += ionpop
 
                         except ValueError:
                             estimators[tsmgi]["populations"][(atomic_number, ion_stage)] = float("NaN")
diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index c742a0867..802af3737 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -12,7 +12,6 @@
 import typing as t
 from collections import namedtuple
 from functools import partial
-from functools import reduce
 from pathlib import Path
 
 import numpy as np
@@ -124,12 +123,12 @@ def parse_estimfile(
     get_ion_values: bool = True,
     get_heatingcooling: bool = True,
     skip_emptycells: bool = False,
-) -> t.Iterator[tuple[int, int, dict]]:  # pylint: disable=unused-argument
+) -> t.Iterator[tuple[int, int, dict[str, t.Any]]]:  # pylint: disable=unused-argument
     """Generate timestep, modelgridindex, dict from estimator file."""
     with at.zopen(estfilepath) as estimfile:
         timestep: int = -1
         modelgridindex: int = -1
-        estimblock: dict[t.Any, t.Any] = {}
+        estimblock: dict[str, t.Any] = {}
         for line in estimfile:
             row: list[str] = line.split()
             if not row:
@@ -170,8 +169,6 @@ def parse_estimfile(
                     startindex = 2
                 elsymbol = at.get_elsymbol(atomic_number)
 
-                estimblock.setdefault(variablename, {})
-
                 for ion_stage_str, value in zip(row[startindex::2], row[startindex + 1 :: 2]):
                     ion_stage_str_strip = ion_stage_str.strip()
                     if ion_stage_str_strip == "(or":
@@ -180,38 +177,37 @@ def parse_estimfile(
                     value_thision = float(value.rstrip(","))
 
                     if ion_stage_str_strip == "SUM:":
-                        estimblock[variablename][atomic_number] = value_thision
+                        estimblock[f"{variablename}_{atomic_number}"] = value_thision
                         continue
 
                     try:
                         ion_stage = int(ion_stage_str.rstrip(":"))
                     except ValueError:
                         if variablename == "populations" and ion_stage_str.startswith(elsymbol):
-                            estimblock[variablename][ion_stage_str.rstrip(":")] = float(value)
+                            estimblock[f"{variablename}_{ion_stage_str.rstrip(':')}"] = float(value)
                         else:
                             print(ion_stage_str, elsymbol)
                             print(f"Cannot parse row: {row}")
                         continue
 
-                    estimblock[variablename][(atomic_number, ion_stage)] = value_thision
+                    estimblock[f"{variablename}_{atomic_number}_{ion_stage}"] = value_thision
 
                     if variablename in {"Alpha_R*nne", "AlphaR*nne"}:
-                        estimblock.setdefault("Alpha_R", {})
-                        estimblock["Alpha_R"][(atomic_number, ion_stage)] = (
+                        estimblock[f"Alpha_R_{atomic_number}_{ion_stage}"] = (
                             value_thision / estimblock["nne"] if estimblock["nne"] > 0.0 else float("inf")
                         )
 
                     else:  # variablename == 'populations':
                         # contribute the ion population to the element population
-                        estimblock[variablename].setdefault(atomic_number, 0.0)
-                        estimblock[variablename][atomic_number] += value_thision
+                        estimblock.setdefault(f"{variablename}_{atomic_number}", 0.0)
+                        estimblock[f"{variablename}_{atomic_number}"] += value_thision
 
                 if variablename == "populations":
                     # contribute the element population to the total population
-                    estimblock["populations"].setdefault("total", 0.0)
-                    estimblock["populations"]["total"] += estimblock["populations"][atomic_number]
+                    estimblock.setdefault("populations_total", 0.0)
+                    estimblock["populations_total"] += estimblock[f"populations_{atomic_number}"]
                     estimblock.setdefault("nntot", 0.0)
-                    estimblock["nntot"] += estimblock["populations"][atomic_number]
+                    estimblock["nntot"] += estimblock[f"populations_{atomic_number}"]
 
             elif row[0] == "heating:" and get_heatingcooling:
                 for heatingtype, value in zip(row[1::2], row[2::2]):
@@ -240,7 +236,7 @@ def read_estimators_from_file(
     get_ion_values: bool = True,
     get_heatingcooling: bool = True,
     skip_emptycells: bool = False,
-) -> dict[tuple[int, int], t.Any]:
+) -> dict[tuple[int, int], dict[str, t.Any]]:
     estimfilename = f"estimators_{mpirank:04d}.out"
     try:
         estfilepath = at.firstexisting(estimfilename, folder=folderpath, tryzipped=True)
@@ -275,7 +271,7 @@ def read_estimators(
     get_heatingcooling: bool = True,
     skip_emptycells: bool = False,
     add_velocity: bool = True,
-) -> dict[tuple[int, int], dict]:
+) -> dict[tuple[int, int], dict[str, t.Any]]:
     """Read estimator files into a nested dictionary structure.
 
     Speed it up by only retrieving estimators for a particular timestep(s) or modelgrid cells.
@@ -364,62 +360,48 @@ def get_averaged_estimators(
     estimators: dict[tuple[int, int], dict],
     timesteps: int | t.Sequence[int],
     modelgridindex: int,
-    keys: str | list,
+    keys: str | list | None,
     avgadjcells: int = 0,
-) -> dict | float:
+) -> dict[str, t.Any]:
     """Get the average of estimators[(timestep, modelgridindex)][keys[0]]...[keys[-1]] across timesteps."""
-    if isinstance(keys, str):
-        keys = [keys]
     modelgridindex = int(modelgridindex)
-
-    # reduce(lambda d, k: d[k], keys, dictionary) returns dictionary[keys[0]][keys[1]]...[keys[-1]]
-    # applying all keys in the keys list
-
-    # if single timestep, no averaging needed
     if isinstance(timesteps, int):
-        return reduce(lambda d, k: d[k], [(timesteps, modelgridindex), *keys], estimators)
+        timesteps = [timesteps]
 
-    firsttimestepvalue = reduce(lambda d, k: d[k], [(timesteps[0], modelgridindex), *keys], estimators)
-    if isinstance(firsttimestepvalue, dict):
-        return {
-            k: get_averaged_estimators(modelpath, estimators, timesteps, modelgridindex, [*keys, k])
-            for k in firsttimestepvalue
-        }
+    if isinstance(keys, str):
+        keys = [keys]
+    elif keys is None or not keys:
+        keys = list(estimators[(timesteps[0], modelgridindex)].keys())
 
+    dictout = {}
     tdeltas = at.get_timestep_times(modelpath, loc="delta")
-    valuesum = 0
-    tdeltasum = 0
-    for timestep, tdelta in zip(timesteps, tdeltas):
-        for mgi in range(modelgridindex - avgadjcells, modelgridindex + avgadjcells + 1):
-            with contextlib.suppress(KeyError):
-                valuesum += reduce(lambda d, k: d[k], [(timestep, mgi), *keys], estimators) * tdelta
+    for k in keys:
+        valuesum = 0
+        tdeltasum = 0
+        for timestep, tdelta in zip(timesteps, tdeltas):
+            for mgi in range(modelgridindex - avgadjcells, modelgridindex + avgadjcells + 1):
+                valuesum += estimators[(timestep, mgi)][k] * tdelta
                 tdeltasum += tdelta
-    return valuesum / tdeltasum
+        dictout[k] = valuesum / tdeltasum
 
-    # except KeyError:
-    #     if (timestep, modelgridindex) in estimators:
-    #         print(f'Unknown x variable: {xvariable} for timestep {timestep} in cell {modelgridindex}')
-    #     else:
-    #         print(f'No data for cell {modelgridindex} at timestep {timestep}')
-    #     print(estimators[(timestep, modelgridindex)])
-    #     sys.exit()
+    return dictout
 
 
-def get_averageionisation(populations: dict[t.Any, float], atomic_number: int) -> float:
+def get_averageionisation(estimatorstsmgi: dict[str, float], atomic_number: int) -> float:
     free_electron_weighted_pop_sum = 0.0
     found = False
     popsum = 0.0
-    for key in populations:
-        if isinstance(key, tuple) and key[0] == atomic_number:
+    for key in estimatorstsmgi:
+        if key.startswith(f"populations_{atomic_number}_"):
             found = True
-            ion_stage = key[1]
-            free_electron_weighted_pop_sum += populations[key] * (ion_stage - 1)
-            popsum += populations[key]
+            ion_stage = int(key.removeprefix(f"populations_{atomic_number}_"))
+            free_electron_weighted_pop_sum += estimatorstsmgi[key] * (ion_stage - 1)
+            popsum += estimatorstsmgi[key]
 
     if not found:
         return float("NaN")
 
-    return free_electron_weighted_pop_sum / populations[atomic_number]
+    return free_electron_weighted_pop_sum / estimatorstsmgi[f"populations_{atomic_number}"]
 
 
 def get_averageexcitation(
@@ -459,7 +441,7 @@ def get_averageexcitation(
     return energypopsum / ionpopsum
 
 
-def get_partiallycompletetimesteps(estimators: dict[tuple[int, int], dict]) -> list[int]:
+def get_partiallycompletetimesteps(estimators: dict[tuple[int, int], dict[str, t.Any]]) -> list[int]:
     """During a simulation, some estimator files can contain information for some cells but not others
     for the current timestep.
     """
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 4a10bf4fc..897f0ad89 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -8,7 +8,6 @@
 import argparse
 import contextlib
 import math
-import sys
 import typing as t
 from pathlib import Path
 
@@ -150,9 +149,7 @@ def plot_average_ionisation_excitation(
             for timestep in timesteps:
                 if seriestype == "averageionisation":
                     valuesum += (
-                        at.estimators.get_averageionisation(
-                            estimators[(timestep, modelgridindex)]["populations"], atomic_number
-                        )
+                        at.estimators.get_averageionisation(estimators[(timestep, modelgridindex)], atomic_number)
                         * arr_tdelta[timestep]
                     )
                 elif seriestype == "averageexcitation":
@@ -360,29 +357,37 @@ def get_iontuple(ionstr):
                 #         f"cell {modelgridindex} timesteps {timesteps}"
                 #     )
 
+                if ion_stage == "ALL":
+                    key = f"populations_{atomic_number}"
+                elif hasattr(ion_stage, "lower") and ion_stage.startswith(at.get_elsymbol(atomic_number)):
+                    # not really an ionstage but maybe isotope?
+                    key = f"populations_{ion_stage}"
+                else:
+                    key = f"populations_{atomic_number}_{ion_stage}"
+
                 try:
                     estimpop = at.estimators.get_averaged_estimators(
-                        modelpath, estimators, timesteps, modelgridindex, ["populations"]
+                        modelpath,
+                        estimators,
+                        timesteps,
+                        modelgridindex,
+                        [key, f"populations_{atomic_number}", "populations_total"],
                     )
                 except KeyError:
+                    print(f"KeyError: {key} not in estimators")
                     ylist.append(float("nan"))
                     continue
 
-                if ion_stage == "ALL":
-                    nionpop = estimpop.get((atomic_number), 0.0)
-                elif hasattr(ion_stage, "lower") and ion_stage.startswith(at.get_elsymbol(atomic_number)):
-                    nionpop = estimpop.get(ion_stage, 0.0)
-                else:
-                    nionpop = estimpop.get((atomic_number, ion_stage), 0.0)
+                nionpop = estimpop.get(key, 0.0)
 
                 try:
                     if args.ionpoptype == "absolute":
                         yvalue = nionpop  # Plot as fraction of element population
                     elif args.ionpoptype == "elpop":
-                        elpop = estimpop.get(atomic_number, 0.0)
+                        elpop = estimpop.get(f"populations_{atomic_number}", 0.0)
                         yvalue = nionpop / elpop  # Plot as fraction of element population
                     elif args.ionpoptype == "totalpop":
-                        totalpop = estimpop["total"]
+                        totalpop = estimpop["populations_total"]
                         yvalue = nionpop / totalpop  # Plot as fraction of total population
                     else:
                         raise AssertionError
@@ -391,32 +396,11 @@ def get_iontuple(ionstr):
 
                 ylist.append(yvalue)
 
-            # elif seriestype == 'Alpha_R':
-            #     ylist.append(estim['Alpha_R*nne'].get((atomic_number, ion_stage), 0.) / estim['nne'])
-            # else:
-            #     ylist.append(estim[seriestype].get((atomic_number, ion_stage), 0.))
             else:
-                # this is very slow!
-                try:
-                    estim = at.estimators.get_averaged_estimators(modelpath, estimators, timesteps, modelgridindex, [])
-                except KeyError:
-                    ylist.append(float("nan"))
-                    continue
-
-                dictvars = {}
-                for k, value in estim.items():
-                    if isinstance(value, dict):
-                        dictvars[k] = value.get((atomic_number, ion_stage), 0.0)
-                    else:
-                        dictvars[k] = value
-
-                # dictvars will now define things like 'Te', 'TR',
-                # as well as 'populations' which applies to the current ion
-
-                try:
-                    yvalue = eval(seriestype, {"__builtins__": math}, dictvars)
-                except ZeroDivisionError:
-                    yvalue = float("NaN")
+                key = f"{seriestype}_{atomic_number}_{ion_stage}"
+                yvalue = at.estimators.get_averaged_estimators(modelpath, estimators, timesteps, modelgridindex, key)[
+                    key
+                ]
                 ylist.append(yvalue)
 
         plotlabel = (
@@ -511,16 +495,17 @@ def plot_series(
 
     ylist: list[float] = []
     for modelgridindex, timesteps in zip(mgilist, timestepslist):
-        estimavg = at.estimators.get_averaged_estimators(modelpath, estimators, timesteps, modelgridindex, [])
-        assert isinstance(estimavg, dict)
         try:
-            ylist.append(eval(variablename, {"__builtins__": math}, estimavg))
+            yvalue = at.estimators.get_averaged_estimators(
+                modelpath, estimators, timesteps, modelgridindex, variablename
+            )[variablename]
+            ylist.append(yvalue)
         except KeyError:
             if (timesteps[0], modelgridindex) in estimators:
                 print(f"Undefined variable: {variablename} in cell {modelgridindex}")
             else:
                 print(f"No data for cell {modelgridindex}")
-            sys.exit()
+            raise
 
     try:
         if math.log10(max(ylist) / min(ylist)) > 2:
@@ -598,7 +583,9 @@ def get_xlist(
         mgilist_out = []
         timestepslist_out = []
         for modelgridindex, timesteps in zip(allnonemptymgilist, timestepslist):
-            xvalue = at.estimators.get_averaged_estimators(modelpath, estimators, timesteps, modelgridindex, xvariable)
+            xvalue = at.estimators.get_averaged_estimators(modelpath, estimators, timesteps, modelgridindex, xvariable)[
+                xvariable
+            ]
             assert isinstance(xvalue, float | int)
             xlist.append(xvalue)
             mgilist_out.append(modelgridindex)
diff --git a/artistools/transitions.py b/artistools/transitions.py
index d36efbbb3..fac24a1a7 100755
--- a/artistools/transitions.py
+++ b/artistools/transitions.py
@@ -126,8 +126,7 @@ def make_plot(
     peak_y_value = -1
     yvalues_combined = np.zeros((len(temperature_list), len(xvalues)))
     for seriesindex, temperature in enumerate(temperature_list):
-        T_exc = eval(temperature, vardict)
-        serieslabel = "NLTE" if T_exc < 0 else f"LTE {temperature} = {T_exc:.0f} K"
+        serieslabel = "NLTE" if temperature == "NOTEMPNLTE" else f"LTE {temperature} = {vardict[temperature]:.0f} K"
         for ion_index, axis in enumerate(axes[: len(ionlist)]):
             # an ion subplot
             yvalues_combined[seriesindex] += yvalues[seriesindex][ion_index]
@@ -326,8 +325,8 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             figure_title += f" ({time_days:.1f}d)"
 
         # -1 means use NLTE populations
-        temperature_list = ["Te", "TR", "-1"]
-        temperature_list = ["-1"]
+        temperature_list = ["Te", "TR", "NOTEMPNLTE"]
+        temperature_list = ["NOTEMPNLTE"]
         vardict = {"Te": Te, "TR": TR}
     else:
         if not args.T:
@@ -408,8 +407,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             )
 
             for seriesindex, temperature in enumerate(temperature_list):
-                T_exc = eval(temperature, vardict)
-                if T_exc < 0:
+                if temperature == "NOTEMPNLTE":
                     dfnltepops_thision = dfnltepops.query("Z==@ion.Z & ion_stage==@ion.ion_stage")
 
                     nltepopdict = {x.level: x["n_NLTE"] for _, x in dfnltepops_thision.iterrows()}
@@ -433,6 +431,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
                     with pd.option_context("display.width", 200):
                         print(dftransitions.nlargest(1, "flux_factor_nlte"))
                 else:
+                    T_exc = vardict[temperature]
                     popcolumnname = f"upper_pop_lte_{T_exc:.0f}K"
                     if args.atomicdatabase == "artis":
                         dftransitions = dftransitions.eval("upper_g = @ion.levels.loc[upper].g.to_numpy()")
@@ -467,7 +466,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
 
         def get_strionfracs(atomic_number, ionstages):
             est_ionfracs = [
-                estimators["populations"][(atomic_number, ionstage)] / estimators["populations"][atomic_number]
+                estimators[f"populations_{atomic_number}_{ionstage}"] / estimators[f"populations_{atomic_number}"]
                 for ionstage in ionstages
             ]
             ionfracs_str = " ".join([f"{pop:6.0e}" if pop < 0.01 else f"{pop:6.2f}" for pop in est_ionfracs])
@@ -489,8 +488,8 @@ def get_strionfracs(atomic_number, ionstages):
             f"{velocity:5.0f} km/s({modelgridindex})      {fe2depcoeff:5.2f}                   "
             f"{ni2depcoeff:.2f}        "
             f"{est_fe_ionfracs_str}   /  {est_ni_ionfracs_str}      {Te:.0f}    "
-            f"{estimators['populations'][(26, 3)] / estimators['populations'][(26, 2)]:.2f}          "
-            f"{estimators['populations'][(28, 3)] / estimators['populations'][(28, 2)]:5.2f}"
+            f"{estimators['populations_26_3'] / estimators['populations_26_2']:.2f}          "
+            f"{estimators['populations_28_3'] / estimators['populations_28_2']:5.2f}"
         )
 
     outputfilename = (
diff --git a/pyproject.toml b/pyproject.toml
index 43528e15f..4027873ae 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -111,7 +111,7 @@ ignore_errors = true
 max-line-length = 120
 disable = """
     broad-exception-caught,
-    eval-used,
+    #eval-used,
     fixme,
     missing-function-docstring,
     missing-module-docstring,
@@ -176,7 +176,7 @@ ignore = [
     "N999",  # invalid-module-name
     "PD901",  # df is a bad variable name
     "PERF203", # try-except-in-loop
-    "PGH001",  # No builtin `eval()` allowed
+    #"PGH001",  # No builtin `eval()` allowed
     "PLR0911",  # too-many-return-statements
     "PLR0912",  # too-many-branches
     "PLR0913",  # too-many-arguments
@@ -186,7 +186,7 @@ ignore = [
     "PYI024", # Use `typing.NamedTuple` instead of `collections.namedtuple`
     "S101",  # Use of assert detected
     "S301",  # suspicious-pickle-usage
-    "S307",  # suspicious-eval-usage
+    #"S307",  # suspicious-eval-usage
     "S311",  # suspicious-non-cryptographic-random-usage
     "S603",  # subprocess-without-shell-equals-true
     "S607",  # start-process-with-partial-path

From 0de4d2c66076d5eee9fc3fd20c350de5581221fa Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 14:28:33 +0000
Subject: [PATCH 007/150] Refactor estimator reading func params

---
 artistools/__init__.py              |  1 +
 artistools/estimators/estimators.py | 43 +++++++++++++----------------
 2 files changed, 20 insertions(+), 24 deletions(-)

diff --git a/artistools/__init__.py b/artistools/__init__.py
index c14037d79..c367af2f8 100644
--- a/artistools/__init__.py
+++ b/artistools/__init__.py
@@ -27,6 +27,7 @@
 from artistools.__main__ import main
 from artistools.configuration import get_config
 from artistools.configuration import set_config
+from artistools.estimators import read_estimators
 from artistools.inputmodel import add_derived_cols_to_modeldata
 from artistools.inputmodel import get_cell_angle
 from artistools.inputmodel import get_dfmodel_dimensions
diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 802af3737..385cbd63b 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -118,8 +118,7 @@ def get_units_string(variable: str) -> str:
 
 
 def parse_estimfile(
-    estfilepath: Path,
-    modelpath: Path,
+    estfilepath: Path | str,
     get_ion_values: bool = True,
     get_heatingcooling: bool = True,
     skip_emptycells: bool = False,
@@ -229,31 +228,22 @@ def parse_estimfile(
 
 
 def read_estimators_from_file(
-    folderpath: Path | str,
+    estfilepath: Path | str,
     modelpath: Path,
-    mpirank: int,
     printfilename: bool = False,
     get_ion_values: bool = True,
     get_heatingcooling: bool = True,
     skip_emptycells: bool = False,
 ) -> dict[tuple[int, int], dict[str, t.Any]]:
-    estimfilename = f"estimators_{mpirank:04d}.out"
-    try:
-        estfilepath = at.firstexisting(estimfilename, folder=folderpath, tryzipped=True)
-    except FileNotFoundError:
-        # not worth printing an error, because ranks with no cells to update do not produce an estimator file
-        # print(f'Warning: Could not find {estfilepath.relative_to(modelpath.parent)}')
-        return {}
-
     if printfilename:
-        filesize = Path(estfilepath).stat().st_size / 1024 / 1024
+        estfilepath = Path(estfilepath)
+        filesize = estfilepath.stat().st_size / 1024 / 1024
         print(f"Reading {estfilepath.relative_to(modelpath.parent)} ({filesize:.2f} MiB)")
 
     return {
-        (fileblock_timestep, fileblock_modelgridindex): file_estimblock
-        for fileblock_timestep, fileblock_modelgridindex, file_estimblock in parse_estimfile(
+        (timestep, mgi): file_estimblock
+        for timestep, mgi, file_estimblock in parse_estimfile(
             estfilepath,
-            modelpath,
             get_ion_values=get_ion_values,
             get_heatingcooling=get_heatingcooling,
             skip_emptycells=skip_emptycells,
@@ -313,15 +303,21 @@ def read_estimators(
 
     estimators: dict[tuple[int, int], dict] = {}
     for folderpath in runfolders:
+        estfilepaths = []
+        for mpirank in mpiranklist:
+            # not worth printing an error, because ranks with no cells to update do not produce an estimator file
+            with contextlib.suppress(FileNotFoundError):
+                estfilepath = at.firstexisting(f"estimators_{mpirank:04d}.out", folder=folderpath, tryzipped=True)
+                estfilepaths.append(estfilepath)
+
         if not printfilename:
             print(
-                f"Reading {len(list(mpiranklist))} estimator files in {folderpath.relative_to(Path(modelpath).parent)}"
+                f"Reading {len(list(estfilepaths))} estimator files in {folderpath.relative_to(Path(modelpath).parent)}"
             )
 
         processfile = partial(
             read_estimators_from_file,
-            folderpath,
-            modelpath,
+            modelpath=modelpath,
             get_ion_values=get_ion_values,
             get_heatingcooling=get_heatingcooling,
             printfilename=printfilename,
@@ -330,22 +326,21 @@ def read_estimators(
 
         if at.get_config()["num_processes"] > 1:
             with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
-                arr_rankestimators = pool.map(processfile, mpiranklist)
+                arr_rankestimators = pool.map(processfile, estfilepaths)
                 pool.close()
                 pool.join()
                 pool.terminate()
         else:
-            arr_rankestimators = [processfile(rank) for rank in mpiranklist]
+            arr_rankestimators = [processfile(estfilepath) for estfilepath in estfilepaths]
 
-        for mpirank, estimators_thisfile in zip(mpiranklist, arr_rankestimators):
+        for estfilepath, estimators_thisfile in zip(estfilepaths, arr_rankestimators):
             dupekeys = sorted([k for k in estimators_thisfile if k in estimators])
             for k in dupekeys:
                 # dropping the lowest timestep is normal for restarts. Only warn about other cases
                 if k[0] != dupekeys[0][0]:
-                    filepath = Path(folderpath, f"estimators_{mpirank:04d}.out")
                     print(
                         f"WARNING: Duplicate estimator block for (timestep, mgi) key {k}. "
-                        f"Dropping block from {filepath}"
+                        f"Dropping block from {estfilepath}"
                     )
 
                 del estimators_thisfile[k]

From 0a9ebbc6c55081b2d4fe3b50ca012942a8515986 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 15:10:23 +0000
Subject: [PATCH 008/150] Fix estimator key errors

---
 artistools/__init__.py                       |  1 +
 artistools/codecomparison.py                 |  2 +-
 artistools/estimators/estimators_classic.py  | 12 +++----
 artistools/estimators/exportmassfractions.py | 10 +++---
 artistools/estimators/plotestimators.py      |  2 +-
 artistools/gsinetwork.py                     | 12 ++++---
 artistools/misc.py                           | 34 ++++++++++++++++++++
 artistools/nltepops/plotnltepops.py          |  2 +-
 artistools/nonthermal/plotnonthermal.py      |  6 ++--
 artistools/nonthermal/solvespencerfanocmd.py |  8 +++--
 artistools/radfield.py                       |  6 ++--
 artistools/writecomparisondata.py            | 10 +++---
 12 files changed, 72 insertions(+), 33 deletions(-)

diff --git a/artistools/__init__.py b/artistools/__init__.py
index c367af2f8..9b59c2d0e 100644
--- a/artistools/__init__.py
+++ b/artistools/__init__.py
@@ -61,6 +61,7 @@
 from artistools.misc import get_filterfunc
 from artistools.misc import get_grid_mapping
 from artistools.misc import get_inputparams
+from artistools.misc import get_ion_tuple
 from artistools.misc import get_ionstring
 from artistools.misc import get_linelist_dataframe
 from artistools.misc import get_linelist_dict
diff --git a/artistools/codecomparison.py b/artistools/codecomparison.py
index 37e150612..701df2204 100644
--- a/artistools/codecomparison.py
+++ b/artistools/codecomparison.py
@@ -161,7 +161,7 @@ def read_reference_estimators(
                                 estimators[tsmgi][f"populations_{atomic_number}"] += ionpop
 
                         except ValueError:
-                            estimators[tsmgi]["populations"][(atomic_number, ion_stage)] = float("NaN")
+                            estimators[tsmgi][f"populations_{atomic_number}_{ion_stage}"] = float("NaN")
 
                     assert np.isclose(float(row[0]), estimators[tsmgi]["vel_mid"], rtol=0.01)
                     assert estimators[key]["vel_mid"]
diff --git a/artistools/estimators/estimators_classic.py b/artistools/estimators/estimators_classic.py
index ecaa598b4..30a92a078 100644
--- a/artistools/estimators/estimators_classic.py
+++ b/artistools/estimators/estimators_classic.py
@@ -26,8 +26,6 @@ def get_atomic_composition(modelpath: Path) -> dict[int, int]:
 
 
 def parse_ion_row_classic(row: list[str], outdict: dict[str, t.Any], atomic_composition: dict[int, int]) -> None:
-    outdict["populations"] = {}
-
     elements = atomic_composition.keys()
 
     i = 6  # skip first 6 numbers in est file. These are n, TR, Te, W, TJ, grey_depth.
@@ -35,14 +33,14 @@ def parse_ion_row_classic(row: list[str], outdict: dict[str, t.Any], atomic_comp
     for atomic_number in elements:
         for ion_stage in range(1, atomic_composition[atomic_number] + 1):
             value_thision = float(row[i])
-            outdict["populations"][(atomic_number, ion_stage)] = value_thision
+            outdict[f"populations_{atomic_number}_{ion_stage}"] = value_thision
             i += 1
 
-            elpop = outdict["populations"].get(atomic_number, 0)
-            outdict["populations"][atomic_number] = elpop + value_thision
+            elpop = outdict.get(f"populations_{atomic_number}", 0)
+            outdict[f"populations_{atomic_number}"] = elpop + value_thision
 
-            totalpop = outdict["populations"].get("total", 0)
-            outdict["populations"]["total"] = totalpop + value_thision
+            totalpop = outdict.get("populations_total", 0)
+            outdict["populations_total"] = totalpop + value_thision
 
 
 def get_first_ts_in_run_directory(modelpath) -> dict[str, int]:
diff --git a/artistools/estimators/exportmassfractions.py b/artistools/estimators/exportmassfractions.py
index 8c444becd..1538155cf 100755
--- a/artistools/estimators/exportmassfractions.py
+++ b/artistools/estimators/exportmassfractions.py
@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 import argparse
-import contextlib
 import typing as t
 from pathlib import Path
 
@@ -32,14 +31,13 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         estimators = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindexlist)
         for modelgridindex in modelgridindexlist:
             tdays = estimators[(timestep, modelgridindex)]["tdays"]
-            popdict = estimators[(timestep, modelgridindex)]["populations"]
 
             numberdens = {}
             totaldens = 0.0  # number density times atomic mass summed over all elements
-            for key in popdict:
-                with contextlib.suppress(ValueError, TypeError):
-                    atomic_number = int(key)
-                    numberdens[atomic_number] = popdict[atomic_number]
+            for key, val in estimators[(timestep, modelgridindex)].items():
+                if key.startswith("populations") and key.removeprefix("populations_").isdigit():
+                    atomic_number = int(key.removeprefix("populations_"))
+                    numberdens[atomic_number] = val
                     totaldens += numberdens[atomic_number] * elmass[atomic_number]
             massfracs = {
                 atomic_number: numberdens[atomic_number] * elmass[atomic_number] / totaldens
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 897f0ad89..6ad3111cf 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -320,7 +320,7 @@ def get_iontuple(ionstr):
         print("WARNING: Could not read an ARTIS compositiondata.txt file")
         for atomic_number, ion_stage in iontuplelist:
             mgits = (timestepslist[0][0], mgilist[0])
-            if (atomic_number, ion_stage) not in estimators[mgits]["populations"]:
+            if f"populations_{atomic_number}_{ion_stage}" not in estimators[mgits]:
                 missingions.add((atomic_number, ion_stage))
 
     if missingions:
diff --git a/artistools/gsinetwork.py b/artistools/gsinetwork.py
index a71b802ec..5d4a88271 100755
--- a/artistools/gsinetwork.py
+++ b/artistools/gsinetwork.py
@@ -501,7 +501,7 @@ def plot_qdot_abund_modelcells(
             rho_cgs = rho_init_cgs * (t_model_init_days / time_days) ** 3
 
             for strnuc, a in zip(arr_strnuc, arr_a):
-                abund = estimators[(nts, mgi)]["populations"][strnuc]
+                abund = estimators[(nts, mgi)][f"populations_{strnuc}"]
                 massfrac = abund * a * MH / rho_cgs
                 massfrac = massfrac + dfmodel.iloc[mgi][f"X_{strnuc}"] * (correction_factors[strnuc] - 1.0)
 
@@ -519,7 +519,7 @@ def plot_qdot_abund_modelcells(
             if "Ye" not in arr_abund_artis[mgi]:
                 arr_abund_artis[mgi]["Ye"] = []
 
-            abund = estimators[(nts, mgi)]["populations"].get(strnuc, 0.0)
+            abund = estimators[(nts, mgi)].get(f"populations_{strnuc}", 0.0)
             if "Ye" in estimators[(nts, mgi)]:
                 cell_Ye = estimators[(nts, mgi)]["Ye"]
                 arr_abund_artis[mgi]["Ye"].append(cell_Ye)
@@ -529,8 +529,12 @@ def plot_qdot_abund_modelcells(
                 cell_protoncount = 0.0
                 cell_nucleoncount = 0.0
                 cellvolume = dfmodel.iloc[mgi].volume
-                for popkey, abund in estimators[(nts, mgi)]["populations"].items():
-                    if isinstance(popkey, str) and abund > 0.0:
+                for popkey, abund in estimators[(nts, mgi)].items():
+                    if (
+                        popkey.startswith("populations_")
+                        and "_" not in popkey.removeprefix("populations_")
+                        and abund > 0.0
+                    ):
                         if popkey.endswith("_otherstable"):
                             # TODO: use mean molecular weight, but this is not needed for kilonova input files anyway
                             print(f"WARNING {popkey}={abund} not contributed")
diff --git a/artistools/misc.py b/artistools/misc.py
index abdba5d50..1fc1547c6 100644
--- a/artistools/misc.py
+++ b/artistools/misc.py
@@ -627,6 +627,40 @@ def get_elsymbol(atomic_number: int | np.int64) -> str:
     return get_elsymbolslist()[atomic_number]
 
 
+def get_ion_tuple(ionstr: str) -> tuple[int, int] | int:
+    """Return a tuple of the atomic number and ionisation stage such as (26,2) for an ion string like 'FeII', 'Fe II', or '26_2'.
+
+    Return the atomic number for a string like 'Fe' or '26'.
+    """
+    ionstr = ionstr.removeprefix("populations_")
+    if ionstr.isdigit():
+        return int(ionstr)
+
+    if ionstr in at.get_elsymbolslist():
+        return at.get_atomic_number(ionstr)
+
+    elem = None
+    if " " in ionstr:
+        elem, strionstage = ionstr.split(" ")
+    elif "_" in ionstr:
+        elem, strionstage = ionstr.split("_")
+    else:
+        for elsym in at.get_elsymbolslist():
+            if ionstr.startswith(elsym):
+                elem = elsym
+                strionstage = ionstr.removeprefix(elsym)
+                break
+
+    if not elem:
+        msg = f"Could not parse ionstr {ionstr}"
+        raise ValueError(msg)
+
+    atomic_number = int(elem) if elem.isdigit() else at.get_atomic_number(elem)
+    ionstage = int(strionstage) if strionstage.isdigit() else at.decode_roman_numeral(strionstage)
+
+    return (atomic_number, ionstage)
+
+
 @lru_cache(maxsize=16)
 def get_ionstring(
     atomic_number: int | np.int64,
diff --git a/artistools/nltepops/plotnltepops.py b/artistools/nltepops/plotnltepops.py
index 17843ae41..de34819f7 100755
--- a/artistools/nltepops/plotnltepops.py
+++ b/artistools/nltepops/plotnltepops.py
@@ -198,7 +198,7 @@ def make_ionsubplot(
         dfpopthision = dfpopthision.query("level <= @args.maxlevel")
 
     ionpopulation = dfpopthision["n_NLTE"].sum()
-    ionpopulation_fromest = estimators[(timestep, modelgridindex)]["populations"].get((atomic_number, ion_stage), 0.0)
+    ionpopulation_fromest = estimators[(timestep, modelgridindex)].get(f"populations_{atomic_number}_{ion_stage}", 0.0)
 
     dfpopthision["parity"] = [
         1 if (row.level != -1 and ion_data.levels.iloc[int(row.level)].levelname.split("[")[0][-1] == "o") else 0
diff --git a/artistools/nonthermal/plotnonthermal.py b/artistools/nonthermal/plotnonthermal.py
index 9bca9b4c9..0cc737d20 100755
--- a/artistools/nonthermal/plotnonthermal.py
+++ b/artistools/nonthermal/plotnonthermal.py
@@ -85,12 +85,12 @@ def plot_contributions(axis, modelpath, timestep, modelgridindex, nonthermaldata
     dfcollion = at.nonthermal.read_colliondata()
 
     elementlist = at.get_composition_data(modelpath)
-    totalpop = estimators[(timestep, modelgridindex)]["populations"]["total"]
+    totalpop = estimators[(timestep, modelgridindex)]["populations_total"]
     nelements = len(elementlist)
     for element in range(nelements):
         Z = elementlist.Z[element]
 
-        elpop = estimators[(timestep, modelgridindex)]["populations"][Z]
+        elpop = estimators[(timestep, modelgridindex)][f"populations_{Z}"]
         if elpop <= 1e-4 * totalpop:
             continue
 
@@ -100,7 +100,7 @@ def plot_contributions(axis, modelpath, timestep, modelgridindex, nonthermaldata
         nions = elementlist.nions[element]
         for ion in range(nions):
             ionstage = ion + elementlist.lowermost_ionstage[element]
-            ionpop = estimators[(timestep, modelgridindex)]["populations"][(Z, ionstage)]
+            ionpop = estimators[(timestep, modelgridindex)][f"populations_{Z}_{ionstage}"]
 
             dfcollion_thision = dfcollion.query("Z == @Z and ionstage == @ionstage")
 
diff --git a/artistools/nonthermal/solvespencerfanocmd.py b/artistools/nonthermal/solvespencerfanocmd.py
index 5f854f3b1..6644bee97 100755
--- a/artistools/nonthermal/solvespencerfanocmd.py
+++ b/artistools/nonthermal/solvespencerfanocmd.py
@@ -186,12 +186,16 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             print(f"ERROR: no NLTE populations for cell {args.modelgridindex} at timestep {args.timestep}")
             raise AssertionError
 
-        nntot = estim["populations"]["total"]
+        nntot = estim["populations_total"]
         # nne = estim["nne"]
         T_e = estim["Te"]
         print("WARNING: Use LTE pops at Te for now")
         deposition_density_ev = estim["heating_dep"] / 1.6021772e-12  # convert erg to eV
-        ionpopdict = estim["populations"]
+        ionpopdict = {
+            at.get_ion_tuple(k): v
+            for k, v in estim.items()
+            if k.startswith("populations_") and "_" in k.removeprefix("populations_")
+        }
 
         velocity = modeldata["vel_r_max_kmps"][args.modelgridindex]
         args.timedays = float(at.get_timestep_time(modelpath, args.timestep))
diff --git a/artistools/radfield.py b/artistools/radfield.py
index d70edb380..36d210c1b 100755
--- a/artistools/radfield.py
+++ b/artistools/radfield.py
@@ -315,7 +315,7 @@ def get_kappa_bf_ion(atomic_number, lower_ion_stage, modelgridindex, timestep, m
     ion_data = adata.query("Z == @atomic_number and ion_stage == @lower_ion_stage").iloc[0]
     upper_ion_data = adata.query("Z == @atomic_number and ion_stage == (@lower_ion_stage + 1)").iloc[0]
 
-    lowerionpopdensity = estimators[(timestep, modelgridindex)]["populations"][(atomic_number, lower_ion_stage)]
+    lowerionpopdensity = estimators[(timestep, modelgridindex)][f"populations_{atomic_number}_{lower_ion_stage}"]
 
     ion_popfactor_sum = sum(
         level.g * math.exp(-level.energy_ev * EV / KB / T_e) for _, level in ion_data.levels[:max_levels].iterrows()
@@ -352,11 +352,11 @@ def get_recombination_emission(
 
     estimators = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex)
 
-    upperionpopdensity = estimators[(timestep, modelgridindex)]["populations"][(atomic_number, upper_ion_stage)]
+    upperionpopdensity = estimators[(timestep, modelgridindex)][f"populations_{atomic_number}_{upper_ion_stage}"]
     T_e = estimators[(timestep, modelgridindex)]["Te"]
     nne = estimators[(timestep, modelgridindex)]["nne"]
 
-    upperionpopdensity = estimators[(timestep, modelgridindex)]["populations"][(atomic_number, upper_ion_stage)]
+    upperionpopdensity = estimators[(timestep, modelgridindex)][f"populations_{atomic_number}_{upper_ion_stage}"]
     print(f"Recombination from {upperionstr} -> {lowerionstr} ({upperionstr} pop = {upperionpopdensity:.1e}/cm3)")
 
     if use_lte_pops:
diff --git a/artistools/writecomparisondata.py b/artistools/writecomparisondata.py
index c826a293d..0b512b950 100755
--- a/artistools/writecomparisondata.py
+++ b/artistools/writecomparisondata.py
@@ -114,11 +114,11 @@ def write_ionfracts(
                         continue
                     v_mid = (cell.vel_r_min_kmps + cell.vel_r_max_kmps) / 2.0
                     f.write(f"{v_mid:.2f}")
-                    elabund = estimators[(timestep, modelgridindex)]["populations"].get(atomic_number, 0)
+                    elabund = estimators[(timestep, modelgridindex)].get(f"populations_{atomic_number}", 0)
                     for ion in range(nions):
                         ion_stage = ion + elementlist.lowermost_ionstage[element]
-                        ionabund = estimators[(timestep, modelgridindex)]["populations"].get(
-                            (atomic_number, ion_stage), 0
+                        ionabund = estimators[(timestep, modelgridindex)].get(
+                            f"populations_{atomic_number}_{ion_stage}", 0
                         )
                         ionfrac = ionabund / elabund if elabund > 0 else 0
                         if ionfrac > 0.0:
@@ -149,8 +149,8 @@ def write_phys(modelpath, model_id, selected_timesteps, estimators, allnonemptym
                     10**cell.logrho * (modelmeta["t_model_init_days"] / times[timestep]) ** 3
                 )
 
-                estimators[(timestep, modelgridindex)]["nntot"] = estimators[(timestep, modelgridindex)]["populations"][
-                    "total"
+                estimators[(timestep, modelgridindex)]["nntot"] = estimators[(timestep, modelgridindex)][
+                    "populations_total"
                 ]
 
                 v_mid = (cell.vel_r_min_kmps + cell.vel_r_max_kmps) / 2.0

From ac83ee0fef2a5d9a974c66ef8e9605f4fe4f5384 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 15:37:59 +0000
Subject: [PATCH 009/150] Update estimators.py

---
 artistools/estimators/estimators.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 385cbd63b..f53139b7c 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -262,9 +262,9 @@ def read_estimators(
     skip_emptycells: bool = False,
     add_velocity: bool = True,
 ) -> dict[tuple[int, int], dict[str, t.Any]]:
-    """Read estimator files into a nested dictionary structure.
+    """Read estimator files into a dictionary of (timestep, modelgridindex): estimators.
 
-    Speed it up by only retrieving estimators for a particular timestep(s) or modelgrid cells.
+    Selecting particular timesteps or modelgrid cells will using speed this up by reducing the number of files that must be read.
     """
     modelpath = Path(modelpath)
     match_modelgridindex: t.Collection[int]
@@ -345,7 +345,12 @@ def read_estimators(
 
                 del estimators_thisfile[k]
 
-            estimators |= estimators_thisfile
+            estimators |= {
+                (ts, mgi): v
+                for (ts, mgi), v in estimators_thisfile.items()
+                if (not match_modelgridindex or mgi in match_modelgridindex)
+                and (not match_timestep or ts in match_timestep)
+            }
 
     return estimators
 

From 0b667193e99444d39fbda02b9cc7418c70d0913d Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 15:40:51 +0000
Subject: [PATCH 010/150] Remove get_heatingcooling option to parse_estimfile()
 and calling functions

---
 artistools/estimators/estimators.py     | 10 ++--------
 artistools/linefluxes.py                |  4 ++--
 artistools/nonthermal/plotnonthermal.py |  2 +-
 3 files changed, 5 insertions(+), 11 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index f53139b7c..579abff29 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -120,7 +120,6 @@ def get_units_string(variable: str) -> str:
 def parse_estimfile(
     estfilepath: Path | str,
     get_ion_values: bool = True,
-    get_heatingcooling: bool = True,
     skip_emptycells: bool = False,
 ) -> t.Iterator[tuple[int, int, dict[str, t.Any]]]:  # pylint: disable=unused-argument
     """Generate timestep, modelgridindex, dict from estimator file."""
@@ -208,7 +207,7 @@ def parse_estimfile(
                     estimblock.setdefault("nntot", 0.0)
                     estimblock["nntot"] += estimblock[f"populations_{atomic_number}"]
 
-            elif row[0] == "heating:" and get_heatingcooling:
+            elif row[0] == "heating:":
                 for heatingtype, value in zip(row[1::2], row[2::2]):
                     key = heatingtype if heatingtype.startswith("heating_") else f"heating_{heatingtype}"
                     estimblock[key] = float(value)
@@ -218,7 +217,7 @@ def parse_estimfile(
                 elif "heating_dep/total_dep" in estimblock and estimblock["heating_dep/total_dep"] > 0:
                     estimblock["total_dep"] = estimblock["heating_dep"] / estimblock["heating_dep/total_dep"]
 
-            elif row[0] == "cooling:" and get_heatingcooling:
+            elif row[0] == "cooling:":
                 for coolingtype, value in zip(row[1::2], row[2::2]):
                     estimblock[f"cooling_{coolingtype}"] = float(value)
 
@@ -232,7 +231,6 @@ def read_estimators_from_file(
     modelpath: Path,
     printfilename: bool = False,
     get_ion_values: bool = True,
-    get_heatingcooling: bool = True,
     skip_emptycells: bool = False,
 ) -> dict[tuple[int, int], dict[str, t.Any]]:
     if printfilename:
@@ -245,7 +243,6 @@ def read_estimators_from_file(
         for timestep, mgi, file_estimblock in parse_estimfile(
             estfilepath,
             get_ion_values=get_ion_values,
-            get_heatingcooling=get_heatingcooling,
             skip_emptycells=skip_emptycells,
         )
     }
@@ -258,7 +255,6 @@ def read_estimators(
     mpirank: int | None = None,
     runfolder: None | str | Path = None,
     get_ion_values: bool = True,
-    get_heatingcooling: bool = True,
     skip_emptycells: bool = False,
     add_velocity: bool = True,
 ) -> dict[tuple[int, int], dict[str, t.Any]]:
@@ -319,7 +315,6 @@ def read_estimators(
             read_estimators_from_file,
             modelpath=modelpath,
             get_ion_values=get_ion_values,
-            get_heatingcooling=get_heatingcooling,
             printfilename=printfilename,
             skip_emptycells=skip_emptycells,
         )
@@ -464,7 +459,6 @@ def get_temperatures(modelpath: str | Path) -> pl.LazyFrame:
         estimators = at.estimators.read_estimators(
             modelpath,
             get_ion_values=False,
-            get_heatingcooling=False,
             skip_emptycells=True,
         )
         assert len(estimators) > 0
diff --git a/artistools/linefluxes.py b/artistools/linefluxes.py
index 032aec7a3..fdd40595b 100755
--- a/artistools/linefluxes.py
+++ b/artistools/linefluxes.py
@@ -455,7 +455,7 @@ def get_packets_with_emission_conditions(
     tend: float,
     maxpacketfiles: int | None = None,
 ) -> pd.DataFrame:
-    estimators = at.estimators.read_estimators(modelpath, get_ion_values=False, get_heatingcooling=False)
+    estimators = at.estimators.read_estimators(modelpath, get_ion_values=False)
 
     modeldata, _ = at.inputmodel.get_modeldata(modelpath)
     ts = at.get_timestep_of_timedays(modelpath, tend)
@@ -655,7 +655,7 @@ def make_emitting_regions_plot(args):
                             "em_Te": dfpackets_selected.em_Te.to_numpy(),
                         }
 
-            estimators = at.estimators.read_estimators(modelpath, get_ion_values=False, get_heatingcooling=False)
+            estimators = at.estimators.read_estimators(modelpath, get_ion_values=False)
             modeldata, _ = at.inputmodel.get_modeldata(modelpath)
             Tedata_all[modelindex] = {}
             log10nnedata_all[modelindex] = {}
diff --git a/artistools/nonthermal/plotnonthermal.py b/artistools/nonthermal/plotnonthermal.py
index 0cc737d20..d89077300 100755
--- a/artistools/nonthermal/plotnonthermal.py
+++ b/artistools/nonthermal/plotnonthermal.py
@@ -70,7 +70,7 @@ def make_xs_plot(axis: plt.Axes, nonthermaldata: pd.DataFrame, args: argparse.Na
 
 def plot_contributions(axis, modelpath, timestep, modelgridindex, nonthermaldata, args):
     estimators = at.estimators.read_estimators(
-        modelpath, get_ion_values=True, get_heatingcooling=True, modelgridindex=modelgridindex, timestep=timestep
+        modelpath, get_ion_values=True, modelgridindex=modelgridindex, timestep=timestep
     )
 
     total_depev = estimators[(timestep, modelgridindex)]["total_dep"] * u.erg.to("eV")

From 7d316fc87fedf4bbe35d2f289ed514f5b2118a6f Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 20:47:19 +0000
Subject: [PATCH 011/150] Update radfield.py

---
 artistools/radfield.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/artistools/radfield.py b/artistools/radfield.py
index 36d210c1b..850d9bdb5 100755
--- a/artistools/radfield.py
+++ b/artistools/radfield.py
@@ -24,9 +24,10 @@
 
 
 @lru_cache(maxsize=4)
-def read_files(modelpath, timestep=-1, modelgridindex=-1):
+def read_files(modelpath: Path | str, timestep: int | None = None, modelgridindex: int | None = None):
     """Read radiation field data from a list of file paths into a pandas DataFrame."""
     radfielddata = pd.DataFrame()
+    modelpath = Path(modelpath)
 
     mpiranklist = at.get_mpiranklist(modelpath, modelgridindex=modelgridindex)
     for folderpath in at.get_runfolders(modelpath, timestep=timestep):
@@ -35,21 +36,21 @@ def read_files(modelpath, timestep=-1, modelgridindex=-1):
             radfieldfilepath = Path(folderpath, radfieldfilename)
             radfieldfilepath = at.firstexisting(radfieldfilename, folder=folderpath, tryzipped=True)
 
-            if modelgridindex > -1:
+            if modelgridindex is not None:
                 filesize = Path(radfieldfilepath).stat().st_size / 1024 / 1024
                 print(f"Reading {Path(radfieldfilepath).relative_to(modelpath.parent)} ({filesize:.2f} MiB)")
 
             radfielddata_thisfile = pd.read_csv(radfieldfilepath, delim_whitespace=True)
             # radfielddata_thisfile[['modelgridindex', 'timestep']].apply(pd.to_numeric)
 
-            if timestep >= 0:
+            if timestep is not None:
                 radfielddata_thisfile = radfielddata_thisfile.query("timestep==@timestep")
 
-            if modelgridindex >= 0:
+            if modelgridindex is not None:
                 radfielddata_thisfile = radfielddata_thisfile.query("modelgridindex==@modelgridindex")
 
             if not radfielddata_thisfile.empty:
-                if timestep >= 0 and modelgridindex >= 0:
+                if timestep is not None and modelgridindex is not None:
                     return radfielddata_thisfile
                 radfielddata = radfielddata.append(radfielddata_thisfile.copy(), ignore_index=True)
 

From 8d4f71534711c2a9cbc1ce8e12c88fea0b5c73dc Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 20:47:42 +0000
Subject: [PATCH 012/150] Update estimators.py

---
 artistools/estimators/estimators.py | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 579abff29..5417f53b0 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -263,20 +263,17 @@ def read_estimators(
     Selecting particular timesteps or modelgrid cells will using speed this up by reducing the number of files that must be read.
     """
     modelpath = Path(modelpath)
-    match_modelgridindex: t.Collection[int]
+    match_modelgridindex: None | t.Sequence[int]
     if modelgridindex is None:
-        match_modelgridindex = []
+        match_modelgridindex = None
     elif isinstance(modelgridindex, int):
         match_modelgridindex = (modelgridindex,)
     else:
         match_modelgridindex = tuple(modelgridindex)
 
-    if -1 in match_modelgridindex:
-        match_modelgridindex = []
-
-    match_timestep: t.Collection[int]
+    match_timestep: None | t.Sequence[int]
     if timestep is None:
-        match_timestep = []
+        match_timestep = None
     elif isinstance(timestep, int):
         match_timestep = (timestep,)
     else:

From c8aca8a0f8a3fe2eae1e7426c5685c0f496976f5 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 20:57:12 +0000
Subject: [PATCH 013/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 6ad3111cf..a104c14b5 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -928,7 +928,7 @@ def addargs(parser: argparse.ArgumentParser) -> None:
     parser.add_argument("--recombrates", action="store_true", help="Make a recombination rate plot")
 
     parser.add_argument(
-        "-modelgridindex", "-cell", "-mgi", type=int, default=-1, help="Modelgridindex for time evolution plot"
+        "-modelgridindex", "-cell", "-mgi", type=int, default=None, help="Modelgridindex for time evolution plot"
     )
 
     parser.add_argument("-timestep", "-ts", nargs="?", help="Timestep number for internal structure plot")
@@ -1024,7 +1024,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
 
     modelname = at.get_model_name(modelpath)
 
-    if not args.timedays and not args.timestep and args.modelgridindex > -1:
+    if not args.timedays and not args.timestep and args.modelgridindex is not None:
         args.timestep = f"0-{len(at.get_timestep_times(modelpath)) - 1}"
 
     (timestepmin, timestepmax, args.timemin, args.timemax) = at.get_time_range(
@@ -1129,7 +1129,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
 
     assoc_cells, mgi_of_propcells = at.get_grid_mapping(modelpath)
 
-    if not args.readonlymgi and (args.modelgridindex > -1 or args.x in {"time", "timestep"}):
+    if not args.readonlymgi and (args.modelgridindex is not None or args.x in {"time", "timestep"}):
         # plot time evolution in specific cell
         if not args.x:
             args.x = "time"

From 0ccdb3139fa1bb9fe61d95ea0ff035d4ba714e0d Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 21:01:24 +0000
Subject: [PATCH 014/150] Update estimators.py

---
 artistools/estimators/estimators.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 5417f53b0..e319ea37c 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -124,8 +124,8 @@ def parse_estimfile(
 ) -> t.Iterator[tuple[int, int, dict[str, t.Any]]]:  # pylint: disable=unused-argument
     """Generate timestep, modelgridindex, dict from estimator file."""
     with at.zopen(estfilepath) as estimfile:
-        timestep: int = -1
-        modelgridindex: int = -1
+        timestep: int | None = None
+        modelgridindex: int | None = None
         estimblock: dict[str, t.Any] = {}
         for line in estimfile:
             row: list[str] = line.split()
@@ -135,8 +135,8 @@ def parse_estimfile(
             if row[0] == "timestep":
                 # yield the previous block before starting a new one
                 if (
-                    timestep >= 0
-                    and modelgridindex >= 0
+                    timestep is not None
+                    and modelgridindex is not None
                     and (not skip_emptycells or not estimblock.get("emptycell", True))
                 ):
                     yield timestep, modelgridindex, estimblock
@@ -222,7 +222,11 @@ def parse_estimfile(
                     estimblock[f"cooling_{coolingtype}"] = float(value)
 
     # reached the end of file
-    if timestep >= 0 and modelgridindex >= 0 and (not skip_emptycells or not estimblock.get("emptycell", True)):
+    if (
+        timestep is not None
+        and modelgridindex is not None
+        and (not skip_emptycells or not estimblock.get("emptycell", True))
+    ):
         yield timestep, modelgridindex, estimblock
 
 

From 7fce659b1e3ae4f1bd7bb1c4ae3f58014ad33315 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Wed, 8 Nov 2023 21:03:35 +0000
Subject: [PATCH 015/150] Update estimators.py

---
 artistools/estimators/estimators.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index e319ea37c..60067dd8f 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -182,7 +182,7 @@ def parse_estimfile(
                         ion_stage = int(ion_stage_str.rstrip(":"))
                     except ValueError:
                         if variablename == "populations" and ion_stage_str.startswith(elsymbol):
-                            estimblock[f"{variablename}_{ion_stage_str.rstrip(':')}"] = float(value)
+                            estimblock[f"populations_{ion_stage_str.rstrip(':')}"] = float(value)
                         else:
                             print(ion_stage_str, elsymbol)
                             print(f"Cannot parse row: {row}")

From d40372e92765cf52d5d20ed326a2e5c58abac419 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 08:40:44 +0000
Subject: [PATCH 016/150] Update updatepackages.yml

---
 .github/workflows/updatepackages.yml  | 39 +++++++++++++++++++++++++++
 .github/workflows/updateprecommit.yml | 30 ---------------------
 2 files changed, 39 insertions(+), 30 deletions(-)
 create mode 100644 .github/workflows/updatepackages.yml
 delete mode 100644 .github/workflows/updateprecommit.yml

diff --git a/.github/workflows/updatepackages.yml b/.github/workflows/updatepackages.yml
new file mode 100644
index 000000000..12b4775af
--- /dev/null
+++ b/.github/workflows/updatepackages.yml
@@ -0,0 +1,39 @@
+name: Update packages
+
+on:
+  # every day at 9am
+  schedule:
+    - cron: "0 9 * * *"
+  workflow_dispatch:
+
+permissions:
+  pull-requests: write
+
+jobs:
+  upgrade:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+            cache: pip
+            python-version: ${{ matrix.python-version }}
+
+      - name: Install requirements
+        run: python -m pip pur
+
+      - name: Run pre-commit autoupdate
+        run: pre-commit autoupdate
+
+      - name: Commit and push
+        run: |
+          git add ".pre-commit-config.yaml"
+          git add "requirements.txt"
+          git commit -m "Upgrade pre-commit dependencies"
+          git push origin upgrade/pre-commit
+
+      - name: Open pull request
+        run: |
+          gh pr create --fill
\ No newline at end of file
diff --git a/.github/workflows/updateprecommit.yml b/.github/workflows/updateprecommit.yml
deleted file mode 100644
index 0ec8bc10d..000000000
--- a/.github/workflows/updateprecommit.yml
+++ /dev/null
@@ -1,30 +0,0 @@
-name: Update pre-commit
-
-on:
-  # every week on monday
-  schedule:
-    - cron: "0 0 * * 1"
-  workflow_dispatch:
-
-permissions:
-  pull-requests: write
-
-jobs:
-  upgrade:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v3
-
-      - name: Run autoupdate
-        run: |
-          pre-commit autoupdate
-
-      - name: Commit and push
-        run: |
-          git add ".pre-commit-config.yaml"
-          git commit -m "Upgrade pre-commit dependencies"
-          git push origin upgrade/pre-commit
-
-      - name: Open pull request
-        run: |
-          gh pr create --fill
\ No newline at end of file

From ab49ea31cace86e5f6980c48f926f63782d65241 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 08:44:33 +0000
Subject: [PATCH 017/150] Update updatepackages.yml

---
 .github/workflows/updatepackages.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/updatepackages.yml b/.github/workflows/updatepackages.yml
index 12b4775af..dfd5dadbd 100644
--- a/.github/workflows/updatepackages.yml
+++ b/.github/workflows/updatepackages.yml
@@ -15,11 +15,11 @@ jobs:
     steps:
       - uses: actions/checkout@v3
 
-      - name: Set up Python ${{ matrix.python-version }}
+      - name: Set up Python
         uses: actions/setup-python@v4
         with:
             cache: pip
-            python-version: ${{ matrix.python-version }}
+            python-version-file: .python-version
 
       - name: Install requirements
         run: python -m pip pur

From 40fa39e8b0f655ae36bdb6f81bf94332c1c1d46e Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 09:18:39 +0000
Subject: [PATCH 018/150] Remove astropy from spectra.py

---
 artistools/spectra/spectra.py | 29 +++++++++++------------------
 1 file changed, 11 insertions(+), 18 deletions(-)

diff --git a/artistools/spectra/spectra.py b/artistools/spectra/spectra.py
index ce2242149..49dc1bbc4 100644
--- a/artistools/spectra/spectra.py
+++ b/artistools/spectra/spectra.py
@@ -15,14 +15,13 @@
 import numpy.typing as npt
 import pandas as pd
 import polars as pl
-from astropy import constants as const
-from astropy import units as u
 
 import artistools as at
 
 fluxcontributiontuple = namedtuple(
     "fluxcontributiontuple", "fluxcontrib linelabel array_flambda_emission array_flambda_absorption color"
 )
+megaparsec_to_cm = 3.0856e24
 
 
 def timeshift_fluxscale_co56law(scaletoreftime: float | None, spectime: float) -> float:
@@ -191,7 +190,6 @@ def get_from_packets(
     dfpackets = dfpackets.select(getcols).collect().lazy()
 
     dfdict = {}
-    megaparsec_to_cm = 3.085677581491367e24
     for dirbin in directionbins:
         if dirbin == -1:
             solidanglefactor = 1.0
@@ -726,7 +724,7 @@ def get_flux_contributions(
         nions = elementlist.nions[element]
         # nions = elementlist.iloc[element].uppermost_ionstage - elementlist.iloc[element].lowermost_ionstage + 1
         for ion in range(nions):
-            ion_stage = ion + elementlist.lowermost_ionstage[element]
+            ionstage = ion + elementlist.lowermost_ionstage[element]
             ionserieslist: list[tuple[int, str]] = [
                 (element * maxion + ion, "bound-bound"),
                 (nelements * maxion + element * maxion + ion, "bound-free"),
@@ -781,11 +779,11 @@ def get_flux_contributions(
                 )
 
                 if emissiontypeclass == "bound-bound":
-                    linelabel = at.get_ionstring(elementlist.Z[element], ion_stage)
+                    linelabel = at.get_ionstring(elementlist.Z[element], ionstage)
                 elif emissiontypeclass == "free-free":
                     linelabel = "free-free"
                 else:
-                    linelabel = f"{at.get_ionstring(elementlist.Z[element], ion_stage)} {emissiontypeclass}"
+                    linelabel = f"{at.get_ionstring(elementlist.Z[element], ionstage)} {emissiontypeclass}"
 
                 contribution_list.append(
                     fluxcontributiontuple(
@@ -839,14 +837,14 @@ def get_emprocesslabel(
 
             if groupby == "terms":
                 upper_config = (
-                    adata.query("Z == @line.atomic_number and ion_stage == @line.ionstage", inplace=False)
+                    adata.query("Z == @line.atomic_number and ionstage == @line.ionstage", inplace=False)
                     .iloc[0]
                     .levels.iloc[line.upperlevelindex]
                     .levelname
                 )
                 upper_term_noj = upper_config.split("_")[-1].split("[")[0]
                 lower_config = (
-                    adata.query("Z == @line.atomic_number and ion_stage == @line.ionstage", inplace=False)
+                    adata.query("Z == @line.atomic_number and ionstage == @line.ionstage", inplace=False)
                     .iloc[0]
                     .levels.iloc[line.lowerlevelindex]
                     .levelname
@@ -856,7 +854,7 @@ def get_emprocesslabel(
 
             if groupby == "upperterm":
                 upper_config = (
-                    adata.query("Z == @line.atomic_number and ion_stage == @line.ionstage", inplace=False)
+                    adata.query("Z == @line.atomic_number and ionstage == @line.ionstage", inplace=False)
                     .iloc[0]
                     .levels.iloc[line.upperlevelindex]
                     .levelname
@@ -902,8 +900,8 @@ def get_absprocesslabel(linelist: dict[int, at.linetuple], abstype: int) -> str:
 
     if use_escapetime:
         modeldata, _ = at.inputmodel.get_modeldata(modelpath)
-        vmax = modeldata.iloc[-1].vel_r_max_kmps * u.km / u.s
-        betafactor = math.sqrt(1 - (vmax / const.c).decompose().value ** 2)
+        vmax = modeldata.iloc[-1].vel_r_max_kmps * 1e5
+        betafactor = math.sqrt(1 - (vmax / 29979245800) ** 2)
 
     packetsfiles = at.packets.get_packetsfilepaths(modelpath, maxpacketfiles)
 
@@ -1032,7 +1030,6 @@ def get_absprocesslabel(linelist: dict[int, at.linetuple], abstype: int) -> str:
             print("volume", volume, "shell volume", volume_shells, "-------------------------------------------------")
         normfactor = c_cgs / 4 / math.pi / delta_lambda / volume / nprocs_read
     else:
-        megaparsec_to_cm = 3.085677581491367e24
         normfactor = 1.0 / delta_lambda / (timehigh - timelow) / 4 / math.pi / (megaparsec_to_cm**2) / nprocs_read
 
     array_flambda_emission_total = energysum_spectrum_emission_total * normfactor
@@ -1169,15 +1166,11 @@ def sortkey(x: fluxcontributiontuple) -> tuple[int, float]:
 def print_integrated_flux(
     arr_f_lambda: np.ndarray | pd.Series, arr_lambda_angstroms: np.ndarray | pd.Series, distance_megaparsec: float = 1.0
 ) -> float:
-    integrated_flux = (
-        abs(np.trapz(np.nan_to_num(arr_f_lambda, nan=0.0), x=arr_lambda_angstroms)) * u.erg / u.s / (u.cm**2)
-    )
+    integrated_flux = abs(np.trapz(np.nan_to_num(arr_f_lambda, nan=0.0), x=arr_lambda_angstroms))
     print(
         f" integrated flux ({arr_lambda_angstroms.min():.1f} to "
-        f"{arr_lambda_angstroms.max():.1f} A): {integrated_flux:.3e}"
+        f"{arr_lambda_angstroms.max():.1f} A): {integrated_flux:.3e} erg/s/cm2"
     )
-    # luminosity = integrated_flux * 4 * math.pi * (distance_megaparsec * u.megaparsec ** 2)
-    # print(f'(L={luminosity.to("Lsun"):.3e})')
     return integrated_flux.value
 
 

From 1aff6447e03859f104c85e0e912856dabca83691 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 09:20:31 +0000
Subject: [PATCH 019/150] Update plotnonthermal.py

---
 artistools/nonthermal/plotnonthermal.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/artistools/nonthermal/plotnonthermal.py b/artistools/nonthermal/plotnonthermal.py
index d89077300..ca57db80d 100755
--- a/artistools/nonthermal/plotnonthermal.py
+++ b/artistools/nonthermal/plotnonthermal.py
@@ -10,11 +10,11 @@
 import numpy as np
 import pandas as pd
 import pynonthermal
-from astropy import units as u
 
 import artistools as at
 
 defaultoutputfile = "plotnonthermal_cell{0:03d}_timestep{1:03d}.pdf"
+ERG_TO_EV = 6.242e11
 
 
 @lru_cache(maxsize=4)
@@ -73,7 +73,7 @@ def plot_contributions(axis, modelpath, timestep, modelgridindex, nonthermaldata
         modelpath, get_ion_values=True, modelgridindex=modelgridindex, timestep=timestep
     )
 
-    total_depev = estimators[(timestep, modelgridindex)]["total_dep"] * u.erg.to("eV")
+    total_depev = estimators[(timestep, modelgridindex)]["total_dep"] * ERG_TO_EV
 
     print(f"Deposition: {total_depev:.1f} [eV/cm3/s]")
 

From c85588c8cdd4d8e363b750d0c12ce3b44a42f4cc Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 11:03:05 +0000
Subject: [PATCH 020/150] Replace estimator keys prefix populations_ with
 nnion_, nniso_, nnelement_

---
 artistools/atomic/_atomic_core.py             |   2 +-
 artistools/codecomparison.py                  |  16 +--
 artistools/deposition.py                      |   2 +-
 artistools/estimators/__init__.py             |   1 +
 artistools/estimators/estimators.py           |  68 +++++++----
 artistools/estimators/estimators_classic.py   |  13 ++-
 artistools/estimators/exportmassfractions.py  |   4 +-
 artistools/estimators/plotestimators.py       | 110 +++++++++---------
 artistools/gsinetwork.py                      |  10 +-
 artistools/linefluxes.py                      |  20 ++--
 artistools/misc.py                            |  61 ++++++----
 artistools/nltepops/nltepops.py               |  18 +--
 artistools/nltepops/plotnltepops.py           |  77 ++++++------
 artistools/nonthermal/_nonthermal_core.py     |  41 +++----
 artistools/nonthermal/plotnonthermal.py       |  17 +--
 artistools/nonthermal/solvespencerfanocmd.py  |  14 +--
 artistools/nonthermal/test_nonthermal.py      |  14 +++
 artistools/plotspherical.py                   |  14 +--
 artistools/radfield.py                        |  96 +++++++--------
 .../spectra/sampleblackbodyfrompacket_tr.py   |   2 +-
 artistools/spectra/spectra.py                 |   2 +-
 artistools/test_artistools.py                 |  28 +++--
 artistools/transitions.py                     |  27 ++---
 artistools/writecomparisondata.py             |  19 ++-
 pyproject.toml                                |   2 +
 25 files changed, 364 insertions(+), 314 deletions(-)
 create mode 100644 artistools/nonthermal/test_nonthermal.py

diff --git a/artistools/atomic/_atomic_core.py b/artistools/atomic/_atomic_core.py
index 8b8e69af8..136b9e55a 100644
--- a/artistools/atomic/_atomic_core.py
+++ b/artistools/atomic/_atomic_core.py
@@ -153,7 +153,7 @@ def get_levels(modelpath, ionlist=None, get_transitions=False, get_photoionisati
                 phixsdict[(Z, lowerionstage, lowerionlevel)] = (phixstargetlist, phixstable)
 
     level_lists = []
-    iontuple = namedtuple("ion", "Z ion_stage level_count ion_pot levels transitions")
+    iontuple = namedtuple("ion", "Z ionstage level_count ion_pot levels transitions")
 
     with at.zopen(adatafilename) as fadata:
         if not quiet:
diff --git a/artistools/codecomparison.py b/artistools/codecomparison.py
index 701df2204..96717232c 100644
--- a/artistools/codecomparison.py
+++ b/artistools/codecomparison.py
@@ -140,9 +140,9 @@ def read_reference_estimators(
                         if ion_startnumber is None:
                             ion_startnumber = ion_number
 
-                        ion_stage = ion_number + 1 if ion_startnumber == 0 else ion_number
+                        ionstage = ion_number + 1 if ion_startnumber == 0 else ion_number
 
-                        iontuples.append((atomic_number, ion_stage))
+                        iontuples.append((atomic_number, ionstage))
 
                 elif not line.lstrip().startswith("#"):
                     cur_modelgridindex += 1
@@ -151,17 +151,19 @@ def read_reference_estimators(
 
                     assert len(row) == nstages + 1
                     assert len(iontuples) == nstages
-                    for (atomic_number, ion_stage), strionfrac in zip(iontuples, row[1:]):
+                    for (atomic_number, ionstage), strionfrac in zip(iontuples, row[1:]):
+                        elsym = at.get_elsymbol(atomic_number)
+                        ionstr = at.get_ionstring(atomic_number, ionstage, sep="_", style="spectral")
                         try:
                             ionfrac = float(strionfrac)
                             ionpop = ionfrac * estimators[tsmgi]["nntot"]
                             if ionpop > 1e-80:
-                                estimators[tsmgi][f"populations_{atomic_number}_{ion_stage}"] = ionpop
-                                estimators[tsmgi].setdefault(f"populations_{atomic_number}", 0.0)
-                                estimators[tsmgi][f"populations_{atomic_number}"] += ionpop
+                                estimators[tsmgi][f"nnion_{ionstr}"] = ionpop
+                                estimators[tsmgi].setdefault(f"nnelement_{elsym}", 0.0)
+                                estimators[tsmgi][f"nnelement_{elsym}"] += ionpop
 
                         except ValueError:
-                            estimators[tsmgi][f"populations_{atomic_number}_{ion_stage}"] = float("NaN")
+                            estimators[tsmgi][f"nnion_{ionstr}"] = float("NaN")
 
                     assert np.isclose(float(row[0]), estimators[tsmgi]["vel_mid"], rtol=0.01)
                     assert estimators[key]["vel_mid"]
diff --git a/artistools/deposition.py b/artistools/deposition.py
index 5d2f9677d..4951821ee 100755
--- a/artistools/deposition.py
+++ b/artistools/deposition.py
@@ -80,7 +80,7 @@ def main_analytical(args: argparse.Namespace | None = None, argsraw: list[str] |
     # dfnltepops = at.nltepops.read_files(
     #     args.modelpath, timestep=timestep).query('Z == 26')
 
-    # phixs = adata.query('Z==26 & ion_stage==1', inplace=False).iloc[0].levels.iloc[0].phixstable[0][1] * 1e-18
+    # phixs = adata.query('Z==26 & ionstage==1', inplace=False).iloc[0].levels.iloc[0].phixstable[0][1] * 1e-18
 
     global_posdep = 0.0 * u.erg / u.s
     for i, row in dfmodel.iterrows():
diff --git a/artistools/estimators/__init__.py b/artistools/estimators/__init__.py
index 0f74b5b7b..fdc8b7f1d 100644
--- a/artistools/estimators/__init__.py
+++ b/artistools/estimators/__init__.py
@@ -17,5 +17,6 @@
 from artistools.estimators.estimators import parse_estimfile
 from artistools.estimators.estimators import read_estimators
 from artistools.estimators.estimators import read_estimators_from_file
+from artistools.estimators.estimators import read_estimators_polars
 from artistools.estimators.plotestimators import addargs
 from artistools.estimators.plotestimators import main as plot
diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 60067dd8f..45ff7d6bd 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -167,45 +167,43 @@ def parse_estimfile(
                     startindex = 2
                 elsymbol = at.get_elsymbol(atomic_number)
 
-                for ion_stage_str, value in zip(row[startindex::2], row[startindex + 1 :: 2]):
-                    ion_stage_str_strip = ion_stage_str.strip()
-                    if ion_stage_str_strip == "(or":
+                for ionstage_str, value in zip(row[startindex::2], row[startindex + 1 :: 2]):
+                    ionstage_str_strip = ionstage_str.strip()
+                    if ionstage_str_strip == "(or":
                         continue
 
                     value_thision = float(value.rstrip(","))
 
-                    if ion_stage_str_strip == "SUM:":
-                        estimblock[f"{variablename}_{atomic_number}"] = value_thision
+                    if ionstage_str_strip == "SUM:":
+                        estimblock[f"nnelement_{elsymbol}"] = value_thision
                         continue
 
                     try:
-                        ion_stage = int(ion_stage_str.rstrip(":"))
+                        ionstage = int(ionstage_str.rstrip(":"))
                     except ValueError:
-                        if variablename == "populations" and ion_stage_str.startswith(elsymbol):
-                            estimblock[f"populations_{ion_stage_str.rstrip(':')}"] = float(value)
+                        if variablename == "populations" and ionstage_str.startswith(elsymbol):
+                            estimblock[f"nniso_{ionstage_str.rstrip(':')}"] = float(value)
                         else:
-                            print(ion_stage_str, elsymbol)
+                            print(ionstage_str, elsymbol)
                             print(f"Cannot parse row: {row}")
                         continue
 
-                    estimblock[f"{variablename}_{atomic_number}_{ion_stage}"] = value_thision
+                    ionstr = at.get_ionstring(atomic_number, ionstage, sep="_", style="spectral")
+                    estimblock[f"{'nnion' if variablename=='populations' else variablename}_{ionstr}"] = value_thision
 
                     if variablename in {"Alpha_R*nne", "AlphaR*nne"}:
-                        estimblock[f"Alpha_R_{atomic_number}_{ion_stage}"] = (
+                        estimblock[f"Alpha_R_{ionstr}"] = (
                             value_thision / estimblock["nne"] if estimblock["nne"] > 0.0 else float("inf")
                         )
 
-                    else:  # variablename == 'populations':
-                        # contribute the ion population to the element population
-                        estimblock.setdefault(f"{variablename}_{atomic_number}", 0.0)
-                        estimblock[f"{variablename}_{atomic_number}"] += value_thision
+                    elif variablename == "populations":
+                        estimblock.setdefault(f"nnelement_{elsymbol}", 0.0)
+                        estimblock[f"nnelement_{elsymbol}"] += value_thision
 
                 if variablename == "populations":
                     # contribute the element population to the total population
-                    estimblock.setdefault("populations_total", 0.0)
-                    estimblock["populations_total"] += estimblock[f"populations_{atomic_number}"]
                     estimblock.setdefault("nntot", 0.0)
-                    estimblock["nntot"] += estimblock[f"populations_{atomic_number}"]
+                    estimblock["nntot"] += estimblock[f"nnelement_{elsymbol}"]
 
             elif row[0] == "heating:":
                 for heatingtype, value in zip(row[1::2], row[2::2]):
@@ -387,25 +385,26 @@ def get_averageionisation(estimatorstsmgi: dict[str, float], atomic_number: int)
     free_electron_weighted_pop_sum = 0.0
     found = False
     popsum = 0.0
+    elsymb = at.get_elsymbol(atomic_number)
     for key in estimatorstsmgi:
-        if key.startswith(f"populations_{atomic_number}_"):
+        if key.startswith(f"nnion_{elsymb}_"):
             found = True
-            ion_stage = int(key.removeprefix(f"populations_{atomic_number}_"))
-            free_electron_weighted_pop_sum += estimatorstsmgi[key] * (ion_stage - 1)
+            ionstage = at.decode_roman_numeral(key.removeprefix(f"nnion_{elsymb}_"))
+            free_electron_weighted_pop_sum += estimatorstsmgi[key] * (ionstage - 1)
             popsum += estimatorstsmgi[key]
 
     if not found:
         return float("NaN")
 
-    return free_electron_weighted_pop_sum / estimatorstsmgi[f"populations_{atomic_number}"]
+    return free_electron_weighted_pop_sum / estimatorstsmgi[f"nnelement_{elsymb}"]
 
 
 def get_averageexcitation(
-    modelpath: Path, modelgridindex: int, timestep: int, atomic_number: int, ion_stage: int, T_exc: float
+    modelpath: Path, modelgridindex: int, timestep: int, atomic_number: int, ionstage: int, T_exc: float
 ) -> float:
     dfnltepops = at.nltepops.read_files(modelpath, modelgridindex=modelgridindex, timestep=timestep)
     adata = at.atomic.get_levels(modelpath)
-    ionlevels = adata.query("Z == @atomic_number and ion_stage == @ion_stage").iloc[0].levels
+    ionlevels = adata.query("Z == @atomic_number and ionstage == @ionstage").iloc[0].levels
 
     energypopsum = 0
     ionpopsum = 0
@@ -413,7 +412,7 @@ def get_averageexcitation(
         return float("NaN")
 
     dfnltepops_ion = dfnltepops.query(
-        "modelgridindex==@modelgridindex and timestep==@timestep and Z==@atomic_number & ion_stage==@ion_stage"
+        "modelgridindex==@modelgridindex and timestep==@timestep and Z==@atomic_number & ionstage==@ionstage"
     )
 
     k_b = 8.617333262145179e-05  # eV / K  # noqa: F841
@@ -474,3 +473,22 @@ def get_temperatures(modelpath: str | Path) -> pl.LazyFrame:
         )
 
     return pl.scan_parquet(dfest_parquetfile)
+
+
+def read_estimators_polars(*args, **kwargs) -> pl.LazyFrame:
+    estimators = read_estimators(*args, **kwargs)
+    pldf = pl.DataFrame(
+        [
+            {
+                "timestep": ts,
+                "modelgridindex": mgi,
+                **estimvals,
+            }
+            for (ts, mgi), estimvals in estimators.items()
+            if not estimvals.get("emptycell", True)
+        ]
+    )
+    print(pldf.columns)
+    print(pldf.transpose(include_header=True))
+
+    return pldf.lazy()
diff --git a/artistools/estimators/estimators_classic.py b/artistools/estimators/estimators_classic.py
index 30a92a078..7972b0c14 100644
--- a/artistools/estimators/estimators_classic.py
+++ b/artistools/estimators/estimators_classic.py
@@ -31,16 +31,17 @@ def parse_ion_row_classic(row: list[str], outdict: dict[str, t.Any], atomic_comp
     i = 6  # skip first 6 numbers in est file. These are n, TR, Te, W, TJ, grey_depth.
     # Numbers after these 6 are populations
     for atomic_number in elements:
-        for ion_stage in range(1, atomic_composition[atomic_number] + 1):
+        for ionstage in range(1, atomic_composition[atomic_number] + 1):
             value_thision = float(row[i])
-            outdict[f"populations_{atomic_number}_{ion_stage}"] = value_thision
+            ionstr = at.get_ionstring(atomic_number, ionstage, sep="_")
+            outdict[f"nnion_{ionstr}"] = value_thision
             i += 1
 
-            elpop = outdict.get(f"populations_{atomic_number}", 0)
-            outdict[f"populations_{atomic_number}"] = elpop + value_thision
+            elpop = outdict.get(f"nnelement_{atomic_number}", 0)
+            outdict[f"nnelement_{atomic_number}"] = elpop + value_thision
 
-            totalpop = outdict.get("populations_total", 0)
-            outdict["populations_total"] = totalpop + value_thision
+            totalpop = outdict.get("nntot", 0)
+            outdict["nntot"] = totalpop + value_thision
 
 
 def get_first_ts_in_run_directory(modelpath) -> dict[str, int]:
diff --git a/artistools/estimators/exportmassfractions.py b/artistools/estimators/exportmassfractions.py
index 1538155cf..68ed8eb12 100755
--- a/artistools/estimators/exportmassfractions.py
+++ b/artistools/estimators/exportmassfractions.py
@@ -35,8 +35,8 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             numberdens = {}
             totaldens = 0.0  # number density times atomic mass summed over all elements
             for key, val in estimators[(timestep, modelgridindex)].items():
-                if key.startswith("populations") and key.removeprefix("populations_").isdigit():
-                    atomic_number = int(key.removeprefix("populations_"))
+                if key.startswith("nnelement_"):
+                    atomic_number = at.get_atomic_number(key.removeprefix("nnelement_"))
                     numberdens[atomic_number] = val
                     totaldens += numberdens[atomic_number] * elmass[atomic_number]
             massfracs = {
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index a104c14b5..4fb6b2bfc 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -141,7 +141,7 @@ def plot_average_ionisation_excitation(
             atomic_number = at.get_atomic_number(paramvalue)
         else:
             atomic_number = at.get_atomic_number(paramvalue.split(" ")[0])
-            ion_stage = at.decode_roman_numeral(paramvalue.split(" ")[1])
+            ionstage = at.decode_roman_numeral(paramvalue.split(" ")[1])
         ylist = []
         for modelgridindex, timesteps in zip(mgilist, timestepslist):
             valuesum = 0
@@ -156,7 +156,7 @@ def plot_average_ionisation_excitation(
                     T_exc = estimators[(timestep, modelgridindex)]["Te"]
                     valuesum += (
                         at.estimators.get_averageexcitation(
-                            modelpath, modelgridindex, timestep, atomic_number, ion_stage, T_exc
+                            modelpath, modelgridindex, timestep, atomic_number, ionstage, T_exc
                         )
                         * arr_tdelta[timestep]
                     )
@@ -206,13 +206,13 @@ def plot_levelpop(
     for paramvalue in params:
         paramsplit = paramvalue.split(" ")
         atomic_number = at.get_atomic_number(paramsplit[0])
-        ion_stage = at.decode_roman_numeral(paramsplit[1])
+        ionstage = at.decode_roman_numeral(paramsplit[1])
         levelindex = int(paramsplit[2])
 
-        ionlevels = adata.query("Z == @atomic_number and ion_stage == @ion_stage").iloc[0].levels
+        ionlevels = adata.query("Z == @atomic_number and ionstage == @ionstage").iloc[0].levels
         levelname = ionlevels.iloc[levelindex].levelname
         label = (
-            f"{at.get_ionstring(atomic_number, ion_stage, spectral=False)} level {levelindex}:"
+            f"{at.get_ionstring(atomic_number, ionstage, style="chargelatex")} level {levelindex}:"
             f" {at.nltepops.texifyconfiguration(levelname)}"
         )
 
@@ -220,7 +220,7 @@ def plot_levelpop(
 
         # level index query goes outside for caching granularity reasons
         dfnltepops = at.nltepops.read_files(
-            modelpath, dfquery=f"Z=={atomic_number:.0f} and ion_stage=={ion_stage:.0f}"
+            modelpath, dfquery=f"Z=={atomic_number:.0f} and ionstage=={ionstage:.0f}"
         ).query("level==@levelindex")
 
         ylist = []
@@ -233,7 +233,7 @@ def plot_levelpop(
                 levelpop = (
                     dfnltepops.query(
                         "modelgridindex==@modelgridindex and timestep==@timestep and Z==@atomic_number"
-                        " and ion_stage==@ion_stage and level==@levelindex"
+                        " and ionstage==@ionstage and level==@levelindex"
                     )
                     .iloc[0]
                     .n_NLTE
@@ -251,9 +251,9 @@ def plot_levelpop(
         if dfalldata is not None:
             elsym = at.get_elsymbol(atomic_number).lower()
             colname = (
-                f"nlevel_on_dv_{elsym}_ionstage{ion_stage}_level{levelindex}"
+                f"nlevel_on_dv_{elsym}_ionstage{ionstage}_level{levelindex}"
                 if seriestype == "levelpopulation_dn_on_dvel"
-                else f"nnlevel_{elsym}_ionstage{ion_stage}_level{levelindex}"
+                else f"nnlevel_{elsym}_ionstage{ionstage}_level{levelindex}"
             )
             dfalldata[colname] = ylist
 
@@ -306,35 +306,39 @@ def get_iontuple(ionstr):
             compositiondata = at.estimators.estimators_classic.get_atomic_composition(modelpath)
         else:
             compositiondata = at.get_composition_data(modelpath)
-        for atomic_number, ion_stage in iontuplelist:
+        for atomic_number, ionstage in iontuplelist:
             if (
-                not hasattr(ion_stage, "lower")
+                not hasattr(ionstage, "lower")
                 and not args.classicartis
                 and compositiondata.query(
-                    "Z == @atomic_number & lowermost_ionstage <= @ion_stage & uppermost_ionstage >= @ion_stage"
+                    "Z == @atomic_number & lowermost_ionstage <= @ionstage & uppermost_ionstage >= @ionstage"
                 ).empty
             ):
-                missingions.add((atomic_number, ion_stage))
+                missingions.add((atomic_number, ionstage))
 
     except FileNotFoundError:
         print("WARNING: Could not read an ARTIS compositiondata.txt file")
-        for atomic_number, ion_stage in iontuplelist:
+        for atomic_number, ionstage in iontuplelist:
             mgits = (timestepslist[0][0], mgilist[0])
-            if f"populations_{atomic_number}_{ion_stage}" not in estimators[mgits]:
-                missingions.add((atomic_number, ion_stage))
+            ionstr = at.get_ionstring(atomic_number, ionstage, sep="_", style="spectral")
+            if f"nnion_{ionstr}" not in estimators[mgits]:
+                missingions.add((atomic_number, ionstage))
 
     if missingions:
         print(f" Warning: Can't plot {seriestype} for {missingions} because these ions are not in compositiondata.txt")
 
     prev_atomic_number = iontuplelist[0][0]
     colorindex = 0
-    for atomic_number, ion_stage in iontuplelist:
-        if (atomic_number, ion_stage) in missingions:
+    for atomic_number, ionstage in iontuplelist:
+        if (atomic_number, ionstage) in missingions:
             continue
 
         if atomic_number != prev_atomic_number:
             colorindex += 1
 
+        elsymbol = at.get_elsymbol(atomic_number)
+        ionstr = at.get_ionstring(atomic_number, ionstage, sep="_", style="spectral")
+
         if seriestype == "populations":
             if args.ionpoptype == "absolute":
                 ax.set_ylabel("X$_{i}$ [/cm3]")
@@ -351,19 +355,19 @@ def get_iontuple(ionstr):
         ylist = []
         for modelgridindex, timesteps in zip(mgilist, timestepslist):
             if seriestype == "populations":
-                # if (atomic_number, ion_stage) not in estimators[timesteps[0], modelgridindex]["populations"]:
+                # if (atomic_number, ionstage) not in estimators[timesteps[0], modelgridindex]["populations"]:
                 #     print(
-                #         f"Note: population for {(atomic_number, ion_stage)} not in estimators for "
+                #         f"Note: population for {(atomic_number, ionstage)} not in estimators for "
                 #         f"cell {modelgridindex} timesteps {timesteps}"
                 #     )
 
-                if ion_stage == "ALL":
-                    key = f"populations_{atomic_number}"
-                elif hasattr(ion_stage, "lower") and ion_stage.startswith(at.get_elsymbol(atomic_number)):
-                    # not really an ionstage but maybe isotope?
-                    key = f"populations_{ion_stage}"
+                if ionstage == "ALL":
+                    key = f"nnelement_{elsymbol}"
+                elif hasattr(ionstage, "lower") and ionstage.startswith(at.get_elsymbol(atomic_number)):
+                    # not really an ionstage but an isotope name
+                    key = f"nniso_{ionstage}"
                 else:
-                    key = f"populations_{atomic_number}_{ion_stage}"
+                    key = f"nnion_{ionstr}"
 
                 try:
                     estimpop = at.estimators.get_averaged_estimators(
@@ -371,7 +375,7 @@ def get_iontuple(ionstr):
                         estimators,
                         timesteps,
                         modelgridindex,
-                        [key, f"populations_{atomic_number}", "populations_total"],
+                        [key, f"nnelement_{elsymbol}", "nntot"],
                     )
                 except KeyError:
                     print(f"KeyError: {key} not in estimators")
@@ -384,10 +388,10 @@ def get_iontuple(ionstr):
                     if args.ionpoptype == "absolute":
                         yvalue = nionpop  # Plot as fraction of element population
                     elif args.ionpoptype == "elpop":
-                        elpop = estimpop.get(f"populations_{atomic_number}", 0.0)
+                        elpop = estimpop.get(f"nnelement_{elsymbol}", 0.0)
                         yvalue = nionpop / elpop  # Plot as fraction of element population
                     elif args.ionpoptype == "totalpop":
-                        totalpop = estimpop["populations_total"]
+                        totalpop = estimpop["nntot"]
                         yvalue = nionpop / totalpop  # Plot as fraction of total population
                     else:
                         raise AssertionError
@@ -397,32 +401,32 @@ def get_iontuple(ionstr):
                 ylist.append(yvalue)
 
             else:
-                key = f"{seriestype}_{atomic_number}_{ion_stage}"
+                key = f"{seriestype}_{ionstr}"
                 yvalue = at.estimators.get_averaged_estimators(modelpath, estimators, timesteps, modelgridindex, key)[
                     key
                 ]
                 ylist.append(yvalue)
 
         plotlabel = (
-            ion_stage
-            if hasattr(ion_stage, "lower") and ion_stage != "ALL"
-            else at.get_ionstring(atomic_number, ion_stage, spectral=False)
+            ionstage
+            if hasattr(ionstage, "lower") and ionstage != "ALL"
+            else at.get_ionstring(atomic_number, ionstage, style="chargelatex")
         )
 
         color = get_elemcolor(atomic_number=atomic_number)
 
-        # linestyle = ['-.', '-', '--', (0, (4, 1, 1, 1)), ':'] + [(0, x) for x in dashes_list][ion_stage - 1]
-        if ion_stage == "ALL":
+        # linestyle = ['-.', '-', '--', (0, (4, 1, 1, 1)), ':'] + [(0, x) for x in dashes_list][ionstage - 1]
+        if ionstage == "ALL":
             dashes = ()
             linewidth = 1.0
         else:
-            if hasattr(ion_stage, "lower") and ion_stage.endswith("stable"):
+            if hasattr(ionstage, "lower") and ionstage.endswith("stable"):
                 index = 8
-            elif hasattr(ion_stage, "lower"):
+            elif hasattr(ionstage, "lower"):
                 # isotopic abundance, use the mass number
-                index = int(ion_stage.lstrip(at.get_elsymbol(atomic_number)))
+                index = int(ionstage.lstrip(at.get_elsymbol(atomic_number)))
             else:
-                index = ion_stage
+                index = ionstage
 
             dashes_list = [(3, 1, 1, 1), (), (1.5, 1.5), (6, 3), (1, 3)]
             dashes = dashes_list[(index - 1) % len(dashes_list)]
@@ -432,7 +436,7 @@ def get_iontuple(ionstr):
 
             if args.colorbyion:
                 color = f"C{index - 1 % 10}"
-                # plotlabel = f'{at.get_elsymbol(atomic_number)} {at.roman_numerals[ion_stage]}'
+                # plotlabel = f'{at.get_elsymbol(atomic_number)} {at.roman_numerals[ionstage]}'
                 dashes = ()
 
         # assert colorindex < 10
@@ -442,11 +446,11 @@ def get_iontuple(ionstr):
         if dfalldata is not None:
             elsym = at.get_elsymbol(atomic_number).lower()
             if args.ionpoptype == "absolute":
-                colname = f"nnion_{elsym}_ionstage{ion_stage}"
+                colname = f"nnion_{elsym}_ionstage{ionstage}"
             elif args.ionpoptype == "elpop":
-                colname = f"nnion_over_nnelem_{elsym}_ionstage{ion_stage}"
+                colname = f"nnion_over_nnelem_{elsym}_ionstage{ionstage}"
             elif args.ionpoptype == "totalpop":
-                colname = f"nnion_over_nntot_{elsym}_ionstage{ion_stage}"
+                colname = f"nnion_over_nntot_{elsym}_ionstage{ionstage}"
             dfalldata[colname] = ylist
 
         ylist.insert(0, ylist[0])
@@ -825,9 +829,9 @@ def make_plot(
     return outfilename
 
 
-def plot_recombrates(modelpath, estimators, atomic_number, ion_stage_list, **plotkwargs):
+def plot_recombrates(modelpath, estimators, atomic_number, ionstage_list, **plotkwargs):
     fig, axes = plt.subplots(
-        nrows=len(ion_stage_list),
+        nrows=len(ionstage_list),
         ncols=1,
         sharex=True,
         figsize=(5, 8),
@@ -838,10 +842,8 @@ def plot_recombrates(modelpath, estimators, atomic_number, ion_stage_list, **plo
 
     recombcalibrationdata = at.atomic.get_ionrecombratecalibration(modelpath)
 
-    for ax, ion_stage in zip(axes, ion_stage_list):
-        ionstr = (
-            f"{at.get_elsymbol(atomic_number)} {at.roman_numerals[ion_stage]} to {at.roman_numerals[ion_stage - 1]}"
-        )
+    for ax, ionstage in zip(axes, ionstage_list):
+        ionstr = f"{at.get_elsymbol(atomic_number)} {at.roman_numerals[ionstage]} to {at.roman_numerals[ionstage - 1]}"
 
         listT_e = []
         list_rrc = []
@@ -849,11 +851,11 @@ def plot_recombrates(modelpath, estimators, atomic_number, ion_stage_list, **plo
         for dicttimestepmodelgrid in estimators.values():
             if (
                 not dicttimestepmodelgrid["emptycell"]
-                and (atomic_number, ion_stage) in dicttimestepmodelgrid["RRC_LTE_Nahar"]
+                and (atomic_number, ionstage) in dicttimestepmodelgrid["RRC_LTE_Nahar"]
             ):
                 listT_e.append(dicttimestepmodelgrid["Te"])
-                list_rrc.append(dicttimestepmodelgrid["RRC_LTE_Nahar"][(atomic_number, ion_stage)])
-                list_rrc2.append(dicttimestepmodelgrid["Alpha_R"][(atomic_number, ion_stage)])
+                list_rrc.append(dicttimestepmodelgrid["RRC_LTE_Nahar"][(atomic_number, ionstage)])
+                list_rrc2.append(dicttimestepmodelgrid["Alpha_R"][(atomic_number, ionstage)])
 
         if not list_rrc:
             continue
@@ -865,7 +867,7 @@ def plot_recombrates(modelpath, estimators, atomic_number, ion_stage_list, **plo
         ax.plot(listT_e, list_rrc2, linewidth=2, label=f"{ionstr} ARTIS Alpha_R", **plotkwargs)
 
         with contextlib.suppress(KeyError):
-            dfrates = recombcalibrationdata[(atomic_number, ion_stage)].query(
+            dfrates = recombcalibrationdata[(atomic_number, ionstage)].query(
                 "T_e > @T_e_min & T_e < @T_e_max", local_dict={"T_e_min": min(listT_e), "T_e_max": max(listT_e)}
             )
 
@@ -880,7 +882,7 @@ def plot_recombrates(modelpath, estimators, atomic_number, ion_stage_list, **plo
             )
         # rrcfiles = glob.glob(
         #     f'/Users/lshingles/Library/Mobile Documents/com~apple~CloudDocs/GitHub/'
-        #     f'artis-atomic/atomic-data-nahar/{at.get_elsymbol(atomic_number).lower()}{ion_stage - 1}.rrc*.txt')
+        #     f'artis-atomic/atomic-data-nahar/{at.get_elsymbol(atomic_number).lower()}{ionstage - 1}.rrc*.txt')
         # if rrcfiles:
         #     dfrecombrates = get_ionrecombrates_fromfile(rrcfiles[0])
         #
diff --git a/artistools/gsinetwork.py b/artistools/gsinetwork.py
index 5d4a88271..c56be8dc8 100755
--- a/artistools/gsinetwork.py
+++ b/artistools/gsinetwork.py
@@ -501,7 +501,7 @@ def plot_qdot_abund_modelcells(
             rho_cgs = rho_init_cgs * (t_model_init_days / time_days) ** 3
 
             for strnuc, a in zip(arr_strnuc, arr_a):
-                abund = estimators[(nts, mgi)][f"populations_{strnuc}"]
+                abund = estimators[(nts, mgi)][f"nniso_{strnuc}"]
                 massfrac = abund * a * MH / rho_cgs
                 massfrac = massfrac + dfmodel.iloc[mgi][f"X_{strnuc}"] * (correction_factors[strnuc] - 1.0)
 
@@ -519,7 +519,7 @@ def plot_qdot_abund_modelcells(
             if "Ye" not in arr_abund_artis[mgi]:
                 arr_abund_artis[mgi]["Ye"] = []
 
-            abund = estimators[(nts, mgi)].get(f"populations_{strnuc}", 0.0)
+            abund = estimators[(nts, mgi)].get(f"nniso_{strnuc}", 0.0)
             if "Ye" in estimators[(nts, mgi)]:
                 cell_Ye = estimators[(nts, mgi)]["Ye"]
                 arr_abund_artis[mgi]["Ye"].append(cell_Ye)
@@ -530,11 +530,7 @@ def plot_qdot_abund_modelcells(
                 cell_nucleoncount = 0.0
                 cellvolume = dfmodel.iloc[mgi].volume
                 for popkey, abund in estimators[(nts, mgi)].items():
-                    if (
-                        popkey.startswith("populations_")
-                        and "_" not in popkey.removeprefix("populations_")
-                        and abund > 0.0
-                    ):
+                    if popkey.startswith("nniso_") and abund > 0.0:
                         if popkey.endswith("_otherstable"):
                             # TODO: use mean molecular weight, but this is not needed for kilonova input files anyway
                             print(f"WARNING {popkey}={abund} not contributed")
diff --git a/artistools/linefluxes.py b/artistools/linefluxes.py
index fdd40595b..df0161898 100755
--- a/artistools/linefluxes.py
+++ b/artistools/linefluxes.py
@@ -151,7 +151,7 @@ def get_line_fluxes_from_pops(emfeatures, modelpath, arr_tstart=None, arr_tend=N
 
     modeldata, _ = at.inputmodel.get_modeldata(modelpath)
 
-    ionlist = [(feature.atomic_number, feature.ion_stage) for feature in emfeatures]
+    ionlist = [(feature.atomic_number, feature.ionstage) for feature in emfeatures]
     adata = at.atomic.get_levels(modelpath, ionlist=tuple(ionlist), get_transitions=True, get_photoionisations=False)
 
     # timearrayplusend = np.concatenate([arr_tstart, [arr_tend[-1]]])
@@ -162,10 +162,10 @@ def get_line_fluxes_from_pops(emfeatures, modelpath, arr_tstart=None, arr_tend=N
         fluxdata = np.zeros_like(arr_tmid, dtype=float)
 
         dfnltepops = at.nltepops.read_files(
-            modelpath, dfquery=f"Z=={feature.atomic_number:.0f} and ion_stage=={feature.ion_stage:.0f}"
+            modelpath, dfquery=f"Z=={feature.atomic_number:.0f} and ionstage=={feature.ionstage:.0f}"
         ).query("level in @feature.upperlevelindicies")
 
-        ion = adata.query("Z == @feature.atomic_number and ion_stage == @feature.ion_stage").iloc[0]
+        ion = adata.query("Z == @feature.atomic_number and ionstage == @feature.ionstage").iloc[0]
 
         for timeindex, timedays in enumerate(arr_tmid):
             v_inner = modeldata.vel_r_min_kmps.to_numpy() * u.km / u.s
@@ -185,7 +185,7 @@ def get_line_fluxes_from_pops(emfeatures, modelpath, arr_tstart=None, arr_tend=N
                         levelpop = (
                             dfnltepops.query(
                                 "modelgridindex==@modelgridindex and timestep==@timestep and Z==@feature.atomic_number"
-                                " and ion_stage==@feature.ion_stage and level==@upperlevelindex"
+                                " and ionstage==@feature.ionstage and level==@upperlevelindex"
                             )
                             .iloc[0]
                             .n_NLTE
@@ -225,7 +225,7 @@ def get_line_fluxes_from_pops(emfeatures, modelpath, arr_tstart=None, arr_tend=N
 def get_closelines(
     modelpath,
     atomic_number: int,
-    ion_stage: int,
+    ionstage: int,
     approxlambdalabel: str | int,
     lambdamin: float | None = None,
     lambdamax: float | None = None,
@@ -233,7 +233,7 @@ def get_closelines(
     upperlevelindex: int | None = None,
 ):
     dflinelist = at.get_linelist_dataframe(modelpath)
-    dflinelistclosematches = dflinelist.query("atomic_number == @atomic_number and ionstage == @ion_stage").copy()
+    dflinelistclosematches = dflinelist.query("atomic_number == @atomic_number and ionstage == @ionstage").copy()
     if lambdamin is not None:
         dflinelistclosematches = dflinelistclosematches.query("@lambdamin < lambda_angstroms")
     if lambdamax is not None:
@@ -249,8 +249,8 @@ def get_closelines(
     lowerlevelindicies = tuple(dflinelistclosematches.lowerlevelindex.to_numpy())
     lowestlambda = dflinelistclosematches.lambda_angstroms.min()
     highestlamba = dflinelistclosematches.lambda_angstroms.max()
-    colname = f"flux_{at.get_ionstring(atomic_number, ion_stage, nospace=True)}_{approxlambdalabel}"
-    featurelabel = f"{at.get_ionstring(atomic_number, ion_stage)} {approxlambdalabel} Å"
+    colname = f"flux_{at.get_ionstring(atomic_number, ionstage, sep='')}_{approxlambdalabel}"
+    featurelabel = f"{at.get_ionstring(atomic_number, ionstage)} {approxlambdalabel} Å"
 
     return (
         colname,
@@ -260,7 +260,7 @@ def get_closelines(
         lowestlambda,
         highestlamba,
         atomic_number,
-        ion_stage,
+        ionstage,
         upperlevelindicies,
         lowerlevelindicies,
     )
@@ -277,7 +277,7 @@ def get_labelandlineindices(modelpath, emfeaturesearch):
             "lowestlambda",
             "highestlamba",
             "atomic_number",
-            "ion_stage",
+            "ionstage",
             "upperlevelindicies",
             "lowerlevelindicies",
         ],
diff --git a/artistools/misc.py b/artistools/misc.py
index 1fc1547c6..4b6f3af42 100644
--- a/artistools/misc.py
+++ b/artistools/misc.py
@@ -576,11 +576,15 @@ def get_model_name(path: Path | str) -> str:
 
 
 def get_z_a_nucname(nucname: str) -> tuple[int, int]:
-    """Return atomic number and mass number from a string like 'Pb208' (returns 92, 208)."""
-    nucname = nucname.removeprefix("X_")
+    """Return atomic number and mass number from a string like 'Pb208', 'X_Pb208', or "nniso_Pb208' (returns 92, 208)."""
+    if "_" in nucname:
+        nucname = nucname.split("_")[1]
+
     z = get_atomic_number(nucname.rstrip("0123456789"))
     assert z > 0
+
     a = int(nucname.lower().lstrip("abcdefghijklmnopqrstuvwxyz"))
+
     return z, a
 
 
@@ -632,7 +636,9 @@ def get_ion_tuple(ionstr: str) -> tuple[int, int] | int:
 
     Return the atomic number for a string like 'Fe' or '26'.
     """
-    ionstr = ionstr.removeprefix("populations_")
+    if "_" in ionstr:
+        ionstr = ionstr.split("_", maxsplit=1)[1]
+
     if ionstr.isdigit():
         return int(ionstr)
 
@@ -665,8 +671,8 @@ def get_ion_tuple(ionstr: str) -> tuple[int, int] | int:
 def get_ionstring(
     atomic_number: int | np.int64,
     ionstage: None | int | np.int64 | t.Literal["ALL"],
-    spectral: bool = True,
-    nospace: bool = False,
+    style: t.Literal["spectral", "chargelatex", "charge"] = "spectral",
+    sep: str = " ",
 ) -> str:
     """Return a string with the element symbol and ionisation stage."""
     if ionstage is None or ionstage == "ALL":
@@ -674,16 +680,21 @@ def get_ionstring(
 
     assert not isinstance(ionstage, str)
 
-    if spectral:
-        return f"{get_elsymbol(atomic_number)}{'' if nospace else ' '}{roman_numerals[ionstage]}"
-
-    # ion notion e.g. Co+, Fe2+
-    if ionstage > 2:
-        strcharge = r"$^{" + str(ionstage - 1) + r"{+}}$"
+    if style == "spectral":
+        return f"{get_elsymbol(atomic_number)}{sep}{roman_numerals[ionstage]}"
+
+    strcharge = ""
+    if style == "chargelatex":
+        # ion notion e.g. Co+, Fe2+
+        if ionstage > 2:
+            strcharge = r"$^{" + f"{ionstage - 1}" + r"{+}}$"
+        elif ionstage == 2:
+            strcharge = r"$^{+}$"
+    elif ionstage > 2:
+        strcharge = f"{ionstage - 1}+"
     elif ionstage == 2:
-        strcharge = r"$^{+}$"
-    else:
-        strcharge = ""
+        strcharge = "+"
+
     return f"{get_elsymbol(atomic_number)}{strcharge}"
 
 
@@ -941,8 +952,8 @@ def get_bflist(modelpath: Path | str) -> dict[int, tuple[int, int, int, int]]:
             i, elementindex, ionindex, level = rowints[:4]
             upperionlevel = rowints[4] if len(rowints) > 4 else -1
             atomic_number = compositiondata.Z[elementindex]
-            ion_stage = ionindex + compositiondata.lowermost_ionstage[elementindex]
-            bflist[i] = (atomic_number, ion_stage, level, upperionlevel)
+            ionstage = ionindex + compositiondata.lowermost_ionstage[elementindex]
+            bflist[i] = (atomic_number, ionstage, level, upperionlevel)
 
     return bflist
 
@@ -961,8 +972,8 @@ def read_linestatfile(filepath: Path | str) -> tuple[int, list[float], list[int]
     atomic_numbers = data[1].astype(int)
     assert len(atomic_numbers) == nlines
 
-    ion_stages = data[2].astype(int)
-    assert len(ion_stages) == nlines
+    ionstages = data[2].astype(int)
+    assert len(ionstages) == nlines
 
     # the file adds one to the levelindex, i.e. lowest level is 1
     upper_levels = data[3].astype(int)
@@ -971,21 +982,21 @@ def read_linestatfile(filepath: Path | str) -> tuple[int, list[float], list[int]
     lower_levels = data[4].astype(int)
     assert len(lower_levels) == nlines
 
-    return nlines, lambda_angstroms, atomic_numbers, ion_stages, upper_levels, lower_levels
+    return nlines, lambda_angstroms, atomic_numbers, ionstages, upper_levels, lower_levels
 
 
 def get_linelist_pldf(modelpath: Path | str) -> pl.LazyFrame:
     textfile = at.firstexisting("linestat.out", folder=modelpath)
     parquetfile = Path(modelpath, "linelist.out.parquet")
     if not parquetfile.is_file() or parquetfile.stat().st_mtime < textfile.stat().st_mtime:
-        _, lambda_angstroms, atomic_numbers, ion_stages, upper_levels, lower_levels = read_linestatfile(textfile)
+        _, lambda_angstroms, atomic_numbers, ionstages, upper_levels, lower_levels = read_linestatfile(textfile)
 
         pldf = (
             pl.DataFrame(
                 {
                     "lambda_angstroms": lambda_angstroms,
                     "atomic_number": atomic_numbers,
-                    "ion_stage": ion_stages,
+                    "ionstage": ionstages,
                     "upper_level": upper_levels,
                     "lower_level": lower_levels,
                 },
@@ -1003,7 +1014,7 @@ def get_linelist_pldf(modelpath: Path | str) -> pl.LazyFrame:
 
 def get_linelist_dict(modelpath: Path | str) -> dict[int, linetuple]:
     """Return a dict of line tuples from linestat.out."""
-    nlines, lambda_angstroms, atomic_numbers, ion_stages, upper_levels, lower_levels = read_linestatfile(
+    nlines, lambda_angstroms, atomic_numbers, ionstages, upper_levels, lower_levels = read_linestatfile(
         Path(modelpath, "linestat.out")
     )
     return {
@@ -1012,7 +1023,7 @@ def get_linelist_dict(modelpath: Path | str) -> dict[int, linetuple]:
             range(nlines),
             lambda_angstroms,
             atomic_numbers,
-            ion_stages,
+            ionstages,
             upper_levels,
             lower_levels,
         )
@@ -1023,7 +1034,7 @@ def get_linelist_dict(modelpath: Path | str) -> dict[int, linetuple]:
 def get_linelist_dataframe(
     modelpath: Path | str,
 ) -> pd.DataFrame:
-    nlines, lambda_angstroms, atomic_numbers, ion_stages, upper_levels, lower_levels = read_linestatfile(
+    nlines, lambda_angstroms, atomic_numbers, ionstages, upper_levels, lower_levels = read_linestatfile(
         Path(modelpath, "linestat.out")
     )
 
@@ -1031,7 +1042,7 @@ def get_linelist_dataframe(
         {
             "lambda_angstroms": lambda_angstroms,
             "atomic_number": atomic_numbers,
-            "ionstage": ion_stages,
+            "ionstage": ionstages,
             "upperlevelindex": upper_levels,
             "lowerlevelindex": lower_levels,
         },
diff --git a/artistools/nltepops/nltepops.py b/artistools/nltepops/nltepops.py
index 2859f88f3..60eba59bd 100644
--- a/artistools/nltepops/nltepops.py
+++ b/artistools/nltepops/nltepops.py
@@ -68,27 +68,27 @@ def add_lte_pops(modelpath, dfpop, columntemperature_tuples, noprint=False, maxl
     """
     k_b = const.k_B.to("eV / K").value
 
-    for _, row in dfpop.drop_duplicates(["modelgridindex", "timestep", "Z", "ion_stage"]).iterrows():
+    for _, row in dfpop.drop_duplicates(["modelgridindex", "timestep", "Z", "ionstage"]).iterrows():
         modelgridindex = int(row.modelgridindex)
         timestep = int(row.timestep)
         Z = int(row.Z)
-        ion_stage = int(row.ion_stage)
+        ionstage = int(row.ionstage)
 
-        ionlevels = at.atomic.get_levels(modelpath).query("Z == @Z and ion_stage == @ion_stage").iloc[0].levels
+        ionlevels = at.atomic.get_levels(modelpath).query("Z == @Z and ionstage == @ionstage").iloc[0].levels
 
         gs_g = ionlevels.iloc[0].g
         gs_energy = ionlevels.iloc[0].energy_ev
 
         # gs_pop = dfpop.query(
         #     "modelgridindex == @modelgridindex and timestep == @timestep "
-        #     "and Z == @Z and ion_stage == @ion_stage and level == 0"
+        #     "and Z == @Z and ionstage == @ionstage and level == 0"
         # ).iloc[0]["n_NLTE"]
 
         masksuperlevel = (
             (dfpop["modelgridindex"] == modelgridindex)
             & (dfpop["timestep"] == timestep)
             & (dfpop["Z"] == Z)
-            & (dfpop["ion_stage"] == ion_stage)
+            & (dfpop["ionstage"] == ionstage)
             & (dfpop["level"] == -1)
         )
 
@@ -96,7 +96,7 @@ def add_lte_pops(modelpath, dfpop, columntemperature_tuples, noprint=False, maxl
             (dfpop["modelgridindex"] == modelgridindex)
             & (dfpop["timestep"] == timestep)
             & (dfpop["Z"] == Z)
-            & (dfpop["ion_stage"] == ion_stage)
+            & (dfpop["ionstage"] == ionstage)
             & (dfpop["level"] != -1)
         )
 
@@ -116,7 +116,7 @@ def f_ltepop(x, T_exc: float, gsg: float, gse: float, ionlevels) -> float:
             levelnumber_sl = (
                 dfpop.query(
                     "modelgridindex == @modelgridindex and timestep == @timestep "
-                    "and Z == @Z and ion_stage == @ion_stage"
+                    "and Z == @Z and ionstage == @ionstage"
                 ).level.max()
                 + 1
             )
@@ -124,7 +124,7 @@ def f_ltepop(x, T_exc: float, gsg: float, gse: float, ionlevels) -> float:
             if maxlevel < 0 or levelnumber_sl <= maxlevel:
                 if not noprint:
                     print(
-                        f"{at.get_elsymbol(Z)} {at.roman_numerals[ion_stage]} "
+                        f"{at.get_elsymbol(Z)} {at.roman_numerals[ionstage]} "
                         f"has a superlevel at level {levelnumber_sl}"
                     )
 
@@ -160,6 +160,8 @@ def read_file(nltefilepath: str | Path) -> pd.DataFrame:
         dfpop = pd.read_csv(nltefilepath, delim_whitespace=True)
     except pd.errors.EmptyDataError:
         return pd.DataFrame()
+    if "ion_stage" in dfpop.columns:
+        dfpop = dfpop.rename(columns={"ion_stage": "ionstage"})
 
     return dfpop
 
diff --git a/artistools/nltepops/plotnltepops.py b/artistools/nltepops/plotnltepops.py
index de34819f7..fc193630e 100755
--- a/artistools/nltepops/plotnltepops.py
+++ b/artistools/nltepops/plotnltepops.py
@@ -41,22 +41,22 @@ def annotate_emission_line(ax: plt.Axes, y: float, upperlevel: int, lowerlevel:
     )
 
 
-def plot_reference_data(ax, atomic_number: int, ion_stage: int, estimators_celltimestep, dfpopthision, annotatelines):
+def plot_reference_data(ax, atomic_number: int, ionstage: int, estimators_celltimestep, dfpopthision, annotatelines):
     nne, Te, TR, W = (estimators_celltimestep[s] for s in ["nne", "Te", "TR", "W"])
     # comparison to Chianti file
     elsym = at.get_elsymbol(atomic_number)
     elsymlower = elsym.lower()
-    if Path("data", f"{elsymlower}_{ion_stage}-levelmap.txt").exists():
+    if Path("data", f"{elsymlower}_{ionstage}-levelmap.txt").exists():
         # ax.set_ylim(bottom=2e-3)
         # ax.set_ylim(top=4)
-        with Path("data", f"{elsymlower}_{ion_stage}-levelmap.txt").open("r") as levelmapfile:
+        with Path("data", f"{elsymlower}_{ionstage}-levelmap.txt").open("r") as levelmapfile:
             levelnumofconfigterm = {}
             for line in levelmapfile:
                 row = line.split()
                 levelnumofconfigterm[(row[0], row[1])] = int(row[2]) - 1
 
         # ax.set_ylim(bottom=5e-4)
-        for depfilepath in sorted(Path("data").rglob(f"chianti_{elsym}_{ion_stage}_*.txt")):
+        for depfilepath in sorted(Path("data").rglob(f"chianti_{elsym}_{ionstage}_*.txt")):
             with depfilepath.open("r") as depfile:
                 firstline = depfile.readline()
                 file_nne = float(firstline[firstline.find("ne = ") + 5 :].split(",")[0])
@@ -91,7 +91,7 @@ def plot_reference_data(ax, atomic_number: int, ion_stage: int, estimators_cellt
                                 if firstdep < 0:
                                     firstdep = float(row[0])
                                 depcoeffs.append(float(row[0]) / firstdep)
-                    ionstr = at.get_ionstring(atomic_number, ion_stage, spectral=False)
+                    ionstr = at.get_ionstring(atomic_number, ionstage, style="chargelatex")
                     ax.plot(
                         levelnums,
                         depcoeffs,
@@ -103,22 +103,22 @@ def plot_reference_data(ax, atomic_number: int, ion_stage: int, estimators_cellt
                         zorder=-1,
                     )
 
-        if annotatelines and atomic_number == 28 and ion_stage == 2:
+        if annotatelines and atomic_number == 28 and ionstage == 2:
             annotate_emission_line(ax=ax, y=0.04, upperlevel=6, lowerlevel=0, label=r"7378$~\mathrm{{\AA}}$")
             annotate_emission_line(ax=ax, y=0.15, upperlevel=6, lowerlevel=2, label=r"1.939 $\mu$m")
             annotate_emission_line(ax=ax, y=0.26, upperlevel=7, lowerlevel=1, label=r"7412$~\mathrm{{\AA}}$")
 
-    if annotatelines and atomic_number == 26 and ion_stage == 2:
+    if annotatelines and atomic_number == 26 and ionstage == 2:
         annotate_emission_line(ax=ax, y=0.66, upperlevel=9, lowerlevel=0, label=r"12570$~\mathrm{{\AA}}$")
         annotate_emission_line(ax=ax, y=0.53, upperlevel=16, lowerlevel=5, label=r"7155$~\mathrm{{\AA}}$")
 
 
-def get_floers_data(dfpopthision, atomic_number, ion_stage, modelpath, T_e, modelgridindex):
+def get_floers_data(dfpopthision, atomic_number, ionstage, modelpath, T_e, modelgridindex):
     floers_levelnums, floers_levelpop_values = None, None
 
     # comparison to Andeas Floers's NLTE pops for Shingles et al. (2022)
-    if atomic_number == 26 and ion_stage in {2, 3}:
-        floersfilename = "andreas_level_populations_fe2.txt" if ion_stage == 2 else "andreas_level_populations_fe3.txt"
+    if atomic_number == 26 and ionstage in {2, 3}:
+        floersfilename = "andreas_level_populations_fe2.txt" if ionstage == 2 else "andreas_level_populations_fe3.txt"
         if Path(modelpath / floersfilename).is_file():
             print(f"reading {floersfilename}")
             floers_levelpops = pd.read_csv(modelpath / floersfilename, comment="#", delim_whitespace=True)
@@ -168,7 +168,7 @@ def make_ionsubplot(
     ax,
     modelpath,
     atomic_number,
-    ion_stage,
+    ionstage,
     dfpop,
     ion_data,
     estimators,
@@ -180,11 +180,11 @@ def make_ionsubplot(
     lastsubplot,
 ):
     """Plot the level populations the specified ion, cell, and timestep."""
-    ionstr = at.get_ionstring(atomic_number, ion_stage, spectral=False)
+    ionstr = at.get_ionstring(atomic_number, ionstage, style="chargelatex")
 
     dfpopthision = dfpop.query(
         "modelgridindex == @modelgridindex and timestep == @timestep "
-        "and Z == @atomic_number and ion_stage == @ion_stage",
+        "and Z == @atomic_number and ionstage == @ionstage",
         inplace=False,
     ).copy()
 
@@ -198,7 +198,8 @@ def make_ionsubplot(
         dfpopthision = dfpopthision.query("level <= @args.maxlevel")
 
     ionpopulation = dfpopthision["n_NLTE"].sum()
-    ionpopulation_fromest = estimators[(timestep, modelgridindex)].get(f"populations_{atomic_number}_{ion_stage}", 0.0)
+    ionstr = at.get_ionstring(atomic_number, ionstage, sep="_", style="spectral")
+    ionpopulation_fromest = estimators[(timestep, modelgridindex)].get(f"nnion_{ionstr}", 0.0)
 
     dfpopthision["parity"] = [
         1 if (row.level != -1 and ion_data.levels.iloc[int(row.level)].levelname.split("[")[0][-1] == "o") else 0
@@ -237,7 +238,7 @@ def make_ionsubplot(
         ax.set_xticklabels("" for _ in configtexlist)
 
     print(
-        f"{at.get_elsymbol(atomic_number)} {at.roman_numerals[ion_stage]} has a summed "
+        f"{at.get_elsymbol(atomic_number)} {at.roman_numerals[ionstage]} has a summed "
         f"level population of {ionpopulation:.1f} (from estimator file ion pop = {ionpopulation_fromest})"
     )
 
@@ -256,7 +257,7 @@ def make_ionsubplot(
     pd.set_option("display.max_columns", 150)
     if len(dfpopthision) < 30:
         # print(dfpopthision[
-        #     ['Z', 'ion_stage', 'level', 'config', 'departure_coeff', 'texname']].to_string(index=False))
+        #     ['Z', 'ionstage', 'level', 'config', 'departure_coeff', 'texname']].to_string(index=False))
         print(
             dfpopthision.loc[
                 :, [c not in {"timestep", "modelgridindex", "Z", "parity", "texname"} for c in dfpopthision.columns]
@@ -291,7 +292,7 @@ def make_ionsubplot(
     ax.set_yscale("log")
 
     floers_levelnums, floers_levelpop_values = get_floers_data(
-        dfpopthision, atomic_number, ion_stage, modelpath, T_e, modelgridindex
+        dfpopthision, atomic_number, ionstage, modelpath, T_e, modelgridindex
     )
 
     if args.departuremode:
@@ -372,7 +373,7 @@ def make_ionsubplot(
 
     if args.plotrefdata:
         plot_reference_data(
-            ax, atomic_number, ion_stage, estimators[(timestep, modelgridindex)], dfpopthision, annotatelines=True
+            ax, atomic_number, ionstage, estimators[(timestep, modelgridindex)], dfpopthision, annotatelines=True
         )
 
 
@@ -386,7 +387,7 @@ def make_plot_populations_with_time_or_velocity(modelpaths, args):
     ionstage = int(args.ionstages[0])
 
     adata = at.atomic.get_levels(args.modelpath[0], get_transitions=True)
-    ion_data = adata.query("Z == @Z and ion_stage == @ionstage").iloc[0]
+    ion_data = adata.query("Z == @Z and ionstage == @ionstage").iloc[0]
     levelconfignames = ion_data["levels"]["levelname"]
     # levelconfignames = [at.nltepops.texifyconfiguration(name) for name in levelconfignames]
 
@@ -481,7 +482,7 @@ def plot_populations_with_time_or_velocity(ax, modelpaths, timedays, ionstage, i
         for timestep, mgi in zip(timesteps, modelgridindex_list):
             dfpop = at.nltepops.read_files(modelpath, timestep=timestep, modelgridindex=mgi)
             try:
-                timesteppops = dfpop.loc[(dfpop["Z"] == Z) & (dfpop["ion_stage"] == ionstage)]
+                timesteppops = dfpop.loc[(dfpop["Z"] == Z) & (dfpop["ionstage"] == ionstage)]
             except KeyError:
                 continue
             for ionlevel in ionlevels:
@@ -527,22 +528,22 @@ def make_plot(modelpath, atomic_number, ionstages_displayed, mgilist, timestep,
     dfpop = dfpop.query("Z == @atomic_number")
 
     # top_ion = 9999
-    max_ion_stage = dfpop.ion_stage.max()
+    max_ionstage = dfpop.ionstage.max()
 
-    if len(dfpop.query("ion_stage == @max_ion_stage")) == 1:  # single-level ion, so skip it
-        max_ion_stage -= 1
+    if len(dfpop.query("ionstage == @max_ionstage")) == 1:  # single-level ion, so skip it
+        max_ionstage -= 1
 
-    ion_stage_list = sorted(
+    ionstage_list = sorted(
         [
             i
-            for i in dfpop.ion_stage.unique()
-            if i <= max_ion_stage and (ionstages_displayed is None or i in ionstages_displayed)
+            for i in dfpop.ionstage.unique()
+            if i <= max_ionstage and (ionstages_displayed is None or i in ionstages_displayed)
         ]
     )
 
     subplotheight = 2.4 / 6 if args.x == "config" else 1.8 / 6
 
-    nrows = len(ion_stage_list) * len(mgilist)
+    nrows = len(ionstage_list) * len(mgilist)
     fig, axes = plt.subplots(
         nrows=nrows,
         ncols=1,
@@ -557,11 +558,11 @@ def make_plot(modelpath, atomic_number, ionstages_displayed, mgilist, timestep,
     if nrows == 1:
         axes = [axes]
 
-    prev_ion_stage = -1
+    prev_ionstage = -1
     assert len(mgilist) > 0
     for mgilistindex, modelgridindex in enumerate(mgilist):
         mgifirstaxindex = mgilistindex
-        mgilastaxindex = mgilistindex + len(ion_stage_list) - 1
+        mgilastaxindex = mgilistindex + len(ionstage_list) - 1
 
         estimators = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex)
         elsymbol = at.get_elsymbol(atomic_number)
@@ -595,10 +596,10 @@ def make_plot(modelpath, atomic_number, ionstages_displayed, mgilist, timestep,
         dfpop = dfpop.query("Z == @atomic_number")
 
         # top_ion = 9999
-        max_ion_stage = dfpop.ion_stage.max()
+        max_ionstage = dfpop.ionstage.max()
 
-        if len(dfpop.query("ion_stage == @max_ion_stage")) == 1:  # single-level ion, so skip it
-            max_ion_stage -= 1
+        if len(dfpop.query("ionstage == @max_ionstage")) == 1:  # single-level ion, so skip it
+            max_ionstage -= 1
 
         # timearray = at.get_timestep_times(modelpath)
         nne = estimators[(timestep, modelgridindex)]["nne"]
@@ -622,14 +623,14 @@ def make_plot(modelpath, atomic_number, ionstages_displayed, mgilist, timestep,
         if not args.notitle:
             axes[mgifirstaxindex].set_title(subplot_title, fontsize=10)
 
-        for ax, ion_stage in zip(axes[mgifirstaxindex : mgilastaxindex + 1], ion_stage_list):
-            ion_data = adata.query("Z == @atomic_number and ion_stage == @ion_stage").iloc[0]
-            lastsubplot = modelgridindex == mgilist[-1] and ion_stage == ion_stage_list[-1]
+        for ax, ionstage in zip(axes[mgifirstaxindex : mgilastaxindex + 1], ionstage_list):
+            ion_data = adata.query("Z == @atomic_number and ionstage == @ionstage").iloc[0]
+            lastsubplot = modelgridindex == mgilist[-1] and ionstage == ionstage_list[-1]
             make_ionsubplot(
                 ax,
                 modelpath,
                 atomic_number,
-                ion_stage,
+                ionstage,
                 dfpop,
                 ion_data,
                 estimators,
@@ -655,10 +656,10 @@ def make_plot(modelpath, atomic_number, ionstages_displayed, mgilist, timestep,
             if args.ymax is not None:
                 ax.set_ylim(top=args.ymax)
 
-            if not args.nolegend and prev_ion_stage != ion_stage:
+            if not args.nolegend and prev_ionstage != ionstage:
                 ax.legend(loc="best", handlelength=1, frameon=True, numpoints=1, edgecolor="0.93", facecolor="0.93")
 
-            prev_ion_stage = ion_stage
+            prev_ionstage = ionstage
 
     if args.x == "index":
         axes[-1].set_xlabel(r"Level index")
diff --git a/artistools/nonthermal/_nonthermal_core.py b/artistools/nonthermal/_nonthermal_core.py
index 192317387..343c4409c 100755
--- a/artistools/nonthermal/_nonthermal_core.py
+++ b/artistools/nonthermal/_nonthermal_core.py
@@ -94,11 +94,11 @@ def read_binding_energies(modelpath: str = ".") -> np.ndarray:
     return electron_binding
 
 
-def get_electronoccupancy(atomic_number: int, ion_stage: int, nt_shells: int) -> np.ndarray:
+def get_electronoccupancy(atomic_number: int, ionstage: int, nt_shells: int) -> np.ndarray:
     # adapted from ARTIS code
     q = np.zeros(nt_shells)
 
-    ioncharge = ion_stage - 1
+    ioncharge = ionstage - 1
     nbound = atomic_number - ioncharge  # number of bound electrons
 
     for _electron_loop in range(nbound):
@@ -144,12 +144,10 @@ def get_electronoccupancy(atomic_number: int, ion_stage: int, nt_shells: int) ->
     return q
 
 
-def get_mean_binding_energy(
-    atomic_number: int, ion_stage: int, electron_binding: np.ndarray, ionpot_ev: float
-) -> float:
+def get_mean_binding_energy(atomic_number: int, ionstage: int, electron_binding: np.ndarray, ionpot_ev: float) -> float:
     # LJS: this came from ARTIS and I'm not sure what this actually is - inverse binding energy? electrons per erg?
     n_z_binding, nt_shells = electron_binding.shape
-    q = get_electronoccupancy(atomic_number, ion_stage, nt_shells)
+    q = get_electronoccupancy(atomic_number, ionstage, nt_shells)
 
     total = 0.0
     for electron_loop in range(nt_shells):
@@ -166,20 +164,20 @@ def get_mean_binding_energy(
                 # is for 8 (corresponding to that shell) then just use the M4 value
                 print(
                     "WARNING: I'm trying to use a binding energy when I have no data. "
-                    f"element {atomic_number} ionstage {ion_stage}\n"
+                    f"element {atomic_number} ionstage {ionstage}\n"
                 )
                 assert electron_loop == 8
-                # print("Z = %d, ion_stage = %d\n", get_element(element), get_ionstage(element, ion));
+                # print("Z = %d, ionstage = %d\n", get_element(element), get_ionstage(element, ion));
         total += electronsinshell / use3 if use2 < use3 else electronsinshell / use2
         # print("total total)
 
     return total
 
 
-def get_mean_binding_energy_alt(atomic_number, ion_stage, electron_binding, ionpot_ev):
+def get_mean_binding_energy_alt(atomic_number, ionstage, electron_binding, ionpot_ev):
     # LJS: this should be mean binding energy [erg] per electron
     n_z_binding, nt_shells = electron_binding.shape
-    q = get_electronoccupancy(atomic_number, ion_stage, nt_shells)
+    q = get_electronoccupancy(atomic_number, ionstage, nt_shells)
 
     total = 0.0
     ecount = 0
@@ -198,17 +196,17 @@ def get_mean_binding_energy_alt(atomic_number, ion_stage, electron_binding, ionp
                 # is for 8 (corresponding to that shell) then just use the M4 value
                 print(
                     "WARNING: I'm trying to use a binding energy when I have no data. "
-                    f"element {atomic_number} ionstage {ion_stage}\n"
+                    f"element {atomic_number} ionstage {ionstage}\n"
                 )
                 assert electron_loop == 8
-                # print("Z = %d, ion_stage = %d\n", get_element(element), get_ionstage(element, ion));
+                # print("Z = %d, ionstage = %d\n", get_element(element), get_ionstage(element, ion));
         total += electronsinshell * use3 if use2 < use3 else electronsinshell * use2
-    assert ecount == (atomic_number - ion_stage + 1)
+    assert ecount == (atomic_number - ionstage + 1)
 
     return total / ecount
 
 
-def get_lotz_xs_ionisation(atomic_number, ion_stage, electron_binding, ionpot_ev, en_ev):
+def get_lotz_xs_ionisation(atomic_number, ionstage, electron_binding, ionpot_ev, en_ev):
     # Axelrod 1980 Eq 3.38
 
     en_erg = en_ev * EV
@@ -218,7 +216,7 @@ def get_lotz_xs_ionisation(atomic_number, ion_stage, electron_binding, ionpot_ev
     # print(f'{gamma=} {beta=}')
 
     n_z_binding, nt_shells = electron_binding.shape
-    q = get_electronoccupancy(atomic_number, ion_stage, nt_shells)
+    q = get_electronoccupancy(atomic_number, ionstage, nt_shells)
 
     part_sigma = 0.0
     for electron_loop in range(nt_shells):
@@ -235,10 +233,10 @@ def get_lotz_xs_ionisation(atomic_number, ion_stage, electron_binding, ionpot_ev
                 # is for 8 (corresponding to that shell) then just use the M4 value
                 print(
                     "WARNING: I'm trying to use a binding energy when I have no data. "
-                    f"element {atomic_number} ionstage {ion_stage}\n"
+                    f"element {atomic_number} ionstage {ionstage}\n"
                 )
                 assert electron_loop == 8
-                # print("Z = %d, ion_stage = %d\n", get_element(element), get_ionstage(element, ion));
+                # print("Z = %d, ionstage = %d\n", get_element(element), get_ionstage(element, ion));
 
         p = max(use2, use3)
 
@@ -833,7 +831,7 @@ def solve_spencerfano_differentialform(
     dftransitions = {}
     for Z, ionstage in ions:
         nnion = ionpopdict[(Z, ionstage)]
-        print(f"  including Z={Z} ion_stage {ionstage} ({at.get_ionstring(Z, ionstage)}) nnion={nnion:.2e} ionisation")
+        print(f"  including Z={Z} ionstage {ionstage} ({at.get_ionstring(Z, ionstage)}) nnion={nnion:.2e} ionisation")
         dfcollion_thision = dfcollion.query("Z == @Z and ionstage == @ionstage", inplace=False)
         # print(dfcollion_thision)
 
@@ -888,7 +886,7 @@ def analyse_ntspectrum(
     for Z, ionstage in ions:
         nnion = ionpopdict[(Z, ionstage)]
         if nnion == 0.0:
-            print(f"skipping Z={Z:2d} ion_stage {ionstage} due to nnion = {nnion}")
+            print(f"skipping Z={Z:2d} ionstage {ionstage} due to nnion = {nnion}")
             continue
         X_ion = nnion / nntot
         dfcollion_thision = dfcollion.query("Z == @Z and ionstage == @ionstage", inplace=False)
@@ -897,7 +895,7 @@ def analyse_ntspectrum(
         ionpot_valence = dfcollion_thision.ionpot_ev.min()
 
         print(
-            f"====> Z={Z:2d} ion_stage {ionstage} {at.get_ionstring(Z, ionstage)} (valence potential"
+            f"====> Z={Z:2d} ionstage {ionstage} {at.get_ionstring(Z, ionstage)} (valence potential"
             f" {ionpot_valence:.1f} eV)"
         )
 
@@ -1187,7 +1185,7 @@ def calculate_Latom_excitation(ions, ionpopdict, nntot, en_ev, adata, T_exc=5000
     for Z, ionstage in ions:
         nnion = ionpopdict[(Z, ionstage)]
 
-        ion = adata.query("Z == @Z and ion_stage == @ionstage").iloc[0]
+        ion = adata.query("Z == @Z and ionstage == @ionstage").iloc[0]
         # filterquery = 'collstr >= 0 or forbidden == False'
         filterquery = "collstr != -999"
         if maxnlevelslower is not None:
@@ -1345,7 +1343,6 @@ def workfunction_tests(modelpath, args):
         #       f'Latom_axelrod: {Latom_axelrod[i]:.3e} ratio(elec/atom): {Lelec_axelrod_nne[i] / Latom_axelrod[i]:.2e} ratio(atom/elec) {Latom_axelrod[i] / Lelec_axelrod_nne[i]:.2e} bound/free {nnebound/nne:.2e}')
 
     for Z, ionstage in ions:
-        # ionstr = at.get_ionstring(Z, ionstage, spectral=True, nospace=False)
         dfcollion_thision = dfcollion.query("Z == @Z and ionstage == @ionstage", inplace=False)
 
         # dfcollion_thision.query('n == 3 and l == 2', inplace=True)
diff --git a/artistools/nonthermal/plotnonthermal.py b/artistools/nonthermal/plotnonthermal.py
index ca57db80d..785d0d1e3 100755
--- a/artistools/nonthermal/plotnonthermal.py
+++ b/artistools/nonthermal/plotnonthermal.py
@@ -69,11 +69,11 @@ def make_xs_plot(axis: plt.Axes, nonthermaldata: pd.DataFrame, args: argparse.Na
 
 
 def plot_contributions(axis, modelpath, timestep, modelgridindex, nonthermaldata, args):
-    estimators = at.estimators.read_estimators(
+    estim_tsmgi = at.estimators.read_estimators(
         modelpath, get_ion_values=True, modelgridindex=modelgridindex, timestep=timestep
-    )
+    )[(timestep, modelgridindex)]
 
-    total_depev = estimators[(timestep, modelgridindex)]["total_dep"] * ERG_TO_EV
+    total_depev = estim_tsmgi["total_dep"] * ERG_TO_EV
 
     print(f"Deposition: {total_depev:.1f} [eV/cm3/s]")
 
@@ -85,12 +85,12 @@ def plot_contributions(axis, modelpath, timestep, modelgridindex, nonthermaldata
     dfcollion = at.nonthermal.read_colliondata()
 
     elementlist = at.get_composition_data(modelpath)
-    totalpop = estimators[(timestep, modelgridindex)]["populations_total"]
+    totalpop = estim_tsmgi["nntot"]
     nelements = len(elementlist)
     for element in range(nelements):
         Z = elementlist.Z[element]
-
-        elpop = estimators[(timestep, modelgridindex)][f"populations_{Z}"]
+        elsymbol = at.get_elsymbol(Z)
+        elpop = estim_tsmgi[f"nnelement_{elsymbol}"]
         if elpop <= 1e-4 * totalpop:
             continue
 
@@ -100,7 +100,8 @@ def plot_contributions(axis, modelpath, timestep, modelgridindex, nonthermaldata
         nions = elementlist.nions[element]
         for ion in range(nions):
             ionstage = ion + elementlist.lowermost_ionstage[element]
-            ionpop = estimators[(timestep, modelgridindex)][f"populations_{Z}_{ionstage}"]
+            ionstr = at.get_ionstring(Z, ionstage, sep="_", style="spectral")
+            ionpop = estim_tsmgi[f"nnion_{ionstr}"]
 
             dfcollion_thision = dfcollion.query("Z == @Z and ionstage == @ionstage")
 
@@ -125,7 +126,7 @@ def plot_contributions(axis, modelpath, timestep, modelgridindex, nonthermaldata
         if frac_ionisation_element > 1e-5:
             axis.plot(arr_enev, arr_ionisation_element, label=f"Ionisation Z={Z}")
 
-    nne = estimators[(timestep, modelgridindex)]["nne"]
+    nne = estim_tsmgi["nne"]
     arr_heating = np.array([at.nonthermal.lossfunction(enev, nne) / total_depev for enev in arr_enev])
 
     frac_heating = np.trapz(x=arr_enev, y=arr_heating)
diff --git a/artistools/nonthermal/solvespencerfanocmd.py b/artistools/nonthermal/solvespencerfanocmd.py
index 6644bee97..5a358e3ab 100755
--- a/artistools/nonthermal/solvespencerfanocmd.py
+++ b/artistools/nonthermal/solvespencerfanocmd.py
@@ -186,16 +186,12 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             print(f"ERROR: no NLTE populations for cell {args.modelgridindex} at timestep {args.timestep}")
             raise AssertionError
 
-        nntot = estim["populations_total"]
+        nntot = estim["nntot"]
         # nne = estim["nne"]
         T_e = estim["Te"]
         print("WARNING: Use LTE pops at Te for now")
         deposition_density_ev = estim["heating_dep"] / 1.6021772e-12  # convert erg to eV
-        ionpopdict = {
-            at.get_ion_tuple(k): v
-            for k, v in estim.items()
-            if k.startswith("populations_") and "_" in k.removeprefix("populations_")
-        }
+        ionpopdict = {at.get_ion_tuple(k): v for k, v in estim.items() if k.startswith(("nnion_", "nnelement_"))}
 
         velocity = modeldata["vel_r_max_kmps"][args.modelgridindex]
         args.timedays = float(at.get_timestep_time(modelpath, args.timestep))
@@ -325,15 +321,15 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         if step == 0 and args.ostat:
             with Path(args.ostat).open("w") as fstat:
                 strheader = "#emin emax npts x_e frac_sum frac_excitation frac_ionization frac_heating"
-                for atomic_number, ion_stage in ions:
-                    strheader += " frac_ionization_" + at.get_ionstring(atomic_number, ion_stage, nospace=True)
+                for atomic_number, ionstage in ions:
+                    strheader += " frac_ionization_" + at.get_ionstring(atomic_number, ionstage, sep="")
                 fstat.write(strheader + "\n")
 
         with pynt.SpencerFanoSolver(emin_ev=emin, emax_ev=emax, npts=npts, verbose=True) as sf:
             for Z, ionstage in ions:
                 nnion = ionpopdict[(Z, ionstage)]
                 if nnion == 0.0:
-                    print(f"   skipping Z={Z} ion_stage {ionstage} due to nnion={nnion:.1e}")
+                    print(f"   skipping Z={Z} ionstage {ionstage} due to nnion={nnion:.1e}")
                     continue
 
                 sf.add_ionisation(Z, ionstage, nnion)
diff --git a/artistools/nonthermal/test_nonthermal.py b/artistools/nonthermal/test_nonthermal.py
new file mode 100644
index 000000000..7dbaebba3
--- /dev/null
+++ b/artistools/nonthermal/test_nonthermal.py
@@ -0,0 +1,14 @@
+import artistools as at
+
+modelpath = at.get_config()["path_testartismodel"]
+outputpath = at.get_config()["path_testoutput"]
+
+
+def test_nonthermal() -> None:
+    at.nonthermal.plot(argsraw=[], modelpath=modelpath, outputfile=outputpath, timestep=70)
+
+
+def test_spencerfano() -> None:
+    at.nonthermal.solvespencerfanocmd.main(
+        argsraw=[], modelpath=modelpath, timedays=300, makeplot=True, npts=200, noexcitation=True, outputfile=outputpath
+    )
diff --git a/artistools/plotspherical.py b/artistools/plotspherical.py
index a1d59430f..20200dfb9 100755
--- a/artistools/plotspherical.py
+++ b/artistools/plotspherical.py
@@ -27,7 +27,7 @@ def plot_spherical(
     ncosthetabins: int,
     maxpacketfiles: int | None = None,
     atomic_number: int | None = None,
-    ion_stage: int | None = None,
+    ionstage: int | None = None,
     gaussian_sigma: int | None = None,
     plotvars: list[str] | None = None,
     figscale: float = 1.0,
@@ -125,14 +125,14 @@ def plot_spherical(
         dfpackets = dfpackets.join(df_estimators, on=["em_timestep", "em_modelgridindex"], how="left")
         aggs.append(((pl.col("em_TR") * pl.col("e_rf")).mean() / pl.col("e_rf").mean()).alias("temperature"))
 
-    if atomic_number is not None or ion_stage is not None:
+    if atomic_number is not None or ionstage is not None:
         dflinelist = at.get_linelist_pldf(modelpath)
         if atomic_number is not None:
             print(f"Including only packets emitted by Z={atomic_number} {at.get_elsymbol(atomic_number)}")
             dflinelist = dflinelist.filter(pl.col("atomic_number") == atomic_number)
-        if ion_stage is not None:
-            print(f"Including only packets emitted by ionisation stage {ion_stage}")
-            dflinelist = dflinelist.filter(pl.col("ion_stage") == ion_stage)
+        if ionstage is not None:
+            print(f"Including only packets emitted by ionisation stage {ionstage}")
+            dflinelist = dflinelist.filter(pl.col("ionstage") == ionstage)
 
         selected_emtypes = dflinelist.select("lineindex").collect().get_column("lineindex")
         dfpackets = dfpackets.filter(pl.col("emissiontype").is_in(selected_emtypes))
@@ -256,7 +256,7 @@ def addargs(parser: argparse.ArgumentParser) -> None:
         "-atomic_number", type=int, default=None, help="Filter emitted packets by element of last emission"
     )
     parser.add_argument(
-        "-ion_stage", type=int, default=None, help="Filter emitted packets by ionistion stage of last emission"
+        "-ionstage", type=int, default=None, help="Filter emitted packets by ionistion stage of last emission"
     )
     parser.add_argument("-cmap", default=None, type=str, help="Matplotlib color map name")
 
@@ -324,7 +324,7 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non
             maxpacketfiles=args.maxpacketfiles,
             gaussian_sigma=args.gaussian_sigma,
             atomic_number=args.atomic_number,
-            ion_stage=args.ion_stage,
+            ionstage=args.ionstage,
             plotvars=args.plotvars,
             cmap=args.cmap,
             figscale=args.figscale,
diff --git a/artistools/radfield.py b/artistools/radfield.py
index 850d9bdb5..a15b0ffe7 100755
--- a/artistools/radfield.py
+++ b/artistools/radfield.py
@@ -9,8 +9,8 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
+import numpy.typing as npt
 import pandas as pd
-from astropy import units as u
 
 import artistools as at
 
@@ -21,6 +21,7 @@
 HOVERKB = 4.799243681748932e-11
 TWOOVERCLIGHTSQUARED = 2.2253001e-21
 SAHACONST = 2.0706659e-16
+MEGAPARSEC = 3.0857e24
 
 
 @lru_cache(maxsize=4)
@@ -102,7 +103,7 @@ def get_binaverage_field(radfielddata, modelgridindex=None, timestep=None):
     return arr_lambda, yvalues
 
 
-def j_nu_dbb(arr_nu_hz: t.Sequence[float], W: float, T: float) -> list[float]:
+def j_nu_dbb(arr_nu_hz: t.Sequence[float] | npt.NDArray, W: float, T: float) -> list[float]:
     """Calculate the spectral energy density of a dilute blackbody radiation field.
 
     Parameters
@@ -276,10 +277,10 @@ def plot_specout(
 
 @lru_cache(maxsize=128)
 def evaluate_phixs(
-    modelpath, atomic_number: int, lower_ion_stage: int, lowerlevelindex: int, nu_threshold: float, arr_nu_hz
+    modelpath, atomic_number: int, lower_ionstage: int, lowerlevelindex: int, nu_threshold: float, arr_nu_hz
 ):
     adata = at.atomic.get_levels(modelpath, get_photoionisations=True)
-    lower_ion_data = adata.query("Z == @atomic_number and ion_stage == @lower_ion_stage").iloc[0]
+    lower_ion_data = adata.query("Z == @atomic_number and ionstage == @lower_ionstage").iloc[0]
     lowerlevel = lower_ion_data.levels.iloc[lowerlevelindex]
 
     from scipy.interpolate import interp1d
@@ -308,15 +309,16 @@ def sigma_bf(nu):
     return np.array([sigma_bf(nu) for nu in arr_nu_hz])
 
 
-def get_kappa_bf_ion(atomic_number, lower_ion_stage, modelgridindex, timestep, modelpath, arr_nu_hz, max_levels):
+def get_kappa_bf_ion(atomic_number, lower_ionstage, modelgridindex, timestep, modelpath, arr_nu_hz, max_levels):
     adata = at.atomic.get_levels(modelpath, get_photoionisations=True)
     estimators = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex)
     T_e = estimators[(timestep, modelgridindex)]["Te"]
 
-    ion_data = adata.query("Z == @atomic_number and ion_stage == @lower_ion_stage").iloc[0]
-    upper_ion_data = adata.query("Z == @atomic_number and ion_stage == (@lower_ion_stage + 1)").iloc[0]
+    ion_data = adata.query("Z == @atomic_number and ionstage == @lower_ionstage").iloc[0]
+    upper_ion_data = adata.query("Z == @atomic_number and ionstage == (@lower_ionstage + 1)").iloc[0]
 
-    lowerionpopdensity = estimators[(timestep, modelgridindex)][f"populations_{atomic_number}_{lower_ion_stage}"]
+    ionstr = at.get_ionstring(atomic_number, lower_ionstage, sep="_", style="spectral")
+    lowerionpopdensity = estimators[(timestep, modelgridindex)][f"nnion_{ionstr}"]
 
     ion_popfactor_sum = sum(
         level.g * math.exp(-level.energy_ev * EV / KB / T_e) for _, level in ion_data.levels[:max_levels].iterrows()
@@ -331,7 +333,7 @@ def get_kappa_bf_ion(atomic_number, lower_ion_stage, modelgridindex, timestep, m
             nu_threshold = ONEOVERH * (ion_data.ion_pot - lowerlevel.energy_ev + upperlevel.energy_ev) * EV
 
             arr_sigma_bf = (
-                evaluate_phixs(modelpath, atomic_number, lower_ion_stage, levelnum, nu_threshold, tuple(arr_nu_hz))
+                evaluate_phixs(modelpath, atomic_number, lower_ionstage, levelnum, nu_threshold, tuple(arr_nu_hz))
                 * phixsfrac
             )
 
@@ -341,23 +343,26 @@ def get_kappa_bf_ion(atomic_number, lower_ion_stage, modelgridindex, timestep, m
 
 
 def get_recombination_emission(
-    atomic_number, upper_ion_stage, arr_nu_hz, modelgridindex, timestep, modelpath, max_levels, use_lte_pops=False
+    atomic_number, upper_ionstage, arr_nu_hz, modelgridindex, timestep, modelpath, max_levels, use_lte_pops=False
 ):
     adata = at.atomic.get_levels(modelpath, get_photoionisations=True)
 
-    lower_ion_stage = upper_ion_stage - 1
-    upperionstr = at.get_ionstring(atomic_number, upper_ion_stage)
-    lowerionstr = at.get_ionstring(atomic_number, lower_ion_stage)
-    upper_ion_data = adata.query("Z == @atomic_number and ion_stage == @upper_ion_stage").iloc[0]
-    lower_ion_data = adata.query("Z == @atomic_number and ion_stage == @lower_ion_stage").iloc[0]
+    lower_ionstage = upper_ionstage - 1
+    upperionstr = at.get_ionstring(atomic_number, upper_ionstage)
+    lowerionstr = at.get_ionstring(atomic_number, lower_ionstage)
+    upper_ion_data = adata.query("Z == @atomic_number and ionstage == @upper_ionstage").iloc[0]
+    lower_ion_data = adata.query("Z == @atomic_number and ionstage == @lower_ionstage").iloc[0]
 
-    estimators = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex)
+    estimtsmgi = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex)[
+        (timestep, modelgridindex)
+    ]
 
-    upperionpopdensity = estimators[(timestep, modelgridindex)][f"populations_{atomic_number}_{upper_ion_stage}"]
-    T_e = estimators[(timestep, modelgridindex)]["Te"]
-    nne = estimators[(timestep, modelgridindex)]["nne"]
+    upperionstr = at.get_ionstring(atomic_number, upper_ionstage, sep="_", style="spectral")
+    upperionpopdensity = estimtsmgi[f"nnion_{upperionstr}"]
+
+    T_e = estimtsmgi["Te"]
+    nne = estimtsmgi["nne"]
 
-    upperionpopdensity = estimators[(timestep, modelgridindex)][f"populations_{atomic_number}_{upper_ion_stage}"]
     print(f"Recombination from {upperionstr} -> {lowerionstr} ({upperionstr} pop = {upperionpopdensity:.1e}/cm3)")
 
     if use_lte_pops:
@@ -367,7 +372,7 @@ def get_recombination_emission(
         )
     else:
         dfnltepops = at.nltepops.read_files(modelpath, modelgridindex=modelgridindex, timestep=timestep)
-        dfnltepops_upperion = dfnltepops.query("Z==@atomic_number & ion_stage==@upper_ion_stage")
+        dfnltepops_upperion = dfnltepops.query("Z==@atomic_number & ionstage==@upper_ionstage")
         upperion_nltepops = {x.level: x["n_NLTE"] for _, x in dfnltepops_upperion.iterrows()}
 
     arr_j_nu_lowerlevel = {}
@@ -390,7 +395,7 @@ def get_recombination_emission(
             nu_threshold = ONEOVERH * (lower_ion_data.ion_pot - lowerlevel.energy_ev + upperlevel.energy_ev) * EV
 
             arr_sigma_bf = (
-                evaluate_phixs(modelpath, atomic_number, lower_ion_stage, levelnum, nu_threshold, tuple(arr_nu_hz))
+                evaluate_phixs(modelpath, atomic_number, lower_ionstage, levelnum, nu_threshold, tuple(arr_nu_hz))
                 * phixsfrac
             )
 
@@ -417,7 +422,7 @@ def get_recombination_emission(
             # arr_nu_hz2 = nu_threshold * lowerlevel.phixstable[:, 0]
             arr_nu_hz2 = nu_threshold * np.linspace(1.0, 1.0 + 0.03 * (100 + 1), num=3 * 100 + 1, endpoint=False)
             arr_sigma_bf2 = (
-                evaluate_phixs(modelpath, atomic_number, lower_ion_stage, levelnum, nu_threshold, tuple(arr_nu_hz2))
+                evaluate_phixs(modelpath, atomic_number, lower_ionstage, levelnum, nu_threshold, tuple(arr_nu_hz2))
                 * phixsfrac
             )
             arr_alpha_level_dnu2 = (
@@ -452,7 +457,7 @@ def get_recombination_emission(
     return arr_j_nu, arr_j_nu_lowerlevel
 
 
-def get_ion_gamma_dnu(modelpath, modelgridindex, timestep, atomic_number, ion_stage, arr_nu_hz, J_nu_arr, max_levels):
+def get_ion_gamma_dnu(modelpath, modelgridindex, timestep, atomic_number, ionstage, arr_nu_hz, J_nu_arr, max_levels):
     """Calculate the contribution to the photoionisation rate coefficient per J_nu at each frequency nu for an ion."""
     estimators = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex)
 
@@ -460,9 +465,9 @@ def get_ion_gamma_dnu(modelpath, modelgridindex, timestep, atomic_number, ion_st
     T_R = estimators[(timestep, modelgridindex)]["TR"]
 
     adata = at.atomic.get_levels(modelpath, get_photoionisations=True)
-    ion_data = adata.query("Z == @atomic_number and ion_stage == @ion_stage").iloc[0]
-    upper_ion_data = adata.query("Z == @atomic_number and ion_stage == (@ion_stage + 1)").iloc[0]
-    ionstr = at.get_ionstring(atomic_number, ion_stage)
+    ion_data = adata.query("Z == @atomic_number and ionstage == @ionstage").iloc[0]
+    upper_ion_data = adata.query("Z == @atomic_number and ionstage == (@ionstage + 1)").iloc[0]
+    ionstr = at.get_ionstring(atomic_number, ionstage)
 
     ion_popfactor_sum = sum(
         level.g * math.exp(-level.energy_ev * EV / KB / T_e) for _, level in ion_data.levels[:max_levels].iterrows()
@@ -476,8 +481,7 @@ def get_ion_gamma_dnu(modelpath, modelgridindex, timestep, atomic_number, ion_st
             nu_threshold = ONEOVERH * (ion_data.ion_pot - level.energy_ev + upperlevel.energy_ev) * EV
 
             arr_sigma_bf = (
-                evaluate_phixs(modelpath, atomic_number, ion_stage, levelnum, nu_threshold, tuple(arr_nu_hz))
-                * phixsfrac
+                evaluate_phixs(modelpath, atomic_number, ionstage, levelnum, nu_threshold, tuple(arr_nu_hz)) * phixsfrac
             )
 
             arr_corrfactors = 1 - np.exp(-HOVERKB * arr_nu_hz / T_R)
@@ -539,21 +543,21 @@ def calculate_photoionrates(axes, modelpath, radfielddata, modelgridindex, times
 
     # calculate bound-free opacity
     array_kappa_bf_nu = np.zeros_like(arr_nu_hz_recomb)
-    for atomic_number, lower_ion_stage in kappalowerionlist:
+    for atomic_number, lower_ionstage in kappalowerionlist:
         array_kappa_bf_nu += get_kappa_bf_ion(
-            atomic_number, lower_ion_stage, modelgridindex, timestep, modelpath, arr_nu_hz_recomb, max_levels
+            atomic_number, lower_ionstage, modelgridindex, timestep, modelpath, arr_nu_hz_recomb, max_levels
         )
 
     # calculate recombination emission
     J_lambda_recomb_total = np.zeros_like(arraylambda_angstrom_recomb)
 
     lw = 1.0
-    for atomic_number, lower_ion_stage in recomblowerionlist:
+    for atomic_number, lower_ionstage in recomblowerionlist:
         # lw -= 0.1
-        upperionstr = at.get_ionstring(atomic_number, lower_ion_stage + 1)
+        upperionstr = at.get_ionstring(atomic_number, lower_ionstage + 1)
 
         j_emiss_nu_recomb, arr_j_nu_lowerleveldict = get_recombination_emission(
-            atomic_number, lower_ion_stage + 1, arr_nu_hz_recomb, modelgridindex, timestep, modelpath, max_levels
+            atomic_number, lower_ionstage + 1, arr_nu_hz_recomb, modelgridindex, timestep, modelpath, max_levels
         )
 
         for (upperlevelnum, lowerlevelnum), arr_j_emiss_nu_lowerlevel in arr_j_nu_lowerleveldict.items():
@@ -568,7 +572,7 @@ def calculate_photoionrates(axes, modelpath, radfielddata, modelgridindex, times
 
             J_lambda_recomb_level = J_nu_recomb * arr_nu_hz_recomb / arraylambda_angstrom_recomb
             fieldlabel = (
-                f"{upperionstr} level {upperlevelnum} -> {at.roman_numerals[lower_ion_stage]} level {lowerlevelnum}"
+                f"{upperionstr} level {upperlevelnum} -> {at.roman_numerals[lower_ionstage]} level {lowerlevelnum}"
             )
             axes[2].plot(arraylambda_angstrom_recomb, J_lambda_recomb_level, label=fieldlabel, lw=lw)
             fieldlist += [(arraylambda_angstrom_recomb, J_lambda_recomb_level, fieldlabel)]
@@ -590,7 +594,7 @@ def calculate_photoionrates(axes, modelpath, radfielddata, modelgridindex, times
         J_lambda_recomb_total += J_lambda_recomb
 
         # contribution of all levels of the ion
-        fieldlabel = f"{upperionstr} -> {at.roman_numerals[lower_ion_stage]} recombination"
+        fieldlabel = f"{upperionstr} -> {at.roman_numerals[lower_ionstage]} recombination"
         axes[2].plot(arraylambda_angstrom_recomb, J_lambda_recomb, label=fieldlabel, lw=lw)
         fieldlist += [(arraylambda_angstrom_recomb, J_lambda_recomb, fieldlabel)]
 
@@ -629,14 +633,14 @@ def calculate_photoionrates(axes, modelpath, radfielddata, modelgridindex, times
     lw = 1.0
     # fieldlist += [(arr_lambda_fitted, j_lambda_fitted, 'binned field')]
 
-    for atomic_number, ion_stage in photoionlist:
-        ionstr = at.get_ionstring(atomic_number, ion_stage)
-        ion_data = adata.query("Z == @atomic_number and ion_stage == @ion_stage").iloc[0]
+    for atomic_number, ionstage in photoionlist:
+        ionstr = at.get_ionstring(atomic_number, ionstage)
+        ion_data = adata.query("Z == @atomic_number and ionstage == @ionstage").iloc[0]
 
         for levelnum, level in ion_data.levels[:max_levels].iterrows():
             nu_threshold = ONEOVERH * (ion_data.ion_pot - level.energy_ev) * EV
             arr_sigma_bf = evaluate_phixs(
-                modelpath, atomic_number, ion_stage, levelnum, nu_threshold, tuple(arr_nu_hz_recomb)
+                modelpath, atomic_number, ionstage, levelnum, nu_threshold, tuple(arr_nu_hz_recomb)
             )
             if levelnum < 5:
                 axes[0].plot(
@@ -650,7 +654,7 @@ def calculate_photoionrates(axes, modelpath, radfielddata, modelgridindex, times
             J_nu_arr = np.array(J_lambda_arr) * arraylambda_angstrom / arr_nu_hz
 
             arr_gamma_dnu = get_ion_gamma_dnu(
-                modelpath, modelgridindex, timestep, atomic_number, ion_stage, arr_nu_hz, J_nu_arr, max_levels
+                modelpath, modelgridindex, timestep, atomic_number, ionstage, arr_nu_hz, J_nu_arr, max_levels
             )
 
             # xlist = arr_lambda_fitted
@@ -740,13 +744,13 @@ def plot_celltimestep(modelpath, timestep, outputfile, xmin, xmax, modelgridinde
         if not normalised:
             modeldata, _, t_model_init = at.inputmodel.get_modeldata_tuple(modelpath)
             # outer velocity
-            v_surface = modeldata.loc[int(radfielddata.modelgridindex.max())].vel_r_max_kmps * u.km / u.s
-            r_surface = (time_days * u.day * v_surface).to("km")
-            r_observer = u.megaparsec.to("km")
+            v_surface = modeldata.loc[int(radfielddata.modelgridindex.max())].vel_r_max_kmps * 1e5
+            r_surface = time_days * 864000 * v_surface
+            r_observer = MEGAPARSEC
             scale_factor = (r_observer / r_surface) ** 2 / (2 * math.pi)
             print(
                 "Scaling emergent spectrum flux at 1 Mpc to specific intensity "
-                f"at surface (v={v_surface:.3e}, r={r_surface:.3e})"
+                f"at surface (v={v_surface:.3e}, r={r_surface:.3e} {r_observer:.3e}) scale_factor: {scale_factor:.3e}"
             )
             plotkwargs["scale_factor"] = scale_factor
         else:
@@ -1014,7 +1018,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
     else:
         modelgridindexlist = at.parse_range_list(args.modelgridindex)
 
-    timesteplast = len(at.get_timestep_times(modelpath))
+    timesteplast = len(at.get_timestep_times(modelpath)) - 1
     if args.timedays:
         timesteplist = [at.get_timestep_of_timedays(modelpath, args.timedays)]
     elif args.timestep:
diff --git a/artistools/spectra/sampleblackbodyfrompacket_tr.py b/artistools/spectra/sampleblackbodyfrompacket_tr.py
index 868e652d4..76945ace9 100644
--- a/artistools/spectra/sampleblackbodyfrompacket_tr.py
+++ b/artistools/spectra/sampleblackbodyfrompacket_tr.py
@@ -12,7 +12,7 @@
 DAY = 86400
 TWOHOVERCLIGHTSQUARED = 1.4745007e-47
 HOVERKB = 4.799243681748932e-11
-PARSEC = 3.086e18
+PARSEC = 3.0857e18
 c_cgs = const.c.to("cm/s").value
 c_ang_s = const.c.to("angstrom/s").value
 
diff --git a/artistools/spectra/spectra.py b/artistools/spectra/spectra.py
index 49dc1bbc4..8a8925352 100644
--- a/artistools/spectra/spectra.py
+++ b/artistools/spectra/spectra.py
@@ -1171,7 +1171,7 @@ def print_integrated_flux(
         f" integrated flux ({arr_lambda_angstroms.min():.1f} to "
         f"{arr_lambda_angstroms.max():.1f} A): {integrated_flux:.3e} erg/s/cm2"
     )
-    return integrated_flux.value
+    return integrated_flux
 
 
 def get_reference_spectrum(filename: Path | str) -> tuple[pd.DataFrame, dict[t.Any, t.Any]]:
diff --git a/artistools/test_artistools.py b/artistools/test_artistools.py
index c173894bf..a7aeab135 100755
--- a/artistools/test_artistools.py
+++ b/artistools/test_artistools.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 import hashlib
 import importlib
+import inspect
 import math
 import typing as t
 
@@ -11,6 +12,15 @@
 outputpath = at.get_config()["path_testoutput"]
 
 
+def funcname() -> str:
+    """Get the name of the calling function."""
+    try:
+        return inspect.currentframe().f_back.f_code.co_name  # type: ignore[union-attr]
+    except AttributeError as e:
+        msg = "Could not get the name of the calling function."
+        raise RuntimeError(msg) from e
+
+
 def test_commands() -> None:
     # ensure that the commands are pointing to valid submodule.function() targets
     for _command, (submodulename, funcname) in sorted(at.commands.get_commandlist().items()):
@@ -76,12 +86,10 @@ def test_nltepops() -> None:
     at.nltepops.plot(argsraw=[], modelpath=modelpath, outputfile=outputpath, timestep=40)
 
 
-def test_nonthermal() -> None:
-    at.nonthermal.plot(argsraw=[], modelpath=modelpath, outputfile=outputpath, timestep=70)
-
-
 def test_radfield() -> None:
-    at.radfield.main(argsraw=[], modelpath=modelpath, modelgridindex=0, outputfile=outputpath)
+    funcoutpath = outputpath / funcname()
+    funcoutpath.mkdir(exist_ok=True)
+    at.radfield.main(argsraw=[], modelpath=modelpath, modelgridindex=0, outputfile=funcoutpath)
 
 
 def test_get_ionrecombratecalibration() -> None:
@@ -89,19 +97,15 @@ def test_get_ionrecombratecalibration() -> None:
 
 
 def test_plotspherical() -> None:
-    at.plotspherical.main(argsraw=[], modelpath=modelpath, outputfile=outputpath)
+    funcoutpath = outputpath / funcname()
+    funcoutpath.mkdir(exist_ok=True)
+    at.plotspherical.main(argsraw=[], modelpath=modelpath, outputfile=funcoutpath)
 
 
 def test_plotspherical_gif() -> None:
     at.plotspherical.main(argsraw=[], modelpath=modelpath, makegif=True, timemax=270, outputfile=outputpath)
 
 
-def test_spencerfano() -> None:
-    at.nonthermal.solvespencerfanocmd.main(
-        argsraw=[], modelpath=modelpath, timedays=300, makeplot=True, npts=200, noexcitation=True, outputfile=outputpath
-    )
-
-
 def test_transitions() -> None:
     at.transitions.main(argsraw=[], modelpath=modelpath, outputfile=outputpath, timedays=300)
 
diff --git a/artistools/transitions.py b/artistools/transitions.py
index fac24a1a7..f12a8b1a5 100755
--- a/artistools/transitions.py
+++ b/artistools/transitions.py
@@ -17,7 +17,7 @@
 
 defaultoutputfile = "plottransitions_cell{cell:03d}_ts{timestep:02d}_{time_days:.0f}d.pdf"
 
-iontuple = namedtuple("iontuple", "Z ion_stage")
+iontuple = namedtuple("iontuple", "Z ionstage")
 
 
 def get_kurucz_transitions() -> tuple[pd.DataFrame, list[iontuple]]:
@@ -142,8 +142,8 @@ def make_plot(
             peak_y_value = max(peak_y_value, **yvalues_combined[seriesindex])
 
     axislabels = [
-        f"{at.get_elsymbol(Z)} {at.roman_numerals[ion_stage]}\n(pop={ionpopdict[(Z, ion_stage)]:.1e}/cm3)"
-        for (Z, ion_stage) in ionlist
+        f"{at.get_elsymbol(Z)} {at.roman_numerals[ionstage]}\n(pop={ionpopdict[(Z, ionstage)]:.1e}/cm3)"
+        for (Z, ionstage) in ionlist
     ]
     axislabels += ["Total"]
 
@@ -307,8 +307,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             sys.exit(1)
 
         ionpopdict = {
-            (Z, ion_stage): dfnltepops.query("Z==@Z and ion_stage==@ion_stage")["n_NLTE"].sum()
-            for Z, ion_stage in ionlist
+            (Z, ionstage): dfnltepops.query("Z==@Z and ionstage==@ionstage")["n_NLTE"].sum() for Z, ionstage in ionlist
         }
 
         modelname = at.get_model_name(modelpath)
@@ -357,21 +356,21 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
 
     fe2depcoeff, ni2depcoeff = None, None
     for _, ion in adata.iterrows() if args.atomicdatabase == "artis" else enumerate(ionlist):
-        ionid = iontuple(ion.Z, ion.ion_stage)
+        ionid = iontuple(ion.Z, ion.ionstage)
         if ionid not in ionlist:
             continue
 
         ionindex = ionlist.index(ionid)
 
         if args.atomicdatabase == "kurucz":
-            dftransitions = dftransgfall.query("Z == @ion.Z and ionstage == @ion.ion_stage", inplace=False).copy()
+            dftransitions = dftransgfall.query("Z == @ion.Z and ionstage == @ion.ionstage", inplace=False).copy()
         elif args.atomicdatabase == "nist":
-            dftransitions = get_nist_transitions(f"nist/nist-{ion.Z:02d}-{ion.ion_stage:02d}.txt")
+            dftransitions = get_nist_transitions(f"nist/nist-{ion.Z:02d}-{ion.ionstage:02d}.txt")
         else:
             dftransitions = ion.transitions
 
         print(
-            f"\n======> {at.get_elsymbol(ion.Z)} {at.roman_numerals[ion.ion_stage]:3s} "
+            f"\n======> {at.get_elsymbol(ion.Z)} {at.roman_numerals[ion.ionstage]:3s} "
             f"(pop={ionpopdict[ionid]:.2e} / cm3, {len(dftransitions):6d} transitions)"
         )
 
@@ -408,7 +407,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
 
             for seriesindex, temperature in enumerate(temperature_list):
                 if temperature == "NOTEMPNLTE":
-                    dfnltepops_thision = dfnltepops.query("Z==@ion.Z & ion_stage==@ion.ion_stage")
+                    dfnltepops_thision = dfnltepops.query("Z==@ion.Z & ionstage==@ion.ionstage")
 
                     nltepopdict = {x.level: x["n_NLTE"] for _, x in dfnltepops_thision.iterrows()}
 
@@ -465,8 +464,10 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         feions = [2, 3]
 
         def get_strionfracs(atomic_number, ionstages):
+            elsym = at.get_elsymbol(atomic_number)
             est_ionfracs = [
-                estimators[f"populations_{atomic_number}_{ionstage}"] / estimators[f"populations_{atomic_number}"]
+                estimators[f"nnion_{at.get_ionstring(atomic_number, ionstage, sep='_', style='spectral')}"]
+                / estimators[f"nnelement_{elsym}"]
                 for ionstage in ionstages
             ]
             ionfracs_str = " ".join([f"{pop:6.0e}" if pop < 0.01 else f"{pop:6.2f}" for pop in est_ionfracs])
@@ -488,8 +489,8 @@ def get_strionfracs(atomic_number, ionstages):
             f"{velocity:5.0f} km/s({modelgridindex})      {fe2depcoeff:5.2f}                   "
             f"{ni2depcoeff:.2f}        "
             f"{est_fe_ionfracs_str}   /  {est_ni_ionfracs_str}      {Te:.0f}    "
-            f"{estimators['populations_26_3'] / estimators['populations_26_2']:.2f}          "
-            f"{estimators['populations_28_3'] / estimators['populations_28_2']:5.2f}"
+            f"{estimators['nnion_Fe_III'] / estimators['nnion_Fe_II']:.2f}          "
+            f"{estimators['nnion_Ni_III'] / estimators['nnion_Ni_II']:5.2f}"
         )
 
     outputfilename = (
diff --git a/artistools/writecomparisondata.py b/artistools/writecomparisondata.py
index 0b512b950..283b46fa0 100755
--- a/artistools/writecomparisondata.py
+++ b/artistools/writecomparisondata.py
@@ -96,9 +96,9 @@ def write_ionfracts(
     nelements = len(elementlist)
     for element in range(nelements):
         atomic_number = elementlist.Z[element]
-        elsymb = at.get_elsymbol(atomic_number).lower()
+        elsymb = at.get_elsymbol(atomic_number)
         nions = elementlist.nions[element]
-        pathfileout = Path(outputpath, f"ionfrac_{elsymb}_{model_id}_artisnebular.txt")
+        pathfileout = Path(outputpath, f"ionfrac_{elsymb.lower()}_{model_id}_artisnebular.txt")
         fileisallzeros = True  # will be changed when a non-zero is encountered
         with pathfileout.open("w") as f:
             f.write(f"#NTIMES: {len(selected_timesteps)}\n")
@@ -108,18 +108,17 @@ def write_ionfracts(
             for timestep in selected_timesteps:
                 f.write(f"#TIME: {times[timestep]:.2f}\n")
                 f.write(f"#NVEL: {len(allnonemptymgilist)}\n")
-                f.write(f'#vel_mid[km/s] {" ".join([f"{elsymb}{ion}" for ion in range(nions)])}\n')
+                f.write(f'#vel_mid[km/s] {" ".join([f"{elsymb.lower()}{ion}" for ion in range(nions)])}\n')
                 for modelgridindex, cell in modeldata.iterrows():
                     if modelgridindex not in allnonemptymgilist:
                         continue
                     v_mid = (cell.vel_r_min_kmps + cell.vel_r_max_kmps) / 2.0
                     f.write(f"{v_mid:.2f}")
-                    elabund = estimators[(timestep, modelgridindex)].get(f"populations_{atomic_number}", 0)
+                    elabund = estimators[(timestep, modelgridindex)].get(f"nnelement_{elsymb}", 0)
                     for ion in range(nions):
-                        ion_stage = ion + elementlist.lowermost_ionstage[element]
-                        ionabund = estimators[(timestep, modelgridindex)].get(
-                            f"populations_{atomic_number}_{ion_stage}", 0
-                        )
+                        ionstage = ion + elementlist.lowermost_ionstage[element]
+                        ionstr = at.get_ionstring(atomic_number, ionstage, sep="_", style="spectral")
+                        ionabund = estimators[(timestep, modelgridindex)].get(f"nnion_{ionstr}", 0)
                         ionfrac = ionabund / elabund if elabund > 0 else 0
                         if ionfrac > 0.0:
                             fileisallzeros = False
@@ -149,9 +148,7 @@ def write_phys(modelpath, model_id, selected_timesteps, estimators, allnonemptym
                     10**cell.logrho * (modelmeta["t_model_init_days"] / times[timestep]) ** 3
                 )
 
-                estimators[(timestep, modelgridindex)]["nntot"] = estimators[(timestep, modelgridindex)][
-                    "populations_total"
-                ]
+                estimators[(timestep, modelgridindex)]["nntot"] = estimators[(timestep, modelgridindex)]["nntot"]
 
                 v_mid = (cell.vel_r_min_kmps + cell.vel_r_max_kmps) / 2.0
                 f.write(f"{v_mid:.2f}")
diff --git a/pyproject.toml b/pyproject.toml
index 4027873ae..6757f7730 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -215,6 +215,8 @@ show-fixes = true
 "matplotlib.typing" = "mplt"
 "numpy.typing" = "npt"
 "typing" = "t"
+"polars" = "pl"
+"polars.selectors" = "cs"
 
 [tool.ruff.flake8-tidy-imports]
 ban-relative-imports = "all"

From e3edc4b977b4737ea2f3dab3c9a919f2a732e7f4 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 11:03:33 +0000
Subject: [PATCH 021/150] Update .pre-commit-config.yaml

---
 .pre-commit-config.yaml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 5b44c3944..a12ca6892 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -31,13 +31,13 @@ repos:
     #       - id: yamlfmt
 
     - repo: https://github.com/astral-sh/ruff-pre-commit
-      rev: v0.1.4
+      rev: v0.1.5
       hooks:
         - id: ruff
           args: [--fix, --exit-non-zero-on-fix]
 
     - repo: https://github.com/astral-sh/ruff-pre-commit
-      rev: v0.1.4
+      rev: v0.1.5
       hooks:
         - id: ruff-format
 

From 59d3057f5759177e112e72f3fe69a4e9d4e56e60 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 11:03:42 +0000
Subject: [PATCH 022/150] Update requirements.txt

---
 requirements.txt | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index f5e6cc9c4..f835030f4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,17 +10,17 @@ numpy>=1.26.1
 pandas>=2.1.2
 polars>=0.19.12
 pre-commit>=3.5.0
-pyarrow>=14.0.0
+pyarrow>=14.0.1
 pynonthermal>=2021.10.12
 pypdf2>=3.0.1
-pyright>=1.1.334
+pyright>=1.1.335
 pytest>=7.4.3
 pytest-cov>=4.1.0
 python-xz>=0.5
 pyvista>=0.42.3
 PyYAML>=6.0.1
 pyzstd>=0.15.9
-ruff>=0.1.4
+ruff>=0.1.5
 scipy>=1.11.3
 setuptools_scm[toml]>=8.0.4
 typeguard>=4.1.5

From 2a5a4e05513dfd0d1979be210613b5428a9cc5c5 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 11:13:08 +0000
Subject: [PATCH 023/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 4fb6b2bfc..88c7585c5 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -212,7 +212,7 @@ def plot_levelpop(
         ionlevels = adata.query("Z == @atomic_number and ionstage == @ionstage").iloc[0].levels
         levelname = ionlevels.iloc[levelindex].levelname
         label = (
-            f"{at.get_ionstring(atomic_number, ionstage, style="chargelatex")} level {levelindex}:"
+            f"{at.get_ionstring(atomic_number, ionstage, style='chargelatex')} level {levelindex}:"
             f" {at.nltepops.texifyconfiguration(levelname)}"
         )
 

From 29fb6493732c4b7dfaa78b54b1553c4d4fd42afb Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 11:17:41 +0000
Subject: [PATCH 024/150] Remove lz4 support

---
 artistools/misc.py            | 9 ++++-----
 artistools/packets/packets.py | 2 +-
 requirements.txt              | 1 -
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/artistools/misc.py b/artistools/misc.py
index 4b6f3af42..9514ac7bb 100644
--- a/artistools/misc.py
+++ b/artistools/misc.py
@@ -9,7 +9,6 @@
 from itertools import chain
 from pathlib import Path
 
-import lz4.frame
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
@@ -773,9 +772,9 @@ def zopen(filename: Path | str, mode: str = "rt", encoding: str | None = None, f
     """
     if forpolars:
         mode = "r"
-    ext_fopen = [(".lz4", lz4.frame.open), (".zst", pyzstd.open), (".gz", gzip.open), (".xz", xz.open)]
+    ext_fopen: dict[str, t.Callable] = {".zst": pyzstd.open, ".gz": gzip.open, ".xz": xz.open}
 
-    for ext, fopen in ext_fopen:
+    for ext, fopen in ext_fopen.items():
         file_ext = str(filename) if str(filename).endswith(ext) else str(filename) + ext
         if Path(file_ext).exists():
             return fopen(file_ext, mode=mode, encoding=encoding)
@@ -801,7 +800,7 @@ def firstexisting(
         fullpaths.append(Path(folder) / filename)
 
         if tryzipped:
-            for ext in [".lz4", ".zst", ".gz", ".xz"]:
+            for ext in [".zst", ".gz", ".xz"]:
                 filenameext = str(filename) if str(filename).endswith(ext) else str(filename) + ext
                 if filenameext not in filelist:
                     fullpaths.append(folder / filenameext)
@@ -865,7 +864,7 @@ def add_derived_metadata(metadata: dict[str, t.Any]) -> dict[str, t.Any]:
 
     import yaml
 
-    filepath = Path(str(filepath).replace(".xz", "").replace(".gz", "").replace(".lz4", "").replace(".zst", ""))
+    filepath = Path(str(filepath).replace(".xz", "").replace(".gz", "").replace(".zst", ""))
 
     # check if the reference file (e.g. spectrum.txt) has an metadata file (spectrum.txt.meta.yml)
     individualmetafile = filepath.with_suffix(f"{filepath.suffix}.meta.yml")
diff --git a/artistools/packets/packets.py b/artistools/packets/packets.py
index d253ace00..e2affe451 100644
--- a/artistools/packets/packets.py
+++ b/artistools/packets/packets.py
@@ -431,7 +431,7 @@ def get_packetsfilepaths(
 
     searchfolders = [Path(modelpath, "packets"), Path(modelpath)]
     # in descending priority (based on speed of reading)
-    suffix_priority = [".out.zst", ".out.lz4", ".out.zst", ".out", ".out.gz", ".out.xz"]
+    suffix_priority = [".out.zst", ".out.zst", ".out", ".out.gz", ".out.xz"]
     t_lastschemachange = calendar.timegm(time_parquetschemachange)
 
     parquetpacketsfiles = []
diff --git a/requirements.txt b/requirements.txt
index f835030f4..0f2c5d68e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,7 +3,6 @@ astropy>=5.3.4
 coverage>=7.3.2
 extinction>=0.4.6
 imageio>=2.32.0
-lz4>=4.3.2
 matplotlib>=3.8.1
 mypy>=1.6.1
 numpy>=1.26.1

From e47f0ba87fb24131d6360bfa1d4ad1749239b138 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 12:27:07 +0000
Subject: [PATCH 025/150] Store estimators in parquet files

---
 artistools/estimators/__init__.py       |   1 -
 artistools/estimators/estimators.py     | 176 +++++++++++++-----------
 artistools/linefluxes.py                |   4 +-
 artistools/nonthermal/plotnonthermal.py |   6 +-
 4 files changed, 99 insertions(+), 88 deletions(-)

diff --git a/artistools/estimators/__init__.py b/artistools/estimators/__init__.py
index fdc8b7f1d..0f74b5b7b 100644
--- a/artistools/estimators/__init__.py
+++ b/artistools/estimators/__init__.py
@@ -17,6 +17,5 @@
 from artistools.estimators.estimators import parse_estimfile
 from artistools.estimators.estimators import read_estimators
 from artistools.estimators.estimators import read_estimators_from_file
-from artistools.estimators.estimators import read_estimators_polars
 from artistools.estimators.plotestimators import addargs
 from artistools.estimators.plotestimators import main as plot
diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 45ff7d6bd..c3ba6a2f9 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -119,7 +119,6 @@ def get_units_string(variable: str) -> str:
 
 def parse_estimfile(
     estfilepath: Path | str,
-    get_ion_values: bool = True,
     skip_emptycells: bool = False,
 ) -> t.Iterator[tuple[int, int, dict[str, t.Any]]]:  # pylint: disable=unused-argument
     """Generate timestep, modelgridindex, dict from estimator file."""
@@ -157,7 +156,7 @@ def parse_estimfile(
                         estimblock[variablename] = float(value)
                     estimblock["lognne"] = math.log10(estimblock["nne"]) if estimblock["nne"] > 0 else float("-inf")
 
-            elif row[1].startswith("Z=") and get_ion_values:
+            elif row[1].startswith("Z="):
                 variablename = row[0]
                 if row[1].endswith("="):
                     atomic_number = int(row[2])
@@ -230,33 +229,80 @@ def parse_estimfile(
 
 def read_estimators_from_file(
     estfilepath: Path | str,
-    modelpath: Path,
     printfilename: bool = False,
-    get_ion_values: bool = True,
     skip_emptycells: bool = False,
 ) -> dict[tuple[int, int], dict[str, t.Any]]:
     if printfilename:
         estfilepath = Path(estfilepath)
         filesize = estfilepath.stat().st_size / 1024 / 1024
-        print(f"Reading {estfilepath.relative_to(modelpath.parent)} ({filesize:.2f} MiB)")
+        print(f"Reading {estfilepath.relative_to(estfilepath.parent.parent)} ({filesize:.2f} MiB)")
 
     return {
         (timestep, mgi): file_estimblock
         for timestep, mgi, file_estimblock in parse_estimfile(
             estfilepath,
-            get_ion_values=get_ion_values,
             skip_emptycells=skip_emptycells,
         )
     }
 
 
+def get_estimators_in_folder(
+    modelpath: Path,
+    folderpath: Path,
+    match_modelgridindex: None | t.Sequence[int],
+    skip_emptycells: bool,
+) -> dict:
+    mpiranklist = at.get_mpiranklist(modelpath, modelgridindex=match_modelgridindex, only_ranks_withgridcells=True)
+    printfilename = len(mpiranklist) < 10
+
+    estfilepaths = []
+    for mpirank in mpiranklist:
+        # not worth printing an error, because ranks with no cells to update do not produce an estimator file
+        with contextlib.suppress(FileNotFoundError):
+            estfilepath = at.firstexisting(f"estimators_{mpirank:04d}.out", folder=folderpath, tryzipped=True)
+            estfilepaths.append(estfilepath)
+
+    if not printfilename:
+        print(f"Reading {len(list(estfilepaths))} estimator files in {folderpath.relative_to(Path(folderpath).parent)}")
+
+    processfile = partial(
+        read_estimators_from_file,
+        printfilename=printfilename,
+        skip_emptycells=skip_emptycells,
+    )
+
+    if at.get_config()["num_processes"] > 1:
+        with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
+            arr_rankestimators = pool.map(processfile, estfilepaths)
+            pool.close()
+            pool.join()
+            pool.terminate()
+    else:
+        arr_rankestimators = [processfile(estfilepath) for estfilepath in estfilepaths]
+
+    estimators: dict[tuple[int, int], dict[str, t.Any]] = {}
+    for estfilepath, estimators_thisfile in zip(estfilepaths, arr_rankestimators):
+        dupekeys = sorted([k for k in estimators_thisfile if k in estimators])
+        for k in dupekeys:
+            # dropping the lowest timestep is normal for restarts. Only warn about other cases
+            if k[0] != dupekeys[0][0]:
+                print(
+                    f"WARNING: Duplicate estimator block for (timestep, mgi) key {k}. "
+                    f"Dropping block from {estfilepath}"
+                )
+
+            del estimators_thisfile[k]
+
+        estimators |= estimators_thisfile
+
+    return estimators
+
+
 def read_estimators(
     modelpath: Path | str = Path(),
     modelgridindex: None | int | t.Sequence[int] = None,
     timestep: None | int | t.Sequence[int] = None,
-    mpirank: int | None = None,
     runfolder: None | str | Path = None,
-    get_ion_values: bool = True,
     skip_emptycells: bool = False,
     add_velocity: bool = True,
 ) -> dict[tuple[int, int], dict[str, t.Any]]:
@@ -286,65 +332,51 @@ def read_estimators(
 
     # print(f" matching cells {match_modelgridindex} and timesteps {match_timestep}")
 
-    mpiranklist = (
-        at.get_mpiranklist(modelpath, modelgridindex=match_modelgridindex, only_ranks_withgridcells=True)
-        if mpirank is None
-        else [mpirank]
-    )
-
     runfolders = at.get_runfolders(modelpath, timesteps=match_timestep) if runfolder is None else [Path(runfolder)]
 
-    printfilename = len(mpiranklist) < 10
+    estimators: dict[tuple[int, int], dict[str, t.Any]] = {}
+    parquetfiles = [folderpath / "estimators.out.parquet.tmp" for folderpath in runfolders]
 
-    estimators: dict[tuple[int, int], dict] = {}
-    for folderpath in runfolders:
-        estfilepaths = []
-        for mpirank in mpiranklist:
-            # not worth printing an error, because ranks with no cells to update do not produce an estimator file
-            with contextlib.suppress(FileNotFoundError):
-                estfilepath = at.firstexisting(f"estimators_{mpirank:04d}.out", folder=folderpath, tryzipped=True)
-                estfilepaths.append(estfilepath)
-
-        if not printfilename:
-            print(
-                f"Reading {len(list(estfilepaths))} estimator files in {folderpath.relative_to(Path(modelpath).parent)}"
-            )
-
-        processfile = partial(
-            read_estimators_from_file,
-            modelpath=modelpath,
-            get_ion_values=get_ion_values,
-            printfilename=printfilename,
-            skip_emptycells=skip_emptycells,
+    for folderpath, parquetfile in zip(runfolders, parquetfiles):
+        if parquetfile.exists():
+            continue
+
+        estim_folder = get_estimators_in_folder(
+            modelpath,
+            folderpath,
+            match_modelgridindex=None,
+            skip_emptycells=True,
+        )
+
+        pldf = pl.DataFrame(
+            [
+                {
+                    "timestep": ts,
+                    "modelgridindex": mgi,
+                    **estimvals,
+                }
+                for (ts, mgi), estimvals in estim_folder.items()
+                if not estimvals.get("emptycell", True)
+            ]
         )
+        print(f"Writing {parquetfile.relative_to(modelpath.parent)}")
+        pldf.write_parquet(parquetfile, compression="zstd")
+
+    for folderpath, parquetfile in zip(runfolders, parquetfiles):
+        print(f"Reading {parquetfile.relative_to(modelpath.parent)}")
 
-        if at.get_config()["num_processes"] > 1:
-            with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
-                arr_rankestimators = pool.map(processfile, estfilepaths)
-                pool.close()
-                pool.join()
-                pool.terminate()
-        else:
-            arr_rankestimators = [processfile(estfilepath) for estfilepath in estfilepaths]
-
-        for estfilepath, estimators_thisfile in zip(estfilepaths, arr_rankestimators):
-            dupekeys = sorted([k for k in estimators_thisfile if k in estimators])
-            for k in dupekeys:
-                # dropping the lowest timestep is normal for restarts. Only warn about other cases
-                if k[0] != dupekeys[0][0]:
-                    print(
-                        f"WARNING: Duplicate estimator block for (timestep, mgi) key {k}. "
-                        f"Dropping block from {estfilepath}"
-                    )
-
-                del estimators_thisfile[k]
-
-            estimators |= {
-                (ts, mgi): v
-                for (ts, mgi), v in estimators_thisfile.items()
-                if (not match_modelgridindex or mgi in match_modelgridindex)
-                and (not match_timestep or ts in match_timestep)
-            }
+    pldflazy = pl.scan_parquet(parquetfiles)
+
+    if match_modelgridindex is not None:
+        pldflazy = pldflazy.filter(pl.col("modelgridindex").is_in(match_modelgridindex))
+
+    if match_timestep is not None:
+        pldflazy = pldflazy.filter(pl.col("timestep").is_in(match_timestep))
+
+    for estimtsmgi in pldflazy.collect().iter_rows(named=True):
+        estimators[(estimtsmgi["timestep"], estimtsmgi["modelgridindex"])] = {
+            k: v for k, v in estimtsmgi.items() if k not in {"timestep", "modelgridindex"} and v is not None
+        }
 
     return estimators
 
@@ -458,7 +490,6 @@ def get_temperatures(modelpath: str | Path) -> pl.LazyFrame:
     if not dfest_parquetfile.is_file():
         estimators = at.estimators.read_estimators(
             modelpath,
-            get_ion_values=False,
             skip_emptycells=True,
         )
         assert len(estimators) > 0
@@ -473,22 +504,3 @@ def get_temperatures(modelpath: str | Path) -> pl.LazyFrame:
         )
 
     return pl.scan_parquet(dfest_parquetfile)
-
-
-def read_estimators_polars(*args, **kwargs) -> pl.LazyFrame:
-    estimators = read_estimators(*args, **kwargs)
-    pldf = pl.DataFrame(
-        [
-            {
-                "timestep": ts,
-                "modelgridindex": mgi,
-                **estimvals,
-            }
-            for (ts, mgi), estimvals in estimators.items()
-            if not estimvals.get("emptycell", True)
-        ]
-    )
-    print(pldf.columns)
-    print(pldf.transpose(include_header=True))
-
-    return pldf.lazy()
diff --git a/artistools/linefluxes.py b/artistools/linefluxes.py
index df0161898..0f40bc942 100755
--- a/artistools/linefluxes.py
+++ b/artistools/linefluxes.py
@@ -455,7 +455,7 @@ def get_packets_with_emission_conditions(
     tend: float,
     maxpacketfiles: int | None = None,
 ) -> pd.DataFrame:
-    estimators = at.estimators.read_estimators(modelpath, get_ion_values=False)
+    estimators = at.estimators.read_estimators(modelpath)
 
     modeldata, _ = at.inputmodel.get_modeldata(modelpath)
     ts = at.get_timestep_of_timedays(modelpath, tend)
@@ -655,7 +655,7 @@ def make_emitting_regions_plot(args):
                             "em_Te": dfpackets_selected.em_Te.to_numpy(),
                         }
 
-            estimators = at.estimators.read_estimators(modelpath, get_ion_values=False)
+            estimators = at.estimators.read_estimators(modelpath)
             modeldata, _ = at.inputmodel.get_modeldata(modelpath)
             Tedata_all[modelindex] = {}
             log10nnedata_all[modelindex] = {}
diff --git a/artistools/nonthermal/plotnonthermal.py b/artistools/nonthermal/plotnonthermal.py
index 785d0d1e3..6aa927853 100755
--- a/artistools/nonthermal/plotnonthermal.py
+++ b/artistools/nonthermal/plotnonthermal.py
@@ -69,9 +69,9 @@ def make_xs_plot(axis: plt.Axes, nonthermaldata: pd.DataFrame, args: argparse.Na
 
 
 def plot_contributions(axis, modelpath, timestep, modelgridindex, nonthermaldata, args):
-    estim_tsmgi = at.estimators.read_estimators(
-        modelpath, get_ion_values=True, modelgridindex=modelgridindex, timestep=timestep
-    )[(timestep, modelgridindex)]
+    estim_tsmgi = at.estimators.read_estimators(modelpath, modelgridindex=modelgridindex, timestep=timestep)[
+        (timestep, modelgridindex)
+    ]
 
     total_depev = estim_tsmgi["total_dep"] * ERG_TO_EV
 

From f301db3ad263a367db3a960c729ba489d6cc37ba Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 13:17:01 +0000
Subject: [PATCH 026/150] Restructure estimator file reading to lower mem usage

---
 artistools/estimators/__init__.py   |  1 -
 artistools/estimators/estimators.py | 92 ++++++++++-------------------
 2 files changed, 30 insertions(+), 63 deletions(-)

diff --git a/artistools/estimators/__init__.py b/artistools/estimators/__init__.py
index 0f74b5b7b..812a146c9 100644
--- a/artistools/estimators/__init__.py
+++ b/artistools/estimators/__init__.py
@@ -14,7 +14,6 @@
 from artistools.estimators.estimators import get_units_string
 from artistools.estimators.estimators import get_variablelongunits
 from artistools.estimators.estimators import get_variableunits
-from artistools.estimators.estimators import parse_estimfile
 from artistools.estimators.estimators import read_estimators
 from artistools.estimators.estimators import read_estimators_from_file
 from artistools.estimators.plotestimators import addargs
diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index c3ba6a2f9..fb4bed20f 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -117,11 +117,16 @@ def get_units_string(variable: str) -> str:
     return ""
 
 
-def parse_estimfile(
+def read_estimators_from_file(
     estfilepath: Path | str,
+    printfilename: bool = False,
     skip_emptycells: bool = False,
-) -> t.Iterator[tuple[int, int, dict[str, t.Any]]]:  # pylint: disable=unused-argument
-    """Generate timestep, modelgridindex, dict from estimator file."""
+) -> list[dict[str, t.Any]]:
+    if printfilename:
+        estfilepath = Path(estfilepath)
+        filesize = estfilepath.stat().st_size / 1024 / 1024
+        print(f"Reading {estfilepath.relative_to(estfilepath.parent.parent)} ({filesize:.2f} MiB)")
+    estimblocklist: list[dict[str, t.Any]] = []
     with at.zopen(estfilepath) as estimfile:
         timestep: int | None = None
         modelgridindex: int | None = None
@@ -138,7 +143,9 @@ def parse_estimfile(
                     and modelgridindex is not None
                     and (not skip_emptycells or not estimblock.get("emptycell", True))
                 ):
-                    yield timestep, modelgridindex, estimblock
+                    estimblock["timestep"] = timestep
+                    estimblock["modelgridindex"] = modelgridindex
+                    estimblocklist.append(estimblock)
 
                 timestep = int(row[1])
                 # if timestep > itstep:
@@ -224,36 +231,20 @@ def parse_estimfile(
         and modelgridindex is not None
         and (not skip_emptycells or not estimblock.get("emptycell", True))
     ):
-        yield timestep, modelgridindex, estimblock
+        estimblock["timestep"] = timestep
+        estimblock["modelgridindex"] = modelgridindex
+        estimblocklist.append(estimblock)
+    return estimblocklist
 
 
-def read_estimators_from_file(
-    estfilepath: Path | str,
-    printfilename: bool = False,
-    skip_emptycells: bool = False,
-) -> dict[tuple[int, int], dict[str, t.Any]]:
-    if printfilename:
-        estfilepath = Path(estfilepath)
-        filesize = estfilepath.stat().st_size / 1024 / 1024
-        print(f"Reading {estfilepath.relative_to(estfilepath.parent.parent)} ({filesize:.2f} MiB)")
-
-    return {
-        (timestep, mgi): file_estimblock
-        for timestep, mgi, file_estimblock in parse_estimfile(
-            estfilepath,
-            skip_emptycells=skip_emptycells,
-        )
-    }
-
-
-def get_estimators_in_folder(
+def read_estimators_in_folder_polars(
     modelpath: Path,
     folderpath: Path,
     match_modelgridindex: None | t.Sequence[int],
     skip_emptycells: bool,
-) -> dict:
+) -> pl.DataFrame:
     mpiranklist = at.get_mpiranklist(modelpath, modelgridindex=match_modelgridindex, only_ranks_withgridcells=True)
-    printfilename = len(mpiranklist) < 10
+    printfilename = len(mpiranklist) < 4000
 
     estfilepaths = []
     for mpirank in mpiranklist:
@@ -271,31 +262,19 @@ def get_estimators_in_folder(
         skip_emptycells=skip_emptycells,
     )
 
-    if at.get_config()["num_processes"] > 1:
-        with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
-            arr_rankestimators = pool.map(processfile, estfilepaths)
-            pool.close()
-            pool.join()
-            pool.terminate()
-    else:
-        arr_rankestimators = [processfile(estfilepath) for estfilepath in estfilepaths]
-
-    estimators: dict[tuple[int, int], dict[str, t.Any]] = {}
-    for estfilepath, estimators_thisfile in zip(estfilepaths, arr_rankestimators):
-        dupekeys = sorted([k for k in estimators_thisfile if k in estimators])
-        for k in dupekeys:
-            # dropping the lowest timestep is normal for restarts. Only warn about other cases
-            if k[0] != dupekeys[0][0]:
-                print(
-                    f"WARNING: Duplicate estimator block for (timestep, mgi) key {k}. "
-                    f"Dropping block from {estfilepath}"
-                )
+    pldf = pl.DataFrame()
+    with multiprocessing.get_context("fork").Pool(processes=at.get_config()["num_processes"]) as pool:
+        for estim_onefile in pool.imap(processfile, estfilepaths):
+            pldf_file = pl.DataFrame(estim_onefile).with_columns(
+                pl.col(pl.Int64).cast(pl.Int32), pl.col(pl.Float64).cast(pl.Float32)
+            )
+            pldf = pl.concat([pldf, pldf_file], how="diagonal_relaxed")
 
-            del estimators_thisfile[k]
+        pool.close()
+        pool.join()
+        pool.terminate()
 
-        estimators |= estimators_thisfile
-
-    return estimators
+    return pldf
 
 
 def read_estimators(
@@ -341,24 +320,13 @@ def read_estimators(
         if parquetfile.exists():
             continue
 
-        estim_folder = get_estimators_in_folder(
+        pldf = read_estimators_in_folder_polars(
             modelpath,
             folderpath,
             match_modelgridindex=None,
             skip_emptycells=True,
         )
 
-        pldf = pl.DataFrame(
-            [
-                {
-                    "timestep": ts,
-                    "modelgridindex": mgi,
-                    **estimvals,
-                }
-                for (ts, mgi), estimvals in estim_folder.items()
-                if not estimvals.get("emptycell", True)
-            ]
-        )
         print(f"Writing {parquetfile.relative_to(modelpath.parent)}")
         pldf.write_parquet(parquetfile, compression="zstd")
 

From 10dd7e00434a1728e5dfd118e38c2e9067689c6a Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 13:33:18 +0000
Subject: [PATCH 027/150] Update estimators.py

---
 artistools/estimators/estimators.py | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index fb4bed20f..e47efbfd7 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -121,7 +121,7 @@ def read_estimators_from_file(
     estfilepath: Path | str,
     printfilename: bool = False,
     skip_emptycells: bool = False,
-) -> list[dict[str, t.Any]]:
+) -> pl.DataFrame:
     if printfilename:
         estfilepath = Path(estfilepath)
         filesize = estfilepath.stat().st_size / 1024 / 1024
@@ -234,7 +234,10 @@ def read_estimators_from_file(
         estimblock["timestep"] = timestep
         estimblock["modelgridindex"] = modelgridindex
         estimblocklist.append(estimblock)
-    return estimblocklist
+
+    return pl.DataFrame(estimblocklist).with_columns(
+        pl.col(pl.Int64).cast(pl.Int32), pl.col(pl.Float64).cast(pl.Float32)
+    )
 
 
 def read_estimators_in_folder_polars(
@@ -263,12 +266,8 @@ def read_estimators_in_folder_polars(
     )
 
     pldf = pl.DataFrame()
-    with multiprocessing.get_context("fork").Pool(processes=at.get_config()["num_processes"]) as pool:
-        for estim_onefile in pool.imap(processfile, estfilepaths):
-            pldf_file = pl.DataFrame(estim_onefile).with_columns(
-                pl.col(pl.Int64).cast(pl.Int32), pl.col(pl.Float64).cast(pl.Float32)
-            )
-            pldf = pl.concat([pldf, pldf_file], how="diagonal_relaxed")
+    with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
+        pldf = pl.concat(pool.imap(processfile, estfilepaths), how="diagonal_relaxed")
 
         pool.close()
         pool.join()

From c2e6c8f1e04641efe072fe2b94a89147421825f1 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 13:38:34 +0000
Subject: [PATCH 028/150] Update test_artistools.py

---
 artistools/test_artistools.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/artistools/test_artistools.py b/artistools/test_artistools.py
index a7aeab135..2cb30436a 100755
--- a/artistools/test_artistools.py
+++ b/artistools/test_artistools.py
@@ -10,6 +10,7 @@
 modelpath = at.get_config()["path_testartismodel"]
 modelpath_3d = at.get_config()["path_testartismodel"].parent / "testmodel_3d_10^3"
 outputpath = at.get_config()["path_testoutput"]
+outputpath.mkdir(exist_ok=True, parents=True)
 
 
 def funcname() -> str:
@@ -88,7 +89,7 @@ def test_nltepops() -> None:
 
 def test_radfield() -> None:
     funcoutpath = outputpath / funcname()
-    funcoutpath.mkdir(exist_ok=True)
+    funcoutpath.mkdir(exist_ok=True, parents=True)
     at.radfield.main(argsraw=[], modelpath=modelpath, modelgridindex=0, outputfile=funcoutpath)
 
 
@@ -98,7 +99,7 @@ def test_get_ionrecombratecalibration() -> None:
 
 def test_plotspherical() -> None:
     funcoutpath = outputpath / funcname()
-    funcoutpath.mkdir(exist_ok=True)
+    funcoutpath.mkdir(exist_ok=True, parents=True)
     at.plotspherical.main(argsraw=[], modelpath=modelpath, outputfile=funcoutpath)
 
 

From 688db7643ae2761a328b8375ad37a6ecd17087ad Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 13:44:31 +0000
Subject: [PATCH 029/150] Update estimators.py

---
 artistools/estimators/estimators.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index e47efbfd7..62767d67a 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -267,7 +267,8 @@ def read_estimators_in_folder_polars(
 
     pldf = pl.DataFrame()
     with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
-        pldf = pl.concat(pool.imap(processfile, estfilepaths), how="diagonal_relaxed")
+        for pldf_file in pool.imap(processfile, estfilepaths):
+            pldf = pl.concat([pldf, pldf_file], how="diagonal_relaxed")
 
         pool.close()
         pool.join()

From 77b0ebedffa9cc31805b28940304a9804bc09b02 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 13:44:54 +0000
Subject: [PATCH 030/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 88c7585c5..400abbc76 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -1056,10 +1056,10 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         #  ['_yscale', 'linear']],
         # [['initmasses', ['Ni_56', 'He', 'C', 'Mg']]],
         # ['heating_gamma/gamma_dep'],
-        # ['nne'],
+        ["nne"],
         ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 22000]],
         # ['Te'],
-        # [['averageionisation', ['Fe', 'Ni']]],
+        [["averageionisation", ["Fe", "Ni", "Sr"]]],
         # [['averageexcitation', ['Fe II', 'Fe III']]],
         # [['populations', ['Sr89', 'Sr90', 'Sr91', 'Sr92', 'Sr93', 'Sr94', 'Sr95']],
         #  ['_ymin', 1e-3], ['_ymax', 5]],

From 1b49ede26a36d301587a896f7db6ce978ee77cff Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 14:02:30 +0000
Subject: [PATCH 031/150] Use estimator batches

---
 artistools/estimators/estimators.py | 86 ++++++++++++++++++++---------
 1 file changed, 59 insertions(+), 27 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 62767d67a..30918d764 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -6,6 +6,7 @@
 
 import argparse
 import contextlib
+import itertools
 import math
 import multiprocessing
 import sys
@@ -240,6 +241,22 @@ def read_estimators_from_file(
     )
 
 
+def batched_it(iterable, n):
+    """Batch data into iterators of length n. The last batch may be shorter."""
+    # batched('ABCDEFG', 3) --> ABC DEF G
+    if n < 1:
+        msg = "n must be at least one"
+        raise ValueError(msg)
+    it = iter(iterable)
+    while True:
+        chunk_it = itertools.islice(it, n)
+        try:
+            first_el = next(chunk_it)
+        except StopIteration:
+            return
+        yield list(itertools.chain((first_el,), chunk_it))
+
+
 def read_estimators_in_folder_polars(
     modelpath: Path,
     folderpath: Path,
@@ -247,35 +264,50 @@ def read_estimators_in_folder_polars(
     skip_emptycells: bool,
 ) -> pl.DataFrame:
     mpiranklist = at.get_mpiranklist(modelpath, modelgridindex=match_modelgridindex, only_ranks_withgridcells=True)
-    printfilename = len(mpiranklist) < 4000
-
-    estfilepaths = []
-    for mpirank in mpiranklist:
-        # not worth printing an error, because ranks with no cells to update do not produce an estimator file
-        with contextlib.suppress(FileNotFoundError):
-            estfilepath = at.firstexisting(f"estimators_{mpirank:04d}.out", folder=folderpath, tryzipped=True)
-            estfilepaths.append(estfilepath)
-
-    if not printfilename:
-        print(f"Reading {len(list(estfilepaths))} estimator files in {folderpath.relative_to(Path(folderpath).parent)}")
-
-    processfile = partial(
-        read_estimators_from_file,
-        printfilename=printfilename,
-        skip_emptycells=skip_emptycells,
+    printfilename = True
+    mpirank_groups = list(batched_it(list(mpiranklist), 100))
+    group_parquetfiles = [
+        folderpath / f"estimators_{group[0]:05d}_{group[-1]:05d}.out.parquet.tmp" for group in mpirank_groups
+    ]
+    for group, parquetfilename in zip(mpirank_groups, group_parquetfiles):
+        print(parquetfilename)
+
+        if not parquetfilename.exists():
+            estfilepaths = []
+            for mpirank in group:
+                # not worth printing an error, because ranks with no cells to update do not produce an estimator file
+                with contextlib.suppress(FileNotFoundError):
+                    estfilepath = at.firstexisting(f"estimators_{mpirank:04d}.out", folder=folderpath, tryzipped=True)
+                    estfilepaths.append(estfilepath)
+
+            print(
+                f"Reading {len(list(estfilepaths))} estimator files in {folderpath.relative_to(Path(folderpath).parent)}"
+            )
+
+            processfile = partial(
+                read_estimators_from_file,
+                printfilename=printfilename,
+                skip_emptycells=skip_emptycells,
+            )
+
+            pldf_group = pl.DataFrame()
+            with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
+                for pldf_file in pool.imap(processfile, estfilepaths):
+                    pldf_group = pl.concat([pldf_group, pldf_file], how="diagonal_relaxed")
+
+                pool.close()
+                pool.join()
+                pool.terminate()
+            print(f"Writing {parquetfilename.relative_to(modelpath.parent)}")
+            pldf_group.write_parquet(parquetfilename, compression="zstd")
+
+    for parquetfilename in group_parquetfiles:
+        print(f"Reading {parquetfilename.relative_to(modelpath.parent)}")
+
+    return pl.concat(
+        [pl.read_parquet(parquetfilename) for parquetfilename in group_parquetfiles], how="diagonal_relaxed"
     )
 
-    pldf = pl.DataFrame()
-    with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
-        for pldf_file in pool.imap(processfile, estfilepaths):
-            pldf = pl.concat([pldf, pldf_file], how="diagonal_relaxed")
-
-        pool.close()
-        pool.join()
-        pool.terminate()
-
-    return pldf
-
 
 def read_estimators(
     modelpath: Path | str = Path(),

From 03a0add06e80c751766552f38cdfc347b8cd3f2a Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 14:12:13 +0000
Subject: [PATCH 032/150] Update estimators.py

---
 artistools/estimators/estimators.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 30918d764..1a87f8e73 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -374,9 +374,11 @@ def read_estimators(
         pldflazy = pldflazy.filter(pl.col("timestep").is_in(match_timestep))
 
     for estimtsmgi in pldflazy.collect().iter_rows(named=True):
-        estimators[(estimtsmgi["timestep"], estimtsmgi["modelgridindex"])] = {
-            k: v for k, v in estimtsmgi.items() if k not in {"timestep", "modelgridindex"} and v is not None
-        }
+        ts, mgi = estimtsmgi["timestep"], estimtsmgi["modelgridindex"]
+        if ts is not None and mgi is not None:
+            estimators[(ts, mgi)] = {
+                k: v for k, v in estimtsmgi.items() if k not in {"timestep", "modelgridindex"} and v is not None
+            }
 
     return estimators
 

From df5ad2d8b250fb1ddb9f87df2a6e762b8a4bbfb5 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 14:49:11 +0000
Subject: [PATCH 033/150] Add --markersonly

---
 artistools/estimators/estimators.py     | 15 +++++++--------
 artistools/estimators/plotestimators.py | 15 +++++++++++----
 2 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 1a87f8e73..6ff5a0d6d 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -126,7 +126,7 @@ def read_estimators_from_file(
     if printfilename:
         estfilepath = Path(estfilepath)
         filesize = estfilepath.stat().st_size / 1024 / 1024
-        print(f"Reading {estfilepath.relative_to(estfilepath.parent.parent)} ({filesize:.2f} MiB)")
+        print(f"  Reading {estfilepath.relative_to(estfilepath.parent.parent)} ({filesize:.2f} MiB)")
     estimblocklist: list[dict[str, t.Any]] = []
     with at.zopen(estfilepath) as estimfile:
         timestep: int | None = None
@@ -267,21 +267,20 @@ def read_estimators_in_folder_polars(
     printfilename = True
     mpirank_groups = list(batched_it(list(mpiranklist), 100))
     group_parquetfiles = [
-        folderpath / f"estimators_{group[0]:05d}_{group[-1]:05d}.out.parquet.tmp" for group in mpirank_groups
+        folderpath / f"estimators_{mpigroup[0]:05d}_{mpigroup[-1]:05d}.out.parquet.tmp" for mpigroup in mpirank_groups
     ]
-    for group, parquetfilename in zip(mpirank_groups, group_parquetfiles):
-        print(parquetfilename)
-
+    for mpigroup, parquetfilename in zip(mpirank_groups, group_parquetfiles):
         if not parquetfilename.exists():
+            print(f"{parquetfilename.relative_to(modelpath.parent)} does not exist")
             estfilepaths = []
-            for mpirank in group:
+            for mpirank in mpigroup:
                 # not worth printing an error, because ranks with no cells to update do not produce an estimator file
                 with contextlib.suppress(FileNotFoundError):
                     estfilepath = at.firstexisting(f"estimators_{mpirank:04d}.out", folder=folderpath, tryzipped=True)
                     estfilepaths.append(estfilepath)
 
             print(
-                f"Reading {len(list(estfilepaths))} estimator files in {folderpath.relative_to(Path(folderpath).parent)}"
+                f"Reading {len(list(estfilepaths))} estimator files from {folderpath.relative_to(Path(folderpath).parent)}"
             )
 
             processfile = partial(
@@ -363,7 +362,7 @@ def read_estimators(
         pldf.write_parquet(parquetfile, compression="zstd")
 
     for folderpath, parquetfile in zip(runfolders, parquetfiles):
-        print(f"Reading {parquetfile.relative_to(modelpath.parent)}")
+        print(f"Scanning {parquetfile.relative_to(modelpath.parent)}")
 
     pldflazy = pl.scan_parquet(parquetfiles)
 
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 400abbc76..8af00c988 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -675,6 +675,7 @@ def plot_subplot(
                     modelpath,
                     dfalldata=dfalldata,
                     args=args,
+                    **plotkwargs,
                 )
 
             elif seriestype == "_ymin":
@@ -762,6 +763,10 @@ def make_plot(
     xmin = args.xmin if args.xmin >= 0 else min(xlist)
     xmax = args.xmax if args.xmax > 0 else max(xlist)
 
+    if args.markersonly:
+        plotkwargs["linestyle"] = "None"
+        plotkwargs["marker"] = "."
+
     for ax, plotitems in zip(axes, plotlist):
         ax.set_xlim(left=xmin, right=xmax)
         plot_subplot(
@@ -801,7 +806,7 @@ def make_plot(
         else:
             figure_title = f"{modelname}\nTimestep {timestepslist[0]} ({timeavg:.2f}d)"
 
-        defaultoutputfile = Path("plotestimators_ts{timestep:02d}_{timeavg:.0f}d.pdf")
+        defaultoutputfile = Path("plotestimators_ts{timestep:02d}_{timeavg:.2f}d.pdf")
         if Path(args.outputfile).is_dir():
             args.outputfile = str(Path(args.outputfile, defaultoutputfile))
 
@@ -955,6 +960,8 @@ def addargs(parser: argparse.ArgumentParser) -> None:
 
     parser.add_argument("--hidexlabel", action="store_true", help="Hide the bottom horizontal axis label")
 
+    parser.add_argument("--makersonly", action="store_true", help="Plot markers instead of lines")
+
     parser.add_argument("-filtermovingavg", type=int, default=0, help="Smoothing length (1 is same as none)")
 
     parser.add_argument(
@@ -1056,10 +1063,10 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         #  ['_yscale', 'linear']],
         # [['initmasses', ['Ni_56', 'He', 'C', 'Mg']]],
         # ['heating_gamma/gamma_dep'],
-        ["nne"],
-        ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 22000]],
+        ["nne", ["_ymin", 1e5], ["_ymax", 1e11]],
+        ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 26000]],
         # ['Te'],
-        [["averageionisation", ["Fe", "Ni", "Sr"]]],
+        [["averageionisation", ["Sr"]]],
         # [['averageexcitation', ['Fe II', 'Fe III']]],
         # [['populations', ['Sr89', 'Sr90', 'Sr91', 'Sr92', 'Sr93', 'Sr94', 'Sr95']],
         #  ['_ymin', 1e-3], ['_ymax', 5]],

From 9e429f9ebde20486fa5d3cf226f532d487ea5c19 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 14:53:31 +0000
Subject: [PATCH 034/150] Update estimators.py

---
 artistools/estimators/estimators.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 6ff5a0d6d..29676bac6 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -314,7 +314,6 @@ def read_estimators(
     timestep: None | int | t.Sequence[int] = None,
     runfolder: None | str | Path = None,
     skip_emptycells: bool = False,
-    add_velocity: bool = True,
 ) -> dict[tuple[int, int], dict[str, t.Any]]:
     """Read estimator files into a dictionary of (timestep, modelgridindex): estimators.
 

From 7306a07200295e70d1c4784fede14242ef48f710 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 14:55:43 +0000
Subject: [PATCH 035/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 8af00c988..49d440d23 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -960,7 +960,7 @@ def addargs(parser: argparse.ArgumentParser) -> None:
 
     parser.add_argument("--hidexlabel", action="store_true", help="Hide the bottom horizontal axis label")
 
-    parser.add_argument("--makersonly", action="store_true", help="Plot markers instead of lines")
+    parser.add_argument("--markersonly", action="store_true", help="Plot markers instead of lines")
 
     parser.add_argument("-filtermovingavg", type=int, default=0, help="Smoothing length (1 is same as none)")
 

From bf57ba04cf48db7d517745504a400688c16f40e5 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 15:51:22 +0000
Subject: [PATCH 036/150] Fix duplicate keys in polars estimators

---
 artistools/estimators/__init__.py   |   1 +
 artistools/estimators/estimators.py | 109 +++++++++++++++-------------
 artistools/radfield.py              |  14 ++--
 3 files changed, 70 insertions(+), 54 deletions(-)

diff --git a/artistools/estimators/__init__.py b/artistools/estimators/__init__.py
index 812a146c9..8b88ae47e 100644
--- a/artistools/estimators/__init__.py
+++ b/artistools/estimators/__init__.py
@@ -16,5 +16,6 @@
 from artistools.estimators.estimators import get_variableunits
 from artistools.estimators.estimators import read_estimators
 from artistools.estimators.estimators import read_estimators_from_file
+from artistools.estimators.estimators import read_estimators_polars
 from artistools.estimators.plotestimators import addargs
 from artistools.estimators.plotestimators import main as plot
diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 29676bac6..1205bb45f 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -121,7 +121,7 @@ def get_units_string(variable: str) -> str:
 def read_estimators_from_file(
     estfilepath: Path | str,
     printfilename: bool = False,
-    skip_emptycells: bool = False,
+    skip_emptycells: bool = True,
 ) -> pl.DataFrame:
     if printfilename:
         estfilepath = Path(estfilepath)
@@ -261,13 +261,12 @@ def read_estimators_in_folder_polars(
     modelpath: Path,
     folderpath: Path,
     match_modelgridindex: None | t.Sequence[int],
-    skip_emptycells: bool,
-) -> pl.DataFrame:
+) -> pl.LazyFrame:
     mpiranklist = at.get_mpiranklist(modelpath, modelgridindex=match_modelgridindex, only_ranks_withgridcells=True)
     printfilename = True
     mpirank_groups = list(batched_it(list(mpiranklist), 100))
     group_parquetfiles = [
-        folderpath / f"estimators_{mpigroup[0]:05d}_{mpigroup[-1]:05d}.out.parquet.tmp" for mpigroup in mpirank_groups
+        folderpath / f"estimators_{mpigroup[0]:04d}_{mpigroup[-1]:04d}.out.parquet.tmp" for mpigroup in mpirank_groups
     ]
     for mpigroup, parquetfilename in zip(mpirank_groups, group_parquetfiles):
         if not parquetfilename.exists():
@@ -286,17 +285,20 @@ def read_estimators_in_folder_polars(
             processfile = partial(
                 read_estimators_from_file,
                 printfilename=printfilename,
-                skip_emptycells=skip_emptycells,
             )
 
-            pldf_group = pl.DataFrame()
+            pldf_group = None
             with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
                 for pldf_file in pool.imap(processfile, estfilepaths):
-                    pldf_group = pl.concat([pldf_group, pldf_file], how="diagonal_relaxed")
+                    if pldf_group is None:
+                        pldf_group = pldf_file
+                    else:
+                        pldf_group = pl.concat([pldf_group, pldf_file], how="diagonal_relaxed")
 
                 pool.close()
                 pool.join()
                 pool.terminate()
+            assert pldf_group is not None
             print(f"Writing {parquetfilename.relative_to(modelpath.parent)}")
             pldf_group.write_parquet(parquetfilename, compression="zstd")
 
@@ -304,17 +306,16 @@ def read_estimators_in_folder_polars(
         print(f"Reading {parquetfilename.relative_to(modelpath.parent)}")
 
     return pl.concat(
-        [pl.read_parquet(parquetfilename) for parquetfilename in group_parquetfiles], how="diagonal_relaxed"
-    )
+        [pl.scan_parquet(parquetfilename) for parquetfilename in group_parquetfiles], how="diagonal_relaxed"
+    ).sort(["timestep", "modelgridindex"])
 
 
-def read_estimators(
+def read_estimators_polars(
     modelpath: Path | str = Path(),
     modelgridindex: None | int | t.Sequence[int] = None,
     timestep: None | int | t.Sequence[int] = None,
     runfolder: None | str | Path = None,
-    skip_emptycells: bool = False,
-) -> dict[tuple[int, int], dict[str, t.Any]]:
+) -> pl.LazyFrame:
     """Read estimator files into a dictionary of (timestep, modelgridindex): estimators.
 
     Selecting particular timesteps or modelgrid cells will using speed this up by reducing the number of files that must be read.
@@ -337,33 +338,44 @@ def read_estimators(
         match_timestep = tuple(timestep)
 
     if not Path(modelpath).exists() and Path(modelpath).parts[0] == "codecomparison":
-        return at.codecomparison.read_reference_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex)
+        estimators = at.codecomparison.read_reference_estimators(
+            modelpath, timestep=timestep, modelgridindex=modelgridindex
+        )
+        return pl.DataFrame(
+            [
+                {
+                    "timestep": ts,
+                    "modelgridindex": mgi,
+                    **estimvals,
+                }
+                for (ts, mgi), estimvals in estimators.items()
+                if not estimvals.get("emptycell", True)
+            ]
+        ).lazy()
 
     # print(f" matching cells {match_modelgridindex} and timesteps {match_timestep}")
 
     runfolders = at.get_runfolders(modelpath, timesteps=match_timestep) if runfolder is None else [Path(runfolder)]
 
-    estimators: dict[tuple[int, int], dict[str, t.Any]] = {}
     parquetfiles = [folderpath / "estimators.out.parquet.tmp" for folderpath in runfolders]
 
     for folderpath, parquetfile in zip(runfolders, parquetfiles):
-        if parquetfile.exists():
-            continue
-
-        pldf = read_estimators_in_folder_polars(
-            modelpath,
-            folderpath,
-            match_modelgridindex=None,
-            skip_emptycells=True,
-        )
+        if not parquetfile.exists():
+            pldflazy = read_estimators_in_folder_polars(
+                modelpath,
+                folderpath,
+                match_modelgridindex=None,
+            )
 
-        print(f"Writing {parquetfile.relative_to(modelpath.parent)}")
-        pldf.write_parquet(parquetfile, compression="zstd")
+            print(f"Writing {parquetfile.relative_to(modelpath.parent)}")
+            pldflazy.collect().write_parquet(parquetfile, compression="zstd")
 
     for folderpath, parquetfile in zip(runfolders, parquetfiles):
         print(f"Scanning {parquetfile.relative_to(modelpath.parent)}")
 
-    pldflazy = pl.scan_parquet(parquetfiles)
+    pldflazy = pl.concat(
+        [pl.scan_parquet(parquetfilename) for parquetfilename in parquetfiles], how="diagonal_relaxed"
+    ).unique(["timestep", "modelgridindex"], maintain_order=True, keep="first")
 
     if match_modelgridindex is not None:
         pldflazy = pldflazy.filter(pl.col("modelgridindex").is_in(match_modelgridindex))
@@ -371,12 +383,25 @@ def read_estimators(
     if match_timestep is not None:
         pldflazy = pldflazy.filter(pl.col("timestep").is_in(match_timestep))
 
+    return pldflazy
+
+
+def read_estimators(
+    modelpath: Path | str = Path(),
+    modelgridindex: None | int | t.Sequence[int] = None,
+    timestep: None | int | t.Sequence[int] = None,
+    runfolder: None | str | Path = None,
+    keys: t.Collection[str] | None = None,
+) -> dict[tuple[int, int], dict[str, t.Any]]:
+    if keys is None:
+        keys = {}
+    pldflazy = read_estimators_polars(modelpath, modelgridindex, timestep, runfolder)
+    estimators: dict[tuple[int, int], dict[str, t.Any]] = {}
     for estimtsmgi in pldflazy.collect().iter_rows(named=True):
         ts, mgi = estimtsmgi["timestep"], estimtsmgi["modelgridindex"]
-        if ts is not None and mgi is not None:
-            estimators[(ts, mgi)] = {
-                k: v for k, v in estimtsmgi.items() if k not in {"timestep", "modelgridindex"} and v is not None
-            }
+        estimators[(ts, mgi)] = {
+            k: v for k, v in estimtsmgi.items() if k not in {"timestep", "modelgridindex", *keys} and v is not None
+        }
 
     return estimators
 
@@ -485,22 +510,8 @@ def get_partiallycompletetimesteps(estimators: dict[tuple[int, int], dict[str, t
 
 def get_temperatures(modelpath: str | Path) -> pl.LazyFrame:
     """Get a polars DataFrame containing the temperatures at every timestep and modelgridindex."""
-    dfest_parquetfile = Path(modelpath, "temperatures.parquet.tmp")
-
-    if not dfest_parquetfile.is_file():
-        estimators = at.estimators.read_estimators(
-            modelpath,
-            skip_emptycells=True,
-        )
-        assert len(estimators) > 0
-        pl.DataFrame(
-            {
-                "timestep": (ts for ts, _ in estimators),
-                "modelgridindex": (mgi for _, mgi in estimators),
-                "TR": (estimators[tsmgi].get("TR", -1) for tsmgi in estimators),
-            },
-        ).filter(pl.col("TR") >= 0).with_columns(pl.col(pl.Int64).cast(pl.Int32)).write_parquet(
-            dfest_parquetfile, compression="zstd"
-        )
-
-    return pl.scan_parquet(dfest_parquetfile)
+    return (
+        at.estimators.read_estimators_polars(modelpath=modelpath)
+        .select(["timestep", "modelgridindex", "TR"])
+        .drop_nulls()
+    )
diff --git a/artistools/radfield.py b/artistools/radfield.py
index a15b0ffe7..3d9a167b1 100755
--- a/artistools/radfield.py
+++ b/artistools/radfield.py
@@ -311,7 +311,9 @@ def sigma_bf(nu):
 
 def get_kappa_bf_ion(atomic_number, lower_ionstage, modelgridindex, timestep, modelpath, arr_nu_hz, max_levels):
     adata = at.atomic.get_levels(modelpath, get_photoionisations=True)
-    estimators = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex)
+    estimators = at.estimators.read_estimators(
+        modelpath, timestep=timestep, modelgridindex=modelgridindex, use_polars=False
+    )
     T_e = estimators[(timestep, modelgridindex)]["Te"]
 
     ion_data = adata.query("Z == @atomic_number and ionstage == @lower_ionstage").iloc[0]
@@ -353,9 +355,9 @@ def get_recombination_emission(
     upper_ion_data = adata.query("Z == @atomic_number and ionstage == @upper_ionstage").iloc[0]
     lower_ion_data = adata.query("Z == @atomic_number and ionstage == @lower_ionstage").iloc[0]
 
-    estimtsmgi = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex)[
-        (timestep, modelgridindex)
-    ]
+    estimtsmgi = at.estimators.read_estimators(
+        modelpath, timestep=timestep, modelgridindex=modelgridindex, use_polars=False
+    )[(timestep, modelgridindex)]
 
     upperionstr = at.get_ionstring(atomic_number, upper_ionstage, sep="_", style="spectral")
     upperionpopdensity = estimtsmgi[f"nnion_{upperionstr}"]
@@ -459,7 +461,9 @@ def get_recombination_emission(
 
 def get_ion_gamma_dnu(modelpath, modelgridindex, timestep, atomic_number, ionstage, arr_nu_hz, J_nu_arr, max_levels):
     """Calculate the contribution to the photoionisation rate coefficient per J_nu at each frequency nu for an ion."""
-    estimators = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex)
+    estimators = at.estimators.read_estimators(
+        modelpath, timestep=timestep, modelgridindex=modelgridindex, use_polars=False
+    )
 
     T_e = estimators[(timestep, modelgridindex)]["Te"]
     T_R = estimators[(timestep, modelgridindex)]["TR"]

From a037c0465de929c179816181f837bbd3ddddc841 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 15:56:51 +0000
Subject: [PATCH 037/150] Remove get_temperatures()

---
 artistools/estimators/__init__.py       | 1 -
 artistools/estimators/estimators.py     | 8 +++++---
 artistools/estimators/plotestimators.py | 5 +++--
 artistools/plotspherical.py             | 7 +++++--
 4 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/artistools/estimators/__init__.py b/artistools/estimators/__init__.py
index 8b88ae47e..724a6c6d2 100644
--- a/artistools/estimators/__init__.py
+++ b/artistools/estimators/__init__.py
@@ -10,7 +10,6 @@
 from artistools.estimators.estimators import get_dictlabelreplacements
 from artistools.estimators.estimators import get_ionrecombrates_fromfile
 from artistools.estimators.estimators import get_partiallycompletetimesteps
-from artistools.estimators.estimators import get_temperatures
 from artistools.estimators.estimators import get_units_string
 from artistools.estimators.estimators import get_variablelongunits
 from artistools.estimators.estimators import get_variableunits
diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 1205bb45f..e5a054b84 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -393,14 +393,16 @@ def read_estimators(
     runfolder: None | str | Path = None,
     keys: t.Collection[str] | None = None,
 ) -> dict[tuple[int, int], dict[str, t.Any]]:
-    if keys is None:
-        keys = {}
+    if isinstance(keys, str):
+        keys = {keys}
     pldflazy = read_estimators_polars(modelpath, modelgridindex, timestep, runfolder)
     estimators: dict[tuple[int, int], dict[str, t.Any]] = {}
     for estimtsmgi in pldflazy.collect().iter_rows(named=True):
         ts, mgi = estimtsmgi["timestep"], estimtsmgi["modelgridindex"]
         estimators[(ts, mgi)] = {
-            k: v for k, v in estimtsmgi.items() if k not in {"timestep", "modelgridindex", *keys} and v is not None
+            k: v
+            for k, v in estimtsmgi.items()
+            if k not in {"timestep", "modelgridindex"} and (keys is None or k in keys) and v is not None
         }
 
     return estimators
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 49d440d23..a121157cc 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -1110,8 +1110,9 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         modeldata, _ = at.inputmodel.get_modeldata(modelpath)
         estimators = artistools.estimators.estimators_classic.read_classic_estimators(modelpath, modeldata)
     elif temperatures_only:
-        df_estimators = at.estimators.get_temperatures(modelpath).filter(pl.col("timestep").is_in(timesteps_included))
-        estimators = df_estimators.collect().rows_by_key(key=["timestep", "modelgridindex"], named=True, unique=True)
+        estimators = at.estimators.read_estimators(
+            modelpath=modelpath, modelgridindex=args.modelgridindex, timestep=tuple(timesteps_included), keys=["TR"]
+        )
     else:
         estimators = at.estimators.read_estimators(
             modelpath=modelpath, modelgridindex=args.modelgridindex, timestep=tuple(timesteps_included)
diff --git a/artistools/plotspherical.py b/artistools/plotspherical.py
index 20200dfb9..10f46adc7 100755
--- a/artistools/plotspherical.py
+++ b/artistools/plotspherical.py
@@ -119,8 +119,11 @@ def plot_spherical(
             ).alias("em_timestep")
         )
 
-        df_estimators = at.estimators.get_temperatures(modelpath).rename(
-            {"timestep": "em_timestep", "modelgridindex": "em_modelgridindex", "TR": "em_TR"}
+        df_estimators = (
+            at.estimators.read_estimators_polars(modelpath=modelpath)
+            .select(["timestep", "modelgridindex", "TR"])
+            .drop_nulls()
+            .rename({"timestep": "em_timestep", "modelgridindex": "em_modelgridindex", "TR": "em_TR"})
         )
         dfpackets = dfpackets.join(df_estimators, on=["em_timestep", "em_modelgridindex"], how="left")
         aggs.append(((pl.col("em_TR") * pl.col("e_rf")).mean() / pl.col("e_rf").mean()).alias("temperature"))

From b85893bf67910e8fdde8cc57987072c27de77dac Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 15:57:33 +0000
Subject: [PATCH 038/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index a121157cc..4463906b1 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -1109,13 +1109,12 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
 
         modeldata, _ = at.inputmodel.get_modeldata(modelpath)
         estimators = artistools.estimators.estimators_classic.read_classic_estimators(modelpath, modeldata)
-    elif temperatures_only:
-        estimators = at.estimators.read_estimators(
-            modelpath=modelpath, modelgridindex=args.modelgridindex, timestep=tuple(timesteps_included), keys=["TR"]
-        )
     else:
         estimators = at.estimators.read_estimators(
-            modelpath=modelpath, modelgridindex=args.modelgridindex, timestep=tuple(timesteps_included)
+            modelpath=modelpath,
+            modelgridindex=args.modelgridindex,
+            timestep=tuple(timesteps_included),
+            keys=["TR"] if temperatures_only else None,
         )
     assert estimators is not None
 

From ef8d90fa16070480007ecfacd97b9a2630286ba9 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 16:31:05 +0000
Subject: [PATCH 039/150] Use polars in plotestimators

---
 artistools/estimators/estimators.py     | 55 +++++++++--------
 artistools/estimators/plotestimators.py | 79 +++++++++++++++----------
 artistools/misc.py                      |  4 ++
 3 files changed, 83 insertions(+), 55 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index e5a054b84..8bbf6648f 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -18,6 +18,7 @@
 import numpy as np
 import pandas as pd
 import polars as pl
+import polars.selectors as cs
 
 import artistools as at
 
@@ -127,6 +128,7 @@ def read_estimators_from_file(
         estfilepath = Path(estfilepath)
         filesize = estfilepath.stat().st_size / 1024 / 1024
         print(f"  Reading {estfilepath.relative_to(estfilepath.parent.parent)} ({filesize:.2f} MiB)")
+
     estimblocklist: list[dict[str, t.Any]] = []
     with at.zopen(estfilepath) as estimfile:
         timestep: int | None = None
@@ -263,7 +265,7 @@ def read_estimators_in_folder_polars(
     match_modelgridindex: None | t.Sequence[int],
 ) -> pl.LazyFrame:
     mpiranklist = at.get_mpiranklist(modelpath, modelgridindex=match_modelgridindex, only_ranks_withgridcells=True)
-    printfilename = True
+
     mpirank_groups = list(batched_it(list(mpiranklist), 100))
     group_parquetfiles = [
         folderpath / f"estimators_{mpigroup[0]:04d}_{mpigroup[-1]:04d}.out.parquet.tmp" for mpigroup in mpirank_groups
@@ -282,10 +284,7 @@ def read_estimators_in_folder_polars(
                 f"Reading {len(list(estfilepaths))} estimator files from {folderpath.relative_to(Path(folderpath).parent)}"
             )
 
-            processfile = partial(
-                read_estimators_from_file,
-                printfilename=printfilename,
-            )
+            processfile = partial(read_estimators_from_file)
 
             pldf_group = None
             with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
@@ -410,7 +409,7 @@ def read_estimators(
 
 def get_averaged_estimators(
     modelpath: Path | str,
-    estimators: dict[tuple[int, int], dict],
+    estimators: pl.LazyFrame | pl.DataFrame,
     timesteps: int | t.Sequence[int],
     modelgridindex: int,
     keys: str | list | None,
@@ -424,7 +423,7 @@ def get_averaged_estimators(
     if isinstance(keys, str):
         keys = [keys]
     elif keys is None or not keys:
-        keys = list(estimators[(timesteps[0], modelgridindex)].keys())
+        keys = [c for c in estimators.columns if c not in {"timestep", "modelgridindex"}]
 
     dictout = {}
     tdeltas = at.get_timestep_times(modelpath, loc="delta")
@@ -433,29 +432,44 @@ def get_averaged_estimators(
         tdeltasum = 0
         for timestep, tdelta in zip(timesteps, tdeltas):
             for mgi in range(modelgridindex - avgadjcells, modelgridindex + avgadjcells + 1):
-                valuesum += estimators[(timestep, mgi)][k] * tdelta
+                value = (
+                    estimators.filter(pl.col("timestep") == timestep)
+                    .filter(pl.col("modelgridindex") == mgi)
+                    .select(k)
+                    .lazy()
+                    .collect()
+                    .item(0, 0)
+                )
+                if value is None:
+                    msg = f"{k} not found for timestep {timestep} and modelgridindex {mgi}"
+                    raise ValueError(msg)
+
+                valuesum += value * tdelta
                 tdeltasum += tdelta
         dictout[k] = valuesum / tdeltasum
 
     return dictout
 
 
-def get_averageionisation(estimatorstsmgi: dict[str, float], atomic_number: int) -> float:
+def get_averageionisation(estimatorstsmgi: pl.LazyFrame, atomic_number: int) -> float:
     free_electron_weighted_pop_sum = 0.0
     found = False
     popsum = 0.0
     elsymb = at.get_elsymbol(atomic_number)
-    for key in estimatorstsmgi:
-        if key.startswith(f"nnion_{elsymb}_"):
-            found = True
-            ionstage = at.decode_roman_numeral(key.removeprefix(f"nnion_{elsymb}_"))
-            free_electron_weighted_pop_sum += estimatorstsmgi[key] * (ionstage - 1)
-            popsum += estimatorstsmgi[key]
+    dfselected = estimatorstsmgi.select(
+        cs.starts_with(f"nnion_{elsymb}_") | cs.by_name(f"nnelement_{elsymb}")
+    ).collect()
+    for key in dfselected.columns:
+        found = True
+        nnion = dfselected[key].item(0)
+        ionstage = at.decode_roman_numeral(key.removeprefix(f"nnion_{elsymb}_"))
+        free_electron_weighted_pop_sum += nnion * (ionstage - 1)
+        popsum += nnion
 
     if not found:
         return float("NaN")
 
-    return free_electron_weighted_pop_sum / estimatorstsmgi[f"nnelement_{elsymb}"]
+    return free_electron_weighted_pop_sum / dfselected[f"nnelement_{elsymb}"].item(0)
 
 
 def get_averageexcitation(
@@ -508,12 +522,3 @@ def get_partiallycompletetimesteps(estimators: dict[tuple[int, int], dict[str, t
         all_mgis.add(mgi)
 
     return [nts for nts, mgilist in timestepcells.items() if len(mgilist) < len(all_mgis)]
-
-
-def get_temperatures(modelpath: str | Path) -> pl.LazyFrame:
-    """Get a polars DataFrame containing the temperatures at every timestep and modelgridindex."""
-    return (
-        at.estimators.read_estimators_polars(modelpath=modelpath)
-        .select(["timestep", "modelgridindex", "TR"])
-        .drop_nulls()
-    )
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 4463906b1..0124965aa 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -149,11 +149,23 @@ def plot_average_ionisation_excitation(
             for timestep in timesteps:
                 if seriestype == "averageionisation":
                     valuesum += (
-                        at.estimators.get_averageionisation(estimators[(timestep, modelgridindex)], atomic_number)
+                        at.estimators.get_averageionisation(
+                            estimators.filter(pl.col("timestep") == timestep).filter(
+                                pl.col("modelgridindex") == modelgridindex
+                            ),
+                            atomic_number,
+                        )
                         * arr_tdelta[timestep]
                     )
                 elif seriestype == "averageexcitation":
-                    T_exc = estimators[(timestep, modelgridindex)]["Te"]
+                    T_exc = (
+                        estimators.filter(pl.col("timestep") == timestep)
+                        .filter(pl.col("modelgridindex") == modelgridindex)
+                        .select("Te")
+                        .lazy()
+                        .collect()
+                        .item(0, 0)
+                    )
                     valuesum += (
                         at.estimators.get_averageexcitation(
                             modelpath, modelgridindex, timestep, atomic_number, ionstage, T_exc
@@ -337,6 +349,7 @@ def get_iontuple(ionstr):
             colorindex += 1
 
         elsymbol = at.get_elsymbol(atomic_number)
+
         ionstr = at.get_ionstring(atomic_number, ionstage, sep="_", style="spectral")
 
         if seriestype == "populations":
@@ -541,7 +554,7 @@ def plot_series(
 def get_xlist(
     xvariable: str,
     allnonemptymgilist: t.Sequence[int],
-    estimators: dict,
+    estimators: pl.LazyFrame | pl.DataFrame,
     timestepslist: t.Any,
     modelpath: str | Path,
     args: t.Any,
@@ -722,7 +735,7 @@ def make_plot(
     modelpath: Path | str,
     timestepslist_unfiltered: list[list[int]],
     allnonemptymgilist: list[int],
-    estimators: dict,
+    estimators: pl.LazyFrame | pl.DataFrame,
     xvariable: str,
     plotlist,
     args: t.Any,
@@ -797,7 +810,14 @@ def make_plot(
         timeavg = (args.timemin + args.timemax) / 2.0
         if args.multiplot and not args.classicartis:
             assert isinstance(timestepslist[0], list)
-            tdays = estimators[(timestepslist[0][0], mgilist[0])]["tdays"]
+            tdays = (
+                estimators.filter(pl.col("timestep") == timestepslist[0][0])
+                .filter(pl.col("modelgridindex") == mgilist[0])
+                .select("tdays")
+                .lazy()
+                .collect()
+                .item(0)
+            )
             figure_title = f"{modelname}\nTimestep {timestepslist[0]} ({tdays:.2f}d)"
         elif args.multiplot:
             assert isinstance(timestepslist[0], int)
@@ -914,19 +934,6 @@ def plot_recombrates(modelpath, estimators, atomic_number, ionstage_list, **plot
     plt.close()
 
 
-def plotlist_is_temperatures_only(plotlist: str | t.Sequence) -> bool:
-    if isinstance(plotlist, str):
-        if plotlist != "TR" and not plotlist.startswith("_"):
-            return False
-
-    elif not isinstance(plotlist[0], str) or not plotlist[0].startswith("_"):
-        for item in plotlist:
-            if not plotlist_is_temperatures_only(item):
-                return False
-
-    return True
-
-
 def addargs(parser: argparse.ArgumentParser) -> None:
     parser.add_argument(
         "-modelpath", default=".", help="Paths to ARTIS folder (or virtual path e.g. codecomparison/ddc10/cmfgen)"
@@ -1064,13 +1071,13 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         # [['initmasses', ['Ni_56', 'He', 'C', 'Mg']]],
         # ['heating_gamma/gamma_dep'],
         ["nne", ["_ymin", 1e5], ["_ymax", 1e11]],
-        ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 26000]],
-        # ['Te'],
+        # ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 26000]],
+        ["Te"],
         [["averageionisation", ["Sr"]]],
         # [['averageexcitation', ['Fe II', 'Fe III']]],
-        # [['populations', ['Sr89', 'Sr90', 'Sr91', 'Sr92', 'Sr93', 'Sr94', 'Sr95']],
+        [["populations", ["Sr89", "Sr90", "Sr91", "Sr92", "Sr93", "Sr94"]]],
         #  ['_ymin', 1e-3], ['_ymax', 5]],
-        # [["populations", ["Fe", "Co", "Ni", "Sr", "Nd", "U"]]],
+        [["populations", ["Fe", "Co", "Ni", "Sr", "Nd"]]],
         # [['populations', ['He I', 'He II', 'He III']]],
         # [['populations', ['C I', 'C II', 'C III', 'C IV', 'C V']]],
         # [['populations', ['O I', 'O II', 'O III', 'O IV']]],
@@ -1100,26 +1107,34 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         dfselectedcells = dfselectedcells.query("rho > 0")
         args.modelgridindex = dfselectedcells["inputcellid"]
 
-    if temperatures_only := plotlist_is_temperatures_only(plotlist):
-        print("Plotting temperatures only (from parquet if available)")
-
     timesteps_included = list(range(timestepmin, timestepmax + 1))
     if args.classicartis:
         import artistools.estimators.estimators_classic
 
         modeldata, _ = at.inputmodel.get_modeldata(modelpath)
-        estimators = artistools.estimators.estimators_classic.read_classic_estimators(modelpath, modeldata)
+        estimatorsdict = artistools.estimators.estimators_classic.read_classic_estimators(modelpath, modeldata)
+        assert estimatorsdict is not None
+        estimators = pl.DataFrame(
+            [
+                {
+                    "timestep": ts,
+                    "modelgridindex": mgi,
+                    **estimvals,
+                }
+                for (ts, mgi), estimvals in estimatorsdict.items()
+                if not estimvals.get("emptycell", True)
+            ]
+        ).lazy()
     else:
-        estimators = at.estimators.read_estimators(
+        estimators = at.estimators.read_estimators_polars(
             modelpath=modelpath,
             modelgridindex=args.modelgridindex,
             timestep=tuple(timesteps_included),
-            keys=["TR"] if temperatures_only else None,
         )
     assert estimators is not None
 
     for ts in reversed(timesteps_included):
-        tswithdata = [ts for (ts, mgi) in estimators]
+        tswithdata = estimators.select("timestep").unique().collect().to_numpy()
         for ts in timesteps_included:
             if ts not in tswithdata:
                 timesteps_included.remove(ts)
@@ -1159,7 +1174,11 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             allnonemptymgilist = [
                 modelgridindex
                 for modelgridindex in modeldata.index
-                if (timesteps_included[0], modelgridindex) in estimators
+                if not estimators.filter(pl.col("modelgridindex") == modelgridindex)
+                .select("modelgridindex")
+                .lazy()
+                .collect()
+                .is_empty()
             ]
         else:
             allnonemptymgilist = [mgi for mgi, assocpropcells in assoc_cells.items() if assocpropcells]
diff --git a/artistools/misc.py b/artistools/misc.py
index 9514ac7bb..15e06a17d 100644
--- a/artistools/misc.py
+++ b/artistools/misc.py
@@ -677,6 +677,10 @@ def get_ionstring(
     if ionstage is None or ionstage == "ALL":
         return f"{get_elsymbol(atomic_number)}"
 
+    if isinstance(ionstage, str) and ionstage.startswith(at.get_elsymbol(atomic_number)):
+        # nuclides like Sr89 get passed in as atomic_number=38, ionstage='Sr89'
+        return ionstage
+
     assert not isinstance(ionstage, str)
 
     if style == "spectral":

From a2369ef7962afcfec11416c9192ce39708cfb22a Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 16:43:01 +0000
Subject: [PATCH 040/150] Speed up estimator plotting

---
 artistools/estimators/estimators.py      | 18 +++++++++++-------
 artistools/estimators/plotestimators.py  |  4 ++--
 artistools/inputmodel/inputmodel_misc.py |  2 +-
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 8bbf6648f..cbdda7099 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -427,18 +427,21 @@ def get_averaged_estimators(
 
     dictout = {}
     tdeltas = at.get_timestep_times(modelpath, loc="delta")
+    mgilist = list(range(modelgridindex - avgadjcells, modelgridindex + avgadjcells + 1))
+    estcollect = (
+        estimators.lazy()
+        .filter(pl.col("timestep").is_in(timesteps))
+        .filter(pl.col("modelgridindex").is_in(mgilist))
+        .select({*keys, "timestep", "modelgridindex"})
+        .collect()
+    )
     for k in keys:
         valuesum = 0
         tdeltasum = 0
         for timestep, tdelta in zip(timesteps, tdeltas):
-            for mgi in range(modelgridindex - avgadjcells, modelgridindex + avgadjcells + 1):
+            for mgi in mgilist:
                 value = (
-                    estimators.filter(pl.col("timestep") == timestep)
-                    .filter(pl.col("modelgridindex") == mgi)
-                    .select(k)
-                    .lazy()
-                    .collect()
-                    .item(0, 0)
+                    estcollect.filter(pl.col("timestep") == timestep).filter(pl.col("modelgridindex") == mgi)[k].item(0)
                 )
                 if value is None:
                     msg = f"{k} not found for timestep {timestep} and modelgridindex {mgi}"
@@ -446,6 +449,7 @@ def get_averaged_estimators(
 
                 valuesum += value * tdelta
                 tdeltasum += tdelta
+
         dictout[k] = valuesum / tdeltasum
 
     return dictout
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 0124965aa..8b554ca73 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -589,7 +589,7 @@ def get_xlist(
             dfmodel = dfmodel.filter(pl.col("vel_r_mid") / 1e5 <= args.xmax)
         else:
             dfmodel = dfmodel.filter(pl.col("vel_r_mid") <= modelmeta["vmax_cmps"])
-        dfmodelcollect = dfmodel.collect()
+        dfmodelcollect = dfmodel.select(["vel_r_mid", "modelgridindex"]).collect()
 
         scalefactor = 1e5 if xvariable == "velocity" else 29979245800
         xlist = (dfmodelcollect["vel_r_mid"] / scalefactor).to_list()
@@ -1075,7 +1075,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         ["Te"],
         [["averageionisation", ["Sr"]]],
         # [['averageexcitation', ['Fe II', 'Fe III']]],
-        [["populations", ["Sr89", "Sr90", "Sr91", "Sr92", "Sr93", "Sr94"]]],
+        # [["populations", ["Sr90", "Sr91", "Sr92", "Sr93", "Sr94"]]],
         #  ['_ymin', 1e-3], ['_ymax', 5]],
         [["populations", ["Fe", "Co", "Ni", "Sr", "Nd"]]],
         # [['populations', ['He I', 'He II', 'He III']]],
diff --git a/artistools/inputmodel/inputmodel_misc.py b/artistools/inputmodel/inputmodel_misc.py
index 4a83143ee..24ceda0b4 100644
--- a/artistools/inputmodel/inputmodel_misc.py
+++ b/artistools/inputmodel/inputmodel_misc.py
@@ -364,7 +364,7 @@ def get_modeldata_polars(
         raise AssertionError
 
     mebibyte = 1024 * 1024
-    if isinstance(dfmodel, pl.DataFrame) and filename.stat().st_size > 10 * mebibyte and not getheadersonly:
+    if isinstance(dfmodel, pl.DataFrame) and filename.stat().st_size > 5 * mebibyte and not getheadersonly:
         print(f"Saving {filenameparquet}")
         dfmodel.write_parquet(filenameparquet, compression="zstd")
         print("  Done.")

From c83330e8d449208d0d8f88d82a2eefa8dafabadb Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 16:52:57 +0000
Subject: [PATCH 041/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 8b554ca73..41dcffd6a 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -780,6 +780,23 @@ def make_plot(
         plotkwargs["linestyle"] = "None"
         plotkwargs["marker"] = "."
 
+    # ideally, we could start filtering the columns here, but it's still faster to collect the whole thing
+    allts: set[int] = set()
+    for tspoint in timestepslist:
+        if isinstance(tspoint, int):
+            allts.add(tspoint)
+        else:
+            for ts in tspoint:
+                allts.add(ts)
+
+    estimators = (
+        estimators.filter(pl.col("modelgridindex").is_in(mgilist))
+        .filter(pl.col("timestep").is_in(allts))
+        .lazy()
+        .collect()
+        .lazy()
+    )
+
     for ax, plotitems in zip(axes, plotlist):
         ax.set_xlim(left=xmin, right=xmax)
         plot_subplot(

From f09ea6841976a295771b94c43b896f7357b0ecae Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 16:57:02 +0000
Subject: [PATCH 042/150] Update estimators.py

---
 artistools/estimators/estimators.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index cbdda7099..9853adc48 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -466,6 +466,9 @@ def get_averageionisation(estimatorstsmgi: pl.LazyFrame, atomic_number: int) ->
     for key in dfselected.columns:
         found = True
         nnion = dfselected[key].item(0)
+        if nnion is None:
+            nnion = f"WARNING: {key} is None"
+            nnion = 0
         ionstage = at.decode_roman_numeral(key.removeprefix(f"nnion_{elsymb}_"))
         free_electron_weighted_pop_sum += nnion * (ionstage - 1)
         popsum += nnion

From 26d9bc871a594f25e6234ee2255d3268a8987010 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 16:57:14 +0000
Subject: [PATCH 043/150] Update estimators.py

---
 artistools/estimators/estimators.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 9853adc48..a38324f94 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -468,7 +468,7 @@ def get_averageionisation(estimatorstsmgi: pl.LazyFrame, atomic_number: int) ->
         nnion = dfselected[key].item(0)
         if nnion is None:
             nnion = f"WARNING: {key} is None"
-            nnion = 0
+            nnion = 0.0
         ionstage = at.decode_roman_numeral(key.removeprefix(f"nnion_{elsymb}_"))
         free_electron_weighted_pop_sum += nnion * (ionstage - 1)
         popsum += nnion

From 662ac2d4d538adc2aaa67c53af2a60118b038cde Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 18:54:27 +0000
Subject: [PATCH 044/150] Update estimators.py

---
 artistools/estimators/estimators.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index a38324f94..3c4ef1922 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -457,26 +457,30 @@ def get_averaged_estimators(
 
 def get_averageionisation(estimatorstsmgi: pl.LazyFrame, atomic_number: int) -> float:
     free_electron_weighted_pop_sum = 0.0
-    found = False
-    popsum = 0.0
     elsymb = at.get_elsymbol(atomic_number)
+
     dfselected = estimatorstsmgi.select(
         cs.starts_with(f"nnion_{elsymb}_") | cs.by_name(f"nnelement_{elsymb}")
     ).collect()
+
+    nnelement = dfselected[f"nnelement_{elsymb}"].item(0)
+    if nnelement is None:
+        return float("NaN")
+
+    found = False
+    popsum = 0.0
     for key in dfselected.columns:
         found = True
         nnion = dfselected[key].item(0)
         if nnion is None:
             nnion = f"WARNING: {key} is None"
             nnion = 0.0
+
         ionstage = at.decode_roman_numeral(key.removeprefix(f"nnion_{elsymb}_"))
         free_electron_weighted_pop_sum += nnion * (ionstage - 1)
         popsum += nnion
 
-    if not found:
-        return float("NaN")
-
-    return free_electron_weighted_pop_sum / dfselected[f"nnelement_{elsymb}"].item(0)
+    return free_electron_weighted_pop_sum / nnelement if found else float("NaN")
 
 
 def get_averageexcitation(

From ee69972f55d78358918d6ccaa4b0cf9d045afad6 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 19:00:37 +0000
Subject: [PATCH 045/150] Update estimators.py

---
 artistools/estimators/estimators.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 3c4ef1922..3f8d885b4 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -444,8 +444,8 @@ def get_averaged_estimators(
                     estcollect.filter(pl.col("timestep") == timestep).filter(pl.col("modelgridindex") == mgi)[k].item(0)
                 )
                 if value is None:
-                    msg = f"{k} not found for timestep {timestep} and modelgridindex {mgi}"
-                    raise ValueError(msg)
+                    print(f"{k} not found for timestep {timestep} and modelgridindex {mgi}")
+                    value = 0.0
 
                 valuesum += value * tdelta
                 tdeltasum += tdelta

From 43ff17c440c270a5e5b4603b66d3dd9ac4d7bdc2 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 19:41:42 +0000
Subject: [PATCH 046/150] Update estimators.py

---
 artistools/estimators/estimators.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 3f8d885b4..74ea8d700 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -445,7 +445,7 @@ def get_averaged_estimators(
                 )
                 if value is None:
                     print(f"{k} not found for timestep {timestep} and modelgridindex {mgi}")
-                    value = 0.0
+                    continue
 
                 valuesum += value * tdelta
                 tdeltasum += tdelta
@@ -473,8 +473,7 @@ def get_averageionisation(estimatorstsmgi: pl.LazyFrame, atomic_number: int) ->
         found = True
         nnion = dfselected[key].item(0)
         if nnion is None:
-            nnion = f"WARNING: {key} is None"
-            nnion = 0.0
+            continue
 
         ionstage = at.decode_roman_numeral(key.removeprefix(f"nnion_{elsymb}_"))
         free_electron_weighted_pop_sum += nnion * (ionstage - 1)

From 40e213bed711bd5309d9a3bc9dc7de2a254cc2f7 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 20:35:38 +0000
Subject: [PATCH 047/150] Update estimators.py

---
 artistools/estimators/estimators.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 74ea8d700..6a9212444 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -444,13 +444,12 @@ def get_averaged_estimators(
                     estcollect.filter(pl.col("timestep") == timestep).filter(pl.col("modelgridindex") == mgi)[k].item(0)
                 )
                 if value is None:
-                    print(f"{k} not found for timestep {timestep} and modelgridindex {mgi}")
                     continue
 
                 valuesum += value * tdelta
                 tdeltasum += tdelta
 
-        dictout[k] = valuesum / tdeltasum
+        dictout[k] = valuesum / tdeltasum if tdeltasum > 0 else float("NaN")
 
     return dictout
 
@@ -469,7 +468,7 @@ def get_averageionisation(estimatorstsmgi: pl.LazyFrame, atomic_number: int) ->
 
     found = False
     popsum = 0.0
-    for key in dfselected.columns:
+    for key in dfselected.select(cs.starts_with(f"nnion_{elsymb}_")).columns:
         found = True
         nnion = dfselected[key].item(0)
         if nnion is None:

From eb002024b92a5249524fd959ad3b106c0415c737 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 21:00:58 +0000
Subject: [PATCH 048/150] Update inputmodel_misc.py

---
 artistools/inputmodel/inputmodel_misc.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/artistools/inputmodel/inputmodel_misc.py b/artistools/inputmodel/inputmodel_misc.py
index 24ceda0b4..49ff0898e 100644
--- a/artistools/inputmodel/inputmodel_misc.py
+++ b/artistools/inputmodel/inputmodel_misc.py
@@ -26,7 +26,7 @@ def read_modelfile_text(
 
     modelmeta: dict[str, t.Any] = {"headercommentlines": []}
 
-    if not printwarningsonly:
+    if not printwarningsonly and not getheadersonly:
         print(f"Reading {filename}")
 
     numheaderrows = 0
@@ -116,7 +116,7 @@ def read_modelfile_text(
 
         assert columns is not None
         if ncols_line_even == len(columns):
-            if not printwarningsonly:
+            if not printwarningsonly and not getheadersonly:
                 print("  model file is one line per cell")
             ncols_line_odd = 0
             onelinepercellformat = True
@@ -340,7 +340,7 @@ def get_modeldata_polars(
     dfmodel: pl.LazyFrame | None | pl.DataFrame = None
     if not getheadersonly and filenameparquet.is_file():
         if not printwarningsonly:
-            print(f"Reading data table from {filenameparquet}")
+            print(f"Reading model table from {filenameparquet}")
         try:
             dfmodel = pl.scan_parquet(filenameparquet)
         except pl.exceptions.ComputeError:

From 563cfa750657cd5c0ec2ae243685540e3d2a3e14 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 21:01:47 +0000
Subject: [PATCH 049/150] Update estimators.py

---
 artistools/estimators/estimators.py | 120 ++++++++++++----------------
 1 file changed, 51 insertions(+), 69 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 6a9212444..3646e49e6 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -243,7 +243,7 @@ def read_estimators_from_file(
     )
 
 
-def batched_it(iterable, n):
+def batched(iterable, n):
     """Batch data into iterators of length n. The last batch may be shorter."""
     # batched('ABCDEFG', 3) --> ABC DEF G
     if n < 1:
@@ -259,61 +259,53 @@ def batched_it(iterable, n):
         yield list(itertools.chain((first_el,), chunk_it))
 
 
-def read_estimators_in_folder_polars(
+def get_rankbatch_parquetfile(
     modelpath: Path,
     folderpath: Path,
-    match_modelgridindex: None | t.Sequence[int],
-) -> pl.LazyFrame:
-    mpiranklist = at.get_mpiranklist(modelpath, modelgridindex=match_modelgridindex, only_ranks_withgridcells=True)
-
-    mpirank_groups = list(batched_it(list(mpiranklist), 100))
-    group_parquetfiles = [
-        folderpath / f"estimators_{mpigroup[0]:04d}_{mpigroup[-1]:04d}.out.parquet.tmp" for mpigroup in mpirank_groups
-    ]
-    for mpigroup, parquetfilename in zip(mpirank_groups, group_parquetfiles):
-        if not parquetfilename.exists():
-            print(f"{parquetfilename.relative_to(modelpath.parent)} does not exist")
-            estfilepaths = []
-            for mpirank in mpigroup:
-                # not worth printing an error, because ranks with no cells to update do not produce an estimator file
-                with contextlib.suppress(FileNotFoundError):
-                    estfilepath = at.firstexisting(f"estimators_{mpirank:04d}.out", folder=folderpath, tryzipped=True)
-                    estfilepaths.append(estfilepath)
-
-            print(
-                f"Reading {len(list(estfilepaths))} estimator files from {folderpath.relative_to(Path(folderpath).parent)}"
-            )
-
-            processfile = partial(read_estimators_from_file)
-
-            pldf_group = None
-            with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
-                for pldf_file in pool.imap(processfile, estfilepaths):
-                    if pldf_group is None:
-                        pldf_group = pldf_file
-                    else:
-                        pldf_group = pl.concat([pldf_group, pldf_file], how="diagonal_relaxed")
-
-                pool.close()
-                pool.join()
-                pool.terminate()
-            assert pldf_group is not None
-            print(f"Writing {parquetfilename.relative_to(modelpath.parent)}")
-            pldf_group.write_parquet(parquetfilename, compression="zstd")
-
-    for parquetfilename in group_parquetfiles:
-        print(f"Reading {parquetfilename.relative_to(modelpath.parent)}")
-
-    return pl.concat(
-        [pl.scan_parquet(parquetfilename) for parquetfilename in group_parquetfiles], how="diagonal_relaxed"
-    ).sort(["timestep", "modelgridindex"])
+    mpiranks: t.Sequence[int],
+) -> Path:
+    parquetfilepath = folderpath / f"estimators_{mpiranks[0]:04d}_{mpiranks[-1]:04d}.out.parquet.tmp"
+
+    if not parquetfilepath.exists():
+        print(f"{parquetfilepath.relative_to(modelpath.parent)} does not exist")
+        estfilepaths = []
+        for mpirank in mpiranks:
+            # not worth printing an error, because ranks with no cells to update do not produce an estimator file
+            with contextlib.suppress(FileNotFoundError):
+                estfilepath = at.firstexisting(f"estimators_{mpirank:04d}.out", folder=folderpath, tryzipped=True)
+                estfilepaths.append(estfilepath)
+
+        print(
+            f"  reading {len(list(estfilepaths))} estimator files from {folderpath.relative_to(Path(folderpath).parent)}"
+        )
+
+        processfile = partial(read_estimators_from_file)
+
+        pldf_group = None
+        with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
+            for pldf_file in pool.imap(processfile, estfilepaths):
+                if pldf_group is None:
+                    pldf_group = pldf_file
+                else:
+                    pldf_group = pl.concat([pldf_group, pldf_file], how="diagonal_relaxed")
+
+            pool.close()
+            pool.join()
+            pool.terminate()
+
+        assert pldf_group is not None
+        print(f"  writing {parquetfilepath.relative_to(modelpath.parent)}")
+        pldf_group.write_parquet(parquetfilepath, compression="zstd")
+
+    print(f"Scanning {parquetfilepath.relative_to(modelpath.parent)}")
+
+    return parquetfilepath
 
 
 def read_estimators_polars(
     modelpath: Path | str = Path(),
     modelgridindex: None | int | t.Sequence[int] = None,
     timestep: None | int | t.Sequence[int] = None,
-    runfolder: None | str | Path = None,
 ) -> pl.LazyFrame:
     """Read estimator files into a dictionary of (timestep, modelgridindex): estimators.
 
@@ -354,27 +346,18 @@ def read_estimators_polars(
 
     # print(f" matching cells {match_modelgridindex} and timesteps {match_timestep}")
 
-    runfolders = at.get_runfolders(modelpath, timesteps=match_timestep) if runfolder is None else [Path(runfolder)]
-
-    parquetfiles = [folderpath / "estimators.out.parquet.tmp" for folderpath in runfolders]
+    mpiranklist = at.get_mpiranklist(modelpath, only_ranks_withgridcells=True)
+    mpirank_groups = list(batched(mpiranklist, 100))
+    runfolders = at.get_runfolders(modelpath, timesteps=match_timestep)
 
-    for folderpath, parquetfile in zip(runfolders, parquetfiles):
-        if not parquetfile.exists():
-            pldflazy = read_estimators_in_folder_polars(
-                modelpath,
-                folderpath,
-                match_modelgridindex=None,
-            )
-
-            print(f"Writing {parquetfile.relative_to(modelpath.parent)}")
-            pldflazy.collect().write_parquet(parquetfile, compression="zstd")
-
-    for folderpath, parquetfile in zip(runfolders, parquetfiles):
-        print(f"Scanning {parquetfile.relative_to(modelpath.parent)}")
+    parquetfiles = (
+        get_rankbatch_parquetfile(modelpath, runfolder, mpiranks)
+        for runfolder in runfolders
+        for mpiranks in mpirank_groups
+    )
 
-    pldflazy = pl.concat(
-        [pl.scan_parquet(parquetfilename) for parquetfilename in parquetfiles], how="diagonal_relaxed"
-    ).unique(["timestep", "modelgridindex"], maintain_order=True, keep="first")
+    pldflazy = pl.concat([pl.scan_parquet(pfile) for pfile in parquetfiles], how="diagonal_relaxed")
+    pldflazy = pldflazy.unique(["timestep", "modelgridindex"], maintain_order=True, keep="first")
 
     if match_modelgridindex is not None:
         pldflazy = pldflazy.filter(pl.col("modelgridindex").is_in(match_modelgridindex))
@@ -389,12 +372,11 @@ def read_estimators(
     modelpath: Path | str = Path(),
     modelgridindex: None | int | t.Sequence[int] = None,
     timestep: None | int | t.Sequence[int] = None,
-    runfolder: None | str | Path = None,
     keys: t.Collection[str] | None = None,
 ) -> dict[tuple[int, int], dict[str, t.Any]]:
     if isinstance(keys, str):
         keys = {keys}
-    pldflazy = read_estimators_polars(modelpath, modelgridindex, timestep, runfolder)
+    pldflazy = read_estimators_polars(modelpath, modelgridindex, timestep)
     estimators: dict[tuple[int, int], dict[str, t.Any]] = {}
     for estimtsmgi in pldflazy.collect().iter_rows(named=True):
         ts, mgi = estimtsmgi["timestep"], estimtsmgi["modelgridindex"]

From dbf935a9e2edc93c5247b9fd8749637304ea3d5a Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 21:26:56 +0000
Subject: [PATCH 050/150] Fix plot presentation and units

---
 artistools/estimators/estimators.py     |  7 ++--
 artistools/estimators/plotestimators.py | 44 ++++++++++++-------------
 2 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 3646e49e6..4657e6e7d 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -26,12 +26,12 @@
 def get_variableunits(key: str | None = None) -> str | dict[str, str]:
     variableunits = {
         "time": "days",
-        "gamma_NT": "/s",
-        "gamma_R_bfest": "/s",
+        "gamma_NT": "s^-1",
+        "gamma_R_bfest": "s^-1",
         "TR": "K",
         "Te": "K",
         "TJ": "K",
-        "nne": "e-/cm3",
+        "nne": "e^-/cm3",
         "heating": "erg/s/cm3",
         "heating_dep/total_dep": "Ratio",
         "cooling": "erg/s/cm3",
@@ -54,6 +54,7 @@ def get_variablelongunits(key: str | None = None) -> str | dict[str, str]:
 
 def get_dictlabelreplacements() -> dict[str, str]:
     return {
+        "nne": r"n$_{\rm e}$",
         "lognne": r"Log n$_{\rm e}$",
         "Te": r"T$_{\rm e}$",
         "TR": r"T$_{\rm R}$",
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 41dcffd6a..cd7104f44 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -354,7 +354,7 @@ def get_iontuple(ionstr):
 
         if seriestype == "populations":
             if args.ionpoptype == "absolute":
-                ax.set_ylabel("X$_{i}$ [/cm3]")
+                ax.set_ylabel(r"Number density $\left[\rm{cm}^{-3}\right]$")
             elif args.ionpoptype == "elpop":
                 # elsym = at.get_elsymbol(atomic_number)
                 ax.set_ylabel(r"X$_{i}$/X$_{\rm element}$")
@@ -501,13 +501,14 @@ def plot_series(
     assert len(xlist) - 1 == len(mgilist) == len(timestepslist)
     formattedvariablename = at.estimators.get_dictlabelreplacements().get(variablename, variablename)
     serieslabel = f"{formattedvariablename}"
-    if not nounits:
-        serieslabel += at.estimators.get_units_string(variablename)
+    units_string = at.estimators.get_units_string(variablename)
 
     if showlegend:
         linelabel = serieslabel
+        if not nounits:
+            linelabel += units_string
     else:
-        ax.set_ylabel(serieslabel)
+        ax.set_ylabel(serieslabel + units_string)
         linelabel = None
 
     ylist: list[float] = []
@@ -558,7 +559,7 @@ def get_xlist(
     timestepslist: t.Any,
     modelpath: str | Path,
     args: t.Any,
-) -> tuple[list[float | int], list[int | t.Sequence[int]], list[int | list[int]]]:
+) -> tuple[list[float | int], list[int | t.Sequence[int]], list[list[int]]]:
     xlist: t.Sequence[float | int]
     if xvariable in {"cellid", "modelgridindex"}:
         mgilist_out = [mgi for mgi in allnonemptymgilist if mgi <= args.xmax] if args.xmax >= 0 else allnonemptymgilist
@@ -624,21 +625,21 @@ def plot_subplot(
     # these three lists give the x value, modelgridex, and a list of timesteps (for averaging) for each plot of the plot
     assert len(xlist) - 1 == len(mgilist) == len(timestepslist)
     showlegend = False
-
+    seriescount = 0
     ylabel = None
     sameylabel = True
     for variablename in plotitems:
-        if not isinstance(variablename, str):
-            pass
-        elif ylabel is None:
-            ylabel = get_ylabel(variablename)
-        elif ylabel != get_ylabel(variablename):
-            sameylabel = False
-            break
+        if isinstance(variablename, str):
+            seriescount += 1
+            if ylabel is None:
+                ylabel = get_ylabel(variablename)
+            elif ylabel != get_ylabel(variablename):
+                sameylabel = False
+                break
 
     for plotitem in plotitems:
         if isinstance(plotitem, str):
-            showlegend = len(plotitems) > 1 or len(plotitem) > 20
+            showlegend = seriescount > 1 or len(plotitem) > 20 or not sameylabel
             plot_series(
                 ax,
                 xlist,
@@ -812,9 +813,7 @@ def make_plot(
             **plotkwargs,
         )
 
-    if (
-        len(set(mgilist)) == 1 and not isinstance(timestepslist[0], int) and len(timestepslist[0]) > 1
-    ):  # single grid cell versus time plot
+    if len(set(mgilist)) == 1 and len(timestepslist[0]) > 1:  # single grid cell versus time plot
         figure_title = f"{modelname}\nCell {mgilist[0]}"
 
         defaultoutputfile = Path("plotestimators_cell{modelgridindex:03d}.pdf")
@@ -835,13 +834,13 @@ def make_plot(
                 .collect()
                 .item(0)
             )
-            figure_title = f"{modelname}\nTimestep {timestepslist[0]} ({tdays:.2f}d)"
+            figure_title = f"{modelname}\nTimestep {timestepslist[0][0]} ({tdays:.2f}d)"
         elif args.multiplot:
             assert isinstance(timestepslist[0], int)
             timedays = float(at.get_timestep_time(modelpath, timestepslist[0]))
-            figure_title = f"{modelname}\nTimestep {timestepslist[0]} ({timedays:.2f}d)"
+            figure_title = f"{modelname}\nTimestep {timestepslist[0][0]} ({timedays:.2f}d)"
         else:
-            figure_title = f"{modelname}\nTimestep {timestepslist[0]} ({timeavg:.2f}d)"
+            figure_title = f"{modelname}\nTimestep {timestepslist[0][0]} ({timeavg:.2f}d)"
 
         defaultoutputfile = Path("plotestimators_ts{timestep:02d}_{timeavg:.2f}d.pdf")
         if Path(args.outputfile).is_dir():
@@ -851,7 +850,7 @@ def make_plot(
         outfilename = str(args.outputfile).format(timestep=timestepslist[0][0], timeavg=timeavg)
 
     if not args.notitle:
-        axes[0].set_title(figure_title, fontsize=11)
+        axes[0].set_title(figure_title, fontsize=8)
     # plt.suptitle(figure_title, fontsize=11, verticalalignment='top')
 
     if args.write_data:
@@ -1090,11 +1089,12 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         ["nne", ["_ymin", 1e5], ["_ymax", 1e11]],
         # ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 26000]],
         ["Te"],
+        # ["Te", "TR"],
         [["averageionisation", ["Sr"]]],
         # [['averageexcitation', ['Fe II', 'Fe III']]],
         # [["populations", ["Sr90", "Sr91", "Sr92", "Sr93", "Sr94"]]],
         #  ['_ymin', 1e-3], ['_ymax', 5]],
-        [["populations", ["Fe", "Co", "Ni", "Sr", "Nd"]]],
+        [["populations", ["Sr"]]],
         # [['populations', ['He I', 'He II', 'He III']]],
         # [['populations', ['C I', 'C II', 'C III', 'C IV', 'C V']]],
         # [['populations', ['O I', 'O II', 'O III', 'O IV']]],

From 1f8bf2a7ce3c07f070719832a76c69498b072a49 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 21:38:43 +0000
Subject: [PATCH 051/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index cd7104f44..042b6c992 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -469,8 +469,9 @@ def get_iontuple(ionstr):
         ylist.insert(0, ylist[0])
 
         xlist, ylist = at.estimators.apply_filters(xlist, ylist, args)
-
-        ax.plot(xlist, ylist, linewidth=linewidth, label=plotlabel, color=color, dashes=dashes, **plotkwargs)
+        if plotkwargs["linestyle"] != "None":
+            plotkwargs["dashes"] = dashes
+        ax.plot(xlist, ylist, linewidth=linewidth, label=plotlabel, color=color, **plotkwargs)
         prev_atomic_number = atomic_number
         plotted_something = True
 
@@ -578,6 +579,8 @@ def get_xlist(
         timestepslist_out = timestepslist
     elif xvariable in {"velocity", "beta"}:
         dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["vel_r_mid"])
+        if modelmeta["dimensions"] > 1:
+            args.markersonly = True
         if modelmeta["vmax_cmps"] > 0.3 * 29979245800:
             args.x = "beta"
             xvariable = "beta"
@@ -983,7 +986,9 @@ def addargs(parser: argparse.ArgumentParser) -> None:
 
     parser.add_argument("--hidexlabel", action="store_true", help="Hide the bottom horizontal axis label")
 
-    parser.add_argument("--markersonly", action="store_true", help="Plot markers instead of lines")
+    parser.add_argument(
+        "--markersonly", action="store_true", help="Plot markers instead of lines (always set for 2D and 3D)"
+    )
 
     parser.add_argument("-filtermovingavg", type=int, default=0, help="Smoothing length (1 is same as none)")
 
@@ -1094,7 +1099,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         # [['averageexcitation', ['Fe II', 'Fe III']]],
         # [["populations", ["Sr90", "Sr91", "Sr92", "Sr93", "Sr94"]]],
         #  ['_ymin', 1e-3], ['_ymax', 5]],
-        [["populations", ["Sr"]]],
+        [["populations", ["Sr II"]]],
         # [['populations', ['He I', 'He II', 'He III']]],
         # [['populations', ['C I', 'C II', 'C III', 'C IV', 'C V']]],
         # [['populations', ['O I', 'O II', 'O III', 'O IV']]],

From 587cc76f3f51d37800180fc204dc3ad2de5e009a Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 21:40:38 +0000
Subject: [PATCH 052/150] Update commands.py

---
 artistools/commands.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/artistools/commands.py b/artistools/commands.py
index ed03bd164..a58d9f3cb 100644
--- a/artistools/commands.py
+++ b/artistools/commands.py
@@ -19,7 +19,7 @@
     "plotinitialcomposition": ("initial_composition", "main"),
     "plotlightcurves": ("lightcurve.plotlightcurve", "main"),
     "plotlinefluxes": ("linefluxes", "main"),
-    "plotmodeldensity": ("inputmodel.plotdensity", "main"),
+    "plotdensity": ("inputmodel.plotdensity", "main"),
     "plotmodeldeposition": ("deposition", "main"),
     "plotmacroatom": ("macroatom", "main"),
     "plotnltepops": ("nltepops.plotnltepops", "main"),

From 1907affe1d4116891a7da47a2d46ff8196895a15 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Thu, 9 Nov 2023 21:42:34 +0000
Subject: [PATCH 053/150] Update radfield.py

---
 artistools/radfield.py | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/artistools/radfield.py b/artistools/radfield.py
index 3d9a167b1..a15b0ffe7 100755
--- a/artistools/radfield.py
+++ b/artistools/radfield.py
@@ -311,9 +311,7 @@ def sigma_bf(nu):
 
 def get_kappa_bf_ion(atomic_number, lower_ionstage, modelgridindex, timestep, modelpath, arr_nu_hz, max_levels):
     adata = at.atomic.get_levels(modelpath, get_photoionisations=True)
-    estimators = at.estimators.read_estimators(
-        modelpath, timestep=timestep, modelgridindex=modelgridindex, use_polars=False
-    )
+    estimators = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex)
     T_e = estimators[(timestep, modelgridindex)]["Te"]
 
     ion_data = adata.query("Z == @atomic_number and ionstage == @lower_ionstage").iloc[0]
@@ -355,9 +353,9 @@ def get_recombination_emission(
     upper_ion_data = adata.query("Z == @atomic_number and ionstage == @upper_ionstage").iloc[0]
     lower_ion_data = adata.query("Z == @atomic_number and ionstage == @lower_ionstage").iloc[0]
 
-    estimtsmgi = at.estimators.read_estimators(
-        modelpath, timestep=timestep, modelgridindex=modelgridindex, use_polars=False
-    )[(timestep, modelgridindex)]
+    estimtsmgi = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex)[
+        (timestep, modelgridindex)
+    ]
 
     upperionstr = at.get_ionstring(atomic_number, upper_ionstage, sep="_", style="spectral")
     upperionpopdensity = estimtsmgi[f"nnion_{upperionstr}"]
@@ -461,9 +459,7 @@ def get_recombination_emission(
 
 def get_ion_gamma_dnu(modelpath, modelgridindex, timestep, atomic_number, ionstage, arr_nu_hz, J_nu_arr, max_levels):
     """Calculate the contribution to the photoionisation rate coefficient per J_nu at each frequency nu for an ion."""
-    estimators = at.estimators.read_estimators(
-        modelpath, timestep=timestep, modelgridindex=modelgridindex, use_polars=False
-    )
+    estimators = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex)
 
     T_e = estimators[(timestep, modelgridindex)]["Te"]
     T_R = estimators[(timestep, modelgridindex)]["TR"]

From 4e41ae2e7305aa2e1f98401d765477d376d32504 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 12:07:28 +0000
Subject: [PATCH 054/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 042b6c992..c26da5f30 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -469,7 +469,7 @@ def get_iontuple(ionstr):
         ylist.insert(0, ylist[0])
 
         xlist, ylist = at.estimators.apply_filters(xlist, ylist, args)
-        if plotkwargs["linestyle"] != "None":
+        if plotkwargs.get("linestyle", "solid") != "None":
             plotkwargs["dashes"] = dashes
         ax.plot(xlist, ylist, linewidth=linewidth, label=plotlabel, color=color, **plotkwargs)
         prev_atomic_number = atomic_number

From 32dcf128740cfa3f3957f55060be3db6715a7a0f Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 12:18:40 +0000
Subject: [PATCH 055/150] Remove trailing commas

---
 artistools/commands.py | 25 +++++--------------------
 1 file changed, 5 insertions(+), 20 deletions(-)

diff --git a/artistools/commands.py b/artistools/commands.py
index a58d9f3cb..73e1fd719 100644
--- a/artistools/commands.py
+++ b/artistools/commands.py
@@ -58,21 +58,12 @@ def get_commandlist() -> dict[str, tuple[str, str]]:
     return {
         "at": ("artistools", "main"),
         "artistools": ("artistools", "main"),
-        "makeartismodel1dslicefromcone": (
-            "artistools.inputmodel.slice1dfromconein3dmodel",
-            "main",
-        ),
+        "makeartismodel1dslicefromcone": ("artistools.inputmodel.slice1dfromconein3dmodel", "main"),
         "makeartismodel": ("artistools.inputmodel.makeartismodel", "main"),
         "plotartismodeldensity": ("artistools.inputmodel.plotdensity", "main"),
         "plotartismodeldeposition": ("artistools.deposition", "main"),
-        "plotartisestimators": (
-            "artistools.estimators.plotestimators",
-            "main",
-        ),
-        "plotartislightcurve": (
-            "artistools.lightcurve.plotlightcurve",
-            "main",
-        ),
+        "plotartisestimators": ("artistools.estimators.plotestimators", "main"),
+        "plotartislightcurve": ("artistools.lightcurve.plotlightcurve", "main"),
         "plotartislinefluxes": ("artistools.linefluxes", "main"),
         "plotartismacroatom": ("artistools.macroatom", "main"),
         "plotartisnltepops": ("artistools.nltepops.plotnltepops", "main"),
@@ -80,14 +71,8 @@ def get_commandlist() -> dict[str, tuple[str, str]]:
         "plotartisradfield": ("artistools.radfield", "main"),
         "plotartisspectrum": ("artistools.spectra.plotspectra", "main"),
         "plotartistransitions": ("artistools.transitions", "main"),
-        "plotartisinitialcomposition": (
-            "artistools.initial_composition",
-            "main",
-        ),
-        "plotartisviewingangles": (
-            "artistools.viewing_angles_visualization",
-            "main",
-        ),
+        "plotartisinitialcomposition": ("artistools.initial_composition", "main"),
+        "plotartisviewingangles": ("artistools.viewing_angles_visualization", "main"),
     }
 
 

From 6ea1418e792185cfb64a8052087c8e306fcb87c9 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 12:19:45 +0000
Subject: [PATCH 056/150] Update commands.py

---
 artistools/commands.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/artistools/commands.py b/artistools/commands.py
index 73e1fd719..7050eb58b 100644
--- a/artistools/commands.py
+++ b/artistools/commands.py
@@ -20,7 +20,7 @@
     "plotlightcurves": ("lightcurve.plotlightcurve", "main"),
     "plotlinefluxes": ("linefluxes", "main"),
     "plotdensity": ("inputmodel.plotdensity", "main"),
-    "plotmodeldeposition": ("deposition", "main"),
+    "plotdeposition": ("deposition", "main"),
     "plotmacroatom": ("macroatom", "main"),
     "plotnltepops": ("nltepops.plotnltepops", "main"),
     "plotnonthermal": ("nonthermal.plotnonthermal", "main"),
@@ -60,8 +60,8 @@ def get_commandlist() -> dict[str, tuple[str, str]]:
         "artistools": ("artistools", "main"),
         "makeartismodel1dslicefromcone": ("artistools.inputmodel.slice1dfromconein3dmodel", "main"),
         "makeartismodel": ("artistools.inputmodel.makeartismodel", "main"),
-        "plotartismodeldensity": ("artistools.inputmodel.plotdensity", "main"),
-        "plotartismodeldeposition": ("artistools.deposition", "main"),
+        "plotartisdensity": ("artistools.inputmodel.plotdensity", "main"),
+        "plotartisdeposition": ("artistools.deposition", "main"),
         "plotartisestimators": ("artistools.estimators.plotestimators", "main"),
         "plotartislightcurve": ("artistools.lightcurve.plotlightcurve", "main"),
         "plotartislinefluxes": ("artistools.linefluxes", "main"),

From cb189d717b81ce3e12379557a62803d9262bb3bf Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 12:24:34 +0000
Subject: [PATCH 057/150] Move plotinitialcomposition to inputmodel submodule

---
 artistools/__init__.py                                      | 1 -
 artistools/commands.py                                      | 4 ++--
 artistools/inputmodel/__init__.py                           | 1 +
 .../plotinitialcomposition.py}                              | 0
 artistools/inputmodel/test_inputmodel.py                    | 6 ++++++
 artistools/test_artistools.py                               | 4 ----
 6 files changed, 9 insertions(+), 7 deletions(-)
 rename artistools/{initial_composition.py => inputmodel/plotinitialcomposition.py} (100%)

diff --git a/artistools/__init__.py b/artistools/__init__.py
index 9b59c2d0e..f9ff09b45 100644
--- a/artistools/__init__.py
+++ b/artistools/__init__.py
@@ -10,7 +10,6 @@
 from artistools import commands
 from artistools import deposition
 from artistools import estimators
-from artistools import initial_composition
 from artistools import inputmodel
 from artistools import lightcurve
 from artistools import macroatom
diff --git a/artistools/commands.py b/artistools/commands.py
index 7050eb58b..bca0b11e4 100644
--- a/artistools/commands.py
+++ b/artistools/commands.py
@@ -16,7 +16,7 @@
     "makeartismodelfromparticlegridmap": ("inputmodel.modelfromhydro", "main"),
     "maptogrid": ("inputmodel.maptogrid", "main"),
     "plotestimators": ("estimators.plotestimators", "main"),
-    "plotinitialcomposition": ("initial_composition", "main"),
+    "plotinitialcomposition": ("inputmodel.plotinitialcomposition", "main"),
     "plotlightcurves": ("lightcurve.plotlightcurve", "main"),
     "plotlinefluxes": ("linefluxes", "main"),
     "plotdensity": ("inputmodel.plotdensity", "main"),
@@ -71,7 +71,7 @@ def get_commandlist() -> dict[str, tuple[str, str]]:
         "plotartisradfield": ("artistools.radfield", "main"),
         "plotartisspectrum": ("artistools.spectra.plotspectra", "main"),
         "plotartistransitions": ("artistools.transitions", "main"),
-        "plotartisinitialcomposition": ("artistools.initial_composition", "main"),
+        "plotartisinitialcomposition": ("artistools.inputmodel.plotinitialcomposition", "main"),
         "plotartisviewingangles": ("artistools.viewing_angles_visualization", "main"),
     }
 
diff --git a/artistools/inputmodel/__init__.py b/artistools/inputmodel/__init__.py
index a7d78d148..3dfb58603 100644
--- a/artistools/inputmodel/__init__.py
+++ b/artistools/inputmodel/__init__.py
@@ -8,6 +8,7 @@
 from artistools.inputmodel import modelfromhydro
 from artistools.inputmodel import opacityinputfile
 from artistools.inputmodel import plotdensity
+from artistools.inputmodel import plotinitialcomposition
 from artistools.inputmodel import rprocess_from_trajectory
 from artistools.inputmodel import slice1dfromconein3dmodel
 from artistools.inputmodel import to_tardis
diff --git a/artistools/initial_composition.py b/artistools/inputmodel/plotinitialcomposition.py
similarity index 100%
rename from artistools/initial_composition.py
rename to artistools/inputmodel/plotinitialcomposition.py
diff --git a/artistools/inputmodel/test_inputmodel.py b/artistools/inputmodel/test_inputmodel.py
index 87aeb5724..07ee2501e 100644
--- a/artistools/inputmodel/test_inputmodel.py
+++ b/artistools/inputmodel/test_inputmodel.py
@@ -228,6 +228,12 @@ def test_plotdensity() -> None:
     at.inputmodel.plotdensity.main(argsraw=[], modelpath=[modelpath], outputpath=outputpath)
 
 
+def test_plotinitialcomposition() -> None:
+    at.inputmodel.plotinitialcomposition.main(
+        argsraw=["-modelpath", str(modelpath_3d), "-o", str(outputpath), "rho", "Fe"]
+    )
+
+
 def test_save_load_3d_model() -> None:
     clear_modelfiles()
     lzdfmodel, modelmeta = at.inputmodel.get_empty_3d_model(ncoordgrid=50, vmax=1000, t_model_init_days=1)
diff --git a/artistools/test_artistools.py b/artistools/test_artistools.py
index 2cb30436a..1d55c23ee 100755
--- a/artistools/test_artistools.py
+++ b/artistools/test_artistools.py
@@ -63,10 +63,6 @@ def test_deposition() -> None:
     at.deposition.main(argsraw=[], modelpath=modelpath)
 
 
-def test_initial_composition() -> None:
-    at.initial_composition.main(argsraw=["-modelpath", str(modelpath_3d), "-o", str(outputpath), "rho", "Fe"])
-
-
 def test_get_inputparams() -> None:
     inputparams = at.get_inputparams(modelpath)
     dicthash = hashlib.sha256(str(sorted(inputparams.items())).encode("utf-8")).hexdigest()

From c21821c6fc2c7448791a8bcd93b5fb5dfc49e9d6 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 12:26:54 +0000
Subject: [PATCH 058/150] Fixup

---
 artistools/estimators/plot3destimators_classic.py | 2 +-
 artistools/estimators/plotestimators.py           | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/artistools/estimators/plot3destimators_classic.py b/artistools/estimators/plot3destimators_classic.py
index f8d070a90..3e3357426 100644
--- a/artistools/estimators/plot3destimators_classic.py
+++ b/artistools/estimators/plot3destimators_classic.py
@@ -37,7 +37,7 @@ def get_modelgridcells_along_axis(modelpath, args=None):
 def get_modelgridcells_2D_slice(modeldata, modelpath) -> list[int]:
     sliceaxis: t.Literal["x", "y", "z"] = "x"
 
-    slicedata = at.initial_composition.get_2D_slice_through_3d_model(modeldata, sliceaxis)
+    slicedata = at.inputmodel.plotinitialcomposition.get_2D_slice_through_3d_model(modeldata, sliceaxis)
     return get_mgi_of_modeldata(slicedata, modelpath)
 
 
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index c26da5f30..0d9ebac49 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -59,7 +59,7 @@ def plot_init_abundances(
     if seriestype == "initabundances":
         mergemodelabundata, _ = at.inputmodel.get_modeldata(modelpath, get_elemabundances=True)
     elif seriestype == "initmasses":
-        mergemodelabundata = at.initial_composition.get_model_abundances_Msun_1D(modelpath)
+        mergemodelabundata = at.inputmodel.plotinitialcomposition.get_model_abundances_Msun_1D(modelpath)
     else:
         raise AssertionError
 

From ca80a1e1a32f92089f2d43922389e3e7dea3fcc4 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 12:28:38 +0000
Subject: [PATCH 059/150] Remove unused avgadjcells parameter

---
 artistools/estimators/estimators.py | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 4657e6e7d..5ea3be00b 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -396,7 +396,6 @@ def get_averaged_estimators(
     timesteps: int | t.Sequence[int],
     modelgridindex: int,
     keys: str | list | None,
-    avgadjcells: int = 0,
 ) -> dict[str, t.Any]:
     """Get the average of estimators[(timestep, modelgridindex)][keys[0]]...[keys[-1]] across timesteps."""
     modelgridindex = int(modelgridindex)
@@ -410,11 +409,11 @@ def get_averaged_estimators(
 
     dictout = {}
     tdeltas = at.get_timestep_times(modelpath, loc="delta")
-    mgilist = list(range(modelgridindex - avgadjcells, modelgridindex + avgadjcells + 1))
+
     estcollect = (
         estimators.lazy()
         .filter(pl.col("timestep").is_in(timesteps))
-        .filter(pl.col("modelgridindex").is_in(mgilist))
+        .filter(pl.col("modelgridindex") == modelgridindex)
         .select({*keys, "timestep", "modelgridindex"})
         .collect()
     )
@@ -422,15 +421,16 @@ def get_averaged_estimators(
         valuesum = 0
         tdeltasum = 0
         for timestep, tdelta in zip(timesteps, tdeltas):
-            for mgi in mgilist:
-                value = (
-                    estcollect.filter(pl.col("timestep") == timestep).filter(pl.col("modelgridindex") == mgi)[k].item(0)
-                )
-                if value is None:
-                    continue
-
-                valuesum += value * tdelta
-                tdeltasum += tdelta
+            value = (
+                estcollect.filter(pl.col("timestep") == timestep)
+                .filter(pl.col("modelgridindex") == modelgridindex)[k]
+                .item(0)
+            )
+            if value is None:
+                continue
+
+            valuesum += value * tdelta
+            tdeltasum += tdelta
 
         dictout[k] = valuesum / tdeltasum if tdeltasum > 0 else float("NaN")
 

From 8cf387b6debeb931782f5e8ed0e0928216168175 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 13:09:29 +0000
Subject: [PATCH 060/150] Add more type hints

---
 artistools/estimators/estimators.py     |  11 +--
 artistools/estimators/plotestimators.py | 116 +++++++++++-------------
 artistools/misc.py                      |  10 +-
 3 files changed, 62 insertions(+), 75 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 5ea3be00b..f6e2f2749 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -3,7 +3,6 @@
 
 Examples are temperatures, populations, and heating/cooling rates.
 """
-
 import argparse
 import contextlib
 import itertools
@@ -66,11 +65,9 @@ def get_dictlabelreplacements() -> dict[str, str]:
 
 
 def apply_filters(
-    xlist: list[float] | np.ndarray, ylist: list[float] | np.ndarray, args: argparse.Namespace
-) -> tuple[list[float] | np.ndarray, list[float] | np.ndarray]:
-    filterfunc = at.get_filterfunc(args)
-
-    if filterfunc is not None:
+    xlist: t.Sequence[float] | np.ndarray, ylist: t.Sequence[float] | np.ndarray, args: argparse.Namespace
+) -> tuple[t.Any, t.Any]:
+    if (filterfunc := at.get_filterfunc(args)) is not None:
         ylist = filterfunc(ylist)
 
     return xlist, ylist
@@ -397,7 +394,7 @@ def get_averaged_estimators(
     modelgridindex: int,
     keys: str | list | None,
 ) -> dict[str, t.Any]:
-    """Get the average of estimators[(timestep, modelgridindex)][keys[0]]...[keys[-1]] across timesteps."""
+    """Get the average across timsteps for a cell."""
     modelgridindex = int(modelgridindex)
     if isinstance(timesteps, int):
         timesteps = [timesteps]
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 0d9ebac49..dd4da6215 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -189,17 +189,17 @@ def plot_average_ionisation_excitation(
 
 
 def plot_levelpop(
-    ax,
-    xlist,
-    seriestype,
-    params,
-    timestepslist,
-    mgilist,
-    estimators,
-    modelpath,
-    dfalldata=None,
-    args=None,
-    **plotkwargs,
+    ax: plt.Axes,
+    xlist: t.Sequence[int | float] | np.ndarray,
+    seriestype: str,
+    params: t.Sequence[str],
+    timestepslist: t.Sequence[t.Sequence[int]],
+    mgilist: t.Sequence[int],
+    estimators: pl.LazyFrame | pl.DataFrame,
+    modelpath: str | Path,
+    dfalldata: pd.DataFrame | None,
+    args: argparse.Namespace,
+    **plotkwargs: t.Any,
 ):
     if seriestype == "levelpopulation_dn_on_dvel":
         ax.set_ylabel("dN/dV [{}km$^{{-1}}$ s]")
@@ -277,17 +277,17 @@ def plot_levelpop(
 
 
 def plot_multi_ion_series(
-    ax,
-    xlist,
-    seriestype,
-    ionlist,
-    timestepslist,
-    mgilist,
-    estimators,
-    modelpath,
-    dfalldata,
-    args,
-    **plotkwargs,
+    ax: plt.Axes,
+    xlist: t.Sequence[int | float] | np.ndarray,
+    seriestype: str,
+    ionlist: t.Sequence[str],
+    timestepslist: t.Sequence[t.Sequence[int]],
+    mgilist: t.Sequence[int],
+    estimators: pl.LazyFrame | pl.DataFrame,
+    modelpath: str | Path,
+    dfalldata: pd.DataFrame | None,
+    args: argparse.Namespace,
+    **plotkwargs: t.Any,
 ):
     """Plot an ion-specific property, e.g., populations."""
     assert len(xlist) - 1 == len(mgilist) == len(timestepslist)
@@ -314,26 +314,29 @@ def get_iontuple(ionstr):
 
     missingions = set()
     try:
-        if args.classicartis:
-            compositiondata = at.estimators.estimators_classic.get_atomic_composition(modelpath)
-        else:
+        if not args.classicartis:
             compositiondata = at.get_composition_data(modelpath)
-        for atomic_number, ionstage in iontuplelist:
-            if (
-                not hasattr(ionstage, "lower")
-                and not args.classicartis
-                and compositiondata.query(
-                    "Z == @atomic_number & lowermost_ionstage <= @ionstage & uppermost_ionstage >= @ionstage"
-                ).empty
-            ):
-                missingions.add((atomic_number, ionstage))
+            for atomic_number, ionstage in iontuplelist:
+                if (
+                    not hasattr(ionstage, "lower")
+                    and not args.classicartis
+                    and compositiondata.query(
+                        "Z == @atomic_number & lowermost_ionstage <= @ionstage & uppermost_ionstage >= @ionstage"
+                    ).empty
+                ):
+                    missingions.add((atomic_number, ionstage))
 
     except FileNotFoundError:
         print("WARNING: Could not read an ARTIS compositiondata.txt file")
+        estim_mgits = estimators.filter(pl.col("timestep") == timestepslist[0][0]).filter(
+            pl.col("modelgridindex") == mgilist[0]
+        )
         for atomic_number, ionstage in iontuplelist:
-            mgits = (timestepslist[0][0], mgilist[0])
             ionstr = at.get_ionstring(atomic_number, ionstage, sep="_", style="spectral")
-            if f"nnion_{ionstr}" not in estimators[mgits]:
+            if (
+                f"nnion_{ionstr}" not in estim_mgits.columns
+                or estim_mgits.select(f"nnion_{ionstr}").lazy().collect()[f"nnion_{ionstr}"].is_null().all()
+            ):
                 missingions.add((atomic_number, ionstage))
 
     if missingions:
@@ -365,6 +368,7 @@ def get_iontuple(ionstr):
         else:
             ax.set_ylabel(at.estimators.get_dictlabelreplacements().get(seriestype, seriestype))
 
+        print(f"Plotting {seriestype} {ionstr}")
         ylist = []
         for modelgridindex, timesteps in zip(mgilist, timestepslist):
             if seriestype == "populations":
@@ -429,6 +433,7 @@ def get_iontuple(ionstr):
         color = get_elemcolor(atomic_number=atomic_number)
 
         # linestyle = ['-.', '-', '--', (0, (4, 1, 1, 1)), ':'] + [(0, x) for x in dashes_list][ionstage - 1]
+        dashes: tuple[float, ...]
         if ionstage == "ALL":
             dashes = ()
             linewidth = 1.0
@@ -485,17 +490,17 @@ def get_iontuple(ionstr):
 
 
 def plot_series(
-    ax,
-    xlist,
-    variablename,
+    ax: plt.Axes,
+    xlist: t.Sequence[int | float] | np.ndarray,
+    variablename: str,
     showlegend: bool,
-    timestepslist,
-    mgilist,
+    timestepslist: t.Sequence[t.Sequence[int]],
+    mgilist: t.Sequence[int],
     modelpath: str | Path,
-    estimators,
+    estimators: pl.LazyFrame | pl.DataFrame,
     args: argparse.Namespace,
     nounits: bool = False,
-    dfalldata=None,
+    dfalldata: pd.DataFrame | None = None,
     **plotkwargs,
 ):
     """Plot something like Te or TR."""
@@ -511,25 +516,15 @@ def plot_series(
     else:
         ax.set_ylabel(serieslabel + units_string)
         linelabel = None
-
+    print(f"Plotting {variablename}")
     ylist: list[float] = []
     for modelgridindex, timesteps in zip(mgilist, timestepslist):
-        try:
-            yvalue = at.estimators.get_averaged_estimators(
-                modelpath, estimators, timesteps, modelgridindex, variablename
-            )[variablename]
-            ylist.append(yvalue)
-        except KeyError:
-            if (timesteps[0], modelgridindex) in estimators:
-                print(f"Undefined variable: {variablename} in cell {modelgridindex}")
-            else:
-                print(f"No data for cell {modelgridindex}")
-            raise
+        yvalue = at.estimators.get_averaged_estimators(modelpath, estimators, timesteps, modelgridindex, variablename)[
+            variablename
+        ]
+        ylist.append(yvalue)
 
-    try:
-        if math.log10(max(ylist) / min(ylist)) > 2:
-            ax.set_yscale("log")
-    except ZeroDivisionError:
+    if math.log10(max(ylist) / min(ylist)) > 2 or min(ylist) == 0:
         ax.set_yscale("log")
 
     dictcolors = {
@@ -538,9 +533,6 @@ def plot_series(
         # 'cooling_adiabatic': 'blue'
     }
 
-    # print out the data to stdout. Maybe want to add a CSV export option at some point?
-    # print(f'#cellidorvelocity {variablename}\n' + '\n'.join([f'{x}  {y}' for x, y in zip(xlist, ylist)]))
-
     if dfalldata is not None:
         dfalldata[variablename] = ylist
 
diff --git a/artistools/misc.py b/artistools/misc.py
index 15e06a17d..a19597765 100644
--- a/artistools/misc.py
+++ b/artistools/misc.py
@@ -890,16 +890,14 @@ def add_derived_metadata(metadata: dict[str, t.Any]) -> dict[str, t.Any]:
     return {}
 
 
-def get_filterfunc(
-    args: argparse.Namespace, mode: str = "interp"
-) -> t.Callable[[list[float] | np.ndarray], np.ndarray] | None:
+def get_filterfunc(args: argparse.Namespace, mode: str = "interp") -> t.Callable | None:
     """Use command line arguments to determine the appropriate filter function."""
-    filterfunc: t.Callable[[list[float] | np.ndarray], np.ndarray] | None = None
+    filterfunc = None
     dictargs = vars(args)
 
     if dictargs.get("filtermovingavg", False):
 
-        def movavgfilterfunc(ylist: list[float] | np.ndarray) -> np.ndarray:
+        def movavgfilterfunc(ylist: t.Any) -> t.Any:
             n = args.filtermovingavg
             arr_padded = np.pad(ylist, (n // 2, n - 1 - n // 2), mode="edge")
             return np.convolve(arr_padded, np.ones((n,)) / n, mode="valid")
@@ -912,7 +910,7 @@ def movavgfilterfunc(ylist: list[float] | np.ndarray) -> np.ndarray:
 
         window_length, poly_order = (int(x) for x in args.filtersavgol)
 
-        def savgolfilterfunc(ylist: list[float] | np.ndarray) -> np.ndarray:
+        def savgolfilterfunc(ylist: t.Any) -> t.Any:
             return scipy.signal.savgol_filter(ylist, window_length, poly_order, mode=mode)
 
         assert filterfunc is None

From 880ad772cae8acdf4ee6a253f07364f5e3ae529c Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 13:14:43 +0000
Subject: [PATCH 061/150] Add more type hints

---
 artistools/estimators/estimators.py     |  4 ++--
 artistools/estimators/plotestimators.py | 23 ++++++++++++++++-------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index f6e2f2749..31cf3ad31 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -391,11 +391,11 @@ def get_averaged_estimators(
     modelpath: Path | str,
     estimators: pl.LazyFrame | pl.DataFrame,
     timesteps: int | t.Sequence[int],
-    modelgridindex: int,
+    modelgridindex: int | t.Sequence[int],
     keys: str | list | None,
 ) -> dict[str, t.Any]:
     """Get the average across timsteps for a cell."""
-    modelgridindex = int(modelgridindex)
+    assert isinstance(modelgridindex, int)
     if isinstance(timesteps, int):
         timesteps = [timesteps]
 
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index dd4da6215..c9f36026d 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -194,7 +194,7 @@ def plot_levelpop(
     seriestype: str,
     params: t.Sequence[str],
     timestepslist: t.Sequence[t.Sequence[int]],
-    mgilist: t.Sequence[int],
+    mgilist: t.Sequence[int | t.Sequence[int]],
     estimators: pl.LazyFrame | pl.DataFrame,
     modelpath: str | Path,
     dfalldata: pd.DataFrame | None,
@@ -282,7 +282,7 @@ def plot_multi_ion_series(
     seriestype: str,
     ionlist: t.Sequence[str],
     timestepslist: t.Sequence[t.Sequence[int]],
-    mgilist: t.Sequence[int],
+    mgilist: t.Sequence[int | t.Sequence[int]],
     estimators: pl.LazyFrame | pl.DataFrame,
     modelpath: str | Path,
     dfalldata: pd.DataFrame | None,
@@ -495,14 +495,14 @@ def plot_series(
     variablename: str,
     showlegend: bool,
     timestepslist: t.Sequence[t.Sequence[int]],
-    mgilist: t.Sequence[int],
+    mgilist: t.Sequence[int | t.Sequence[int]],
     modelpath: str | Path,
     estimators: pl.LazyFrame | pl.DataFrame,
     args: argparse.Namespace,
     nounits: bool = False,
     dfalldata: pd.DataFrame | None = None,
-    **plotkwargs,
-):
+    **plotkwargs: t.Any,
+) -> None:
     """Plot something like Te or TR."""
     assert len(xlist) - 1 == len(mgilist) == len(timestepslist)
     formattedvariablename = at.estimators.get_dictlabelreplacements().get(variablename, variablename)
@@ -614,7 +614,16 @@ def get_xlist(
 
 
 def plot_subplot(
-    ax, timestepslist, xlist, plotitems, mgilist, modelpath, estimators, dfalldata=None, args=None, **plotkwargs
+    ax: plt.Axes,
+    timestepslist: list[list[int]],
+    xlist: list[float | int],
+    plotitems: list[t.Any],
+    mgilist: list[int | t.Sequence[int]],
+    modelpath: str | Path,
+    estimators: pl.LazyFrame | pl.DataFrame,
+    dfalldata: pd.DataFrame | None,
+    args: argparse.Namespace,
+    **plotkwargs: t.Any,
 ):
     """Make plot from ARTIS estimators."""
     # these three lists give the x value, modelgridex, and a list of timesteps (for averaging) for each plot of the plot
@@ -649,7 +658,7 @@ def plot_subplot(
                 dfalldata=dfalldata,
                 **plotkwargs,
             )
-            if showlegend and sameylabel:
+            if showlegend and sameylabel and ylabel is not None:
                 ax.set_ylabel(ylabel)
         else:  # it's a sequence of values
             showlegend = True

From 9dc061a70f74c0330638b2976f5997e805356128 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 13:30:36 +0000
Subject: [PATCH 062/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 29 +++++++------------------
 1 file changed, 8 insertions(+), 21 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index c9f36026d..6efe94bb9 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -342,12 +342,10 @@ def get_iontuple(ionstr):
     if missingions:
         print(f" Warning: Can't plot {seriestype} for {missingions} because these ions are not in compositiondata.txt")
 
+    iontuplelist = [iontuple for iontuple in iontuplelist if iontuple not in missingions]
     prev_atomic_number = iontuplelist[0][0]
     colorindex = 0
     for atomic_number, ionstage in iontuplelist:
-        if (atomic_number, ionstage) in missingions:
-            continue
-
         if atomic_number != prev_atomic_number:
             colorindex += 1
 
@@ -372,12 +370,6 @@ def get_iontuple(ionstr):
         ylist = []
         for modelgridindex, timesteps in zip(mgilist, timestepslist):
             if seriestype == "populations":
-                # if (atomic_number, ionstage) not in estimators[timesteps[0], modelgridindex]["populations"]:
-                #     print(
-                #         f"Note: population for {(atomic_number, ionstage)} not in estimators for "
-                #         f"cell {modelgridindex} timesteps {timesteps}"
-                #     )
-
                 if ionstage == "ALL":
                     key = f"nnelement_{elsymbol}"
                 elif hasattr(ionstage, "lower") and ionstage.startswith(at.get_elsymbol(atomic_number)):
@@ -386,18 +378,13 @@ def get_iontuple(ionstr):
                 else:
                     key = f"nnion_{ionstr}"
 
-                try:
-                    estimpop = at.estimators.get_averaged_estimators(
-                        modelpath,
-                        estimators,
-                        timesteps,
-                        modelgridindex,
-                        [key, f"nnelement_{elsymbol}", "nntot"],
-                    )
-                except KeyError:
-                    print(f"KeyError: {key} not in estimators")
-                    ylist.append(float("nan"))
-                    continue
+                estimpop = at.estimators.get_averaged_estimators(
+                    modelpath,
+                    estimators,
+                    timesteps,
+                    modelgridindex,
+                    [key, f"nnelement_{elsymbol}", "nntot"],
+                )
 
                 nionpop = estimpop.get(key, 0.0)
 

From dbe88a3229f4cba0aa52cd8b34c617c54c98852a Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 13:32:58 +0000
Subject: [PATCH 063/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 26 ++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 6efe94bb9..a59d2f9f5 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -353,19 +353,6 @@ def get_iontuple(ionstr):
 
         ionstr = at.get_ionstring(atomic_number, ionstage, sep="_", style="spectral")
 
-        if seriestype == "populations":
-            if args.ionpoptype == "absolute":
-                ax.set_ylabel(r"Number density $\left[\rm{cm}^{-3}\right]$")
-            elif args.ionpoptype == "elpop":
-                # elsym = at.get_elsymbol(atomic_number)
-                ax.set_ylabel(r"X$_{i}$/X$_{\rm element}$")
-            elif args.ionpoptype == "totalpop":
-                ax.set_ylabel(r"X$_{i}$/X$_{rm tot}$")
-            else:
-                raise AssertionError
-        else:
-            ax.set_ylabel(at.estimators.get_dictlabelreplacements().get(seriestype, seriestype))
-
         print(f"Plotting {seriestype} {ionstr}")
         ylist = []
         for modelgridindex, timesteps in zip(mgilist, timestepslist):
@@ -467,6 +454,19 @@ def get_iontuple(ionstr):
         prev_atomic_number = atomic_number
         plotted_something = True
 
+    if seriestype == "populations":
+        if args.ionpoptype == "absolute":
+            ax.set_ylabel(r"Number density $\left[\rm{cm}^{-3}\right]$")
+        elif args.ionpoptype == "elpop":
+            # elsym = at.get_elsymbol(atomic_number)
+            ax.set_ylabel(r"X$_{i}$/X$_{\rm element}$")
+        elif args.ionpoptype == "totalpop":
+            ax.set_ylabel(r"X$_{i}$/X$_{rm tot}$")
+        else:
+            raise AssertionError
+    else:
+        ax.set_ylabel(at.estimators.get_dictlabelreplacements().get(seriestype, seriestype))
+
     if plotted_something:
         ax.set_yscale(args.yscale)
         if args.yscale == "log":

From 0a6a58df427cc72b4842a320997ff253768ee500 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 14:40:54 +0000
Subject: [PATCH 064/150] Use more polars

---
 artistools/inputmodel/maptogrid.py | 86 +++++++++++++-----------------
 1 file changed, 37 insertions(+), 49 deletions(-)

diff --git a/artistools/inputmodel/maptogrid.py b/artistools/inputmodel/maptogrid.py
index beed35416..35216b9dd 100755
--- a/artistools/inputmodel/maptogrid.py
+++ b/artistools/inputmodel/maptogrid.py
@@ -11,6 +11,7 @@
 import argcomplete
 import numpy as np
 import pandas as pd
+import polars as pl
 
 import artistools as at
 
@@ -155,6 +156,8 @@ def logprint(*args, **kwargs):
     if downsamplefactor > 1:
         dfsnapshot = dfsnapshot.sample(len(dfsnapshot) // downsamplefactor)
 
+    dfsnapshot = pl.from_pandas(dfsnapshot)
+
     logprint(dfsnapshot)
 
     assert len(dfsnapshot.columns) == len(snapshot_columns_used)
@@ -172,65 +175,50 @@ def logprint(*args, **kwargs):
     dtextra_seconds = 0.5  # in seconds ---  dtextra = 0.0 # for no extrapolation
 
     dtextra = dtextra_seconds / 4.926e-6  # convert to geom units.
+    dfsnapshot = dfsnapshot.with_columns(dis_orig=(pl.col("x") ** 2 + pl.col("y") ** 2 + pl.col("z") ** 2).sqrt())
+
+    dfsnapshot = dfsnapshot.with_columns(
+        x=pl.col("x") + pl.col("vx") * dtextra,
+        y=pl.col("y") + pl.col("vy") * dtextra,
+        z=pl.col("z") + pl.col("vz") * dtextra,
+    )
+
+    dfsnapshot = dfsnapshot.with_columns(dis=(pl.col("x") ** 2 + pl.col("y") ** 2 + pl.col("z") ** 2).sqrt())
+
+    dfsnapshot = dfsnapshot.with_columns(
+        h=pl.col("h") / pl.col("dis_orig") * pl.col("dis"),
+        vrad=(pl.col("vx") * pl.col("x") + pl.col("vy") * pl.col("y") + pl.col("vz") * pl.col("z")) / pl.col("dis"),
+        vtot=(pl.col("vx") ** 2 + pl.col("vy") ** 2 + pl.col("vz") ** 2).sqrt(),
+    )
+    dfsnapshot = dfsnapshot.with_columns(
+        vperp=pl.when(pl.col("vtot") > pl.col("vrad"))
+        .then((pl.col("vtot") ** 2 - pl.col("vrad") ** 2).sqrt())
+        .otherwise(0.0),
+    )
 
-    particleid = dfsnapshot.id.to_numpy()
-    x = dfsnapshot["x"].to_numpy().copy()
-    y = dfsnapshot["y"].to_numpy().copy()
-    z = dfsnapshot["z"].to_numpy().copy()
+    particleid = dfsnapshot["id"].to_numpy()
+    x = dfsnapshot["x"].to_numpy()
+    y = dfsnapshot["y"].to_numpy()
+    z = dfsnapshot["z"].to_numpy()
     h = dfsnapshot["h"].to_numpy().copy()
-    vx = dfsnapshot["vx"].to_numpy()
-    vy = dfsnapshot["vy"].to_numpy()
-    vz = dfsnapshot["vz"].to_numpy()
     pmass = dfsnapshot["pmass"].to_numpy()
     rho_rst = dfsnapshot["rho_rst"].to_numpy()
     rho = dfsnapshot["rho"].to_numpy()
     Ye = dfsnapshot["ye"].to_numpy()
 
+    totmass = dfsnapshot["pmass"].sum()
+    rmean = dfsnapshot["dis"].mean()
+    hmean = dfsnapshot["h"].mean()
+    vratiomean = dfsnapshot.select(pl.col("vperp") - pl.col("vrad")).mean()
+    hmin = dfsnapshot["h"].min()
+    rmax = dfsnapshot["dis"].max()
     with Path(outputfolderpath, "ejectapartanalysis.dat").open(mode="w", encoding="utf-8") as fpartanalysis:
-        for n in range(npart):
-            totmass = totmass + pmass[n]
-
-            dis = math.sqrt(x[n] ** 2 + y[n] ** 2 + z[n] ** 2)  # original dis
-
-            x[n] += vx[n] * dtextra
-            y[n] += vy[n] * dtextra
-            z[n] += vz[n] * dtextra
-
-            # actually we should also extrapolate smoothing length h unless we disrgard it below
-
-            # extrapolate h such that ratio between dis and h remains constant
-            h[n] = h[n] / dis * math.sqrt(x[n] ** 2 + y[n] ** 2 + z[n] ** 2)
-
-            dis = math.sqrt(x[n] ** 2 + y[n] ** 2 + z[n] ** 2)  # possibly new distance
-
-            rmean = rmean + dis
-
-            rmax = max(rmax, dis)
-
-            hmean = hmean + h[n]
-
-            hmin = min(hmin, h[n])
-
-            vtot = math.sqrt(vx[n] ** 2 + vy[n] ** 2 + vz[n] ** 2)
-
-            vrad = (vx[n] * x[n] + vy[n] * y[n] + vz[n] * z[n]) / dis  # radial velocity
-
-            # velocity perpendicular
-            # if we extrapolate roundoff error can lead to Nan, ly?
-            vperp = math.sqrt(vtot * vtot - vrad * vrad) if vtot > vrad else 0.0
-
-            vratiomean = vratiomean + vperp / vrad
-
-            # output some ejecta properties in file
-
-            fpartanalysis.write(f"{dis} {h[n]} {h[n] / dis} {vrad} {vperp} {vtot}\n")
+        for part in dfsnapshot.select(["dis", "h", "vrad", "vperp", "vtot"]).iter_rows(named=True):
+            fpartanalysis.write(
+                f"{part['dis']} {part['h']} {part['h'] / part['dis']} {part['vrad']} {part['vperp']} {part['vtot']}\n"
+            )
 
     logprint(f"saved {outputfolderpath / 'ejectapartanalysis.dat'}")
-    rmean = rmean / npart
-
-    hmean = hmean / npart
-
-    vratiomean = vratiomean / npart
 
     logprint(f"total mass of sph particle {totmass} max dist {rmax} mean dist {rmean}")
     logprint(f"smoothing length min {hmin} mean {hmean}")

From 99497cc1df28bc7b703179c5f580492ff260bbb5 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 14:43:31 +0000
Subject: [PATCH 065/150] Update maptogrid.py

---
 artistools/inputmodel/maptogrid.py | 13 ++-----------
 1 file changed, 2 insertions(+), 11 deletions(-)

diff --git a/artistools/inputmodel/maptogrid.py b/artistools/inputmodel/maptogrid.py
index 35216b9dd..aed3b414d 100755
--- a/artistools/inputmodel/maptogrid.py
+++ b/artistools/inputmodel/maptogrid.py
@@ -164,13 +164,6 @@ def logprint(*args, **kwargs):
 
     npart = len(dfsnapshot)
 
-    totmass = 0.0
-    rmax = 0.0
-    rmean = 0.0
-    hmean = 0.0
-    hmin = 100000.0
-    vratiomean = 0.0
-
     # Propagate particles to dtextra using velocities
     dtextra_seconds = 0.5  # in seconds ---  dtextra = 0.0 # for no extrapolation
 
@@ -209,8 +202,6 @@ def logprint(*args, **kwargs):
     totmass = dfsnapshot["pmass"].sum()
     rmean = dfsnapshot["dis"].mean()
     hmean = dfsnapshot["h"].mean()
-    vratiomean = dfsnapshot.select(pl.col("vperp") - pl.col("vrad")).mean()
-    hmin = dfsnapshot["h"].min()
     rmax = dfsnapshot["dis"].max()
     with Path(outputfolderpath, "ejectapartanalysis.dat").open(mode="w", encoding="utf-8") as fpartanalysis:
         for part in dfsnapshot.select(["dis", "h", "vrad", "vperp", "vtot"]).iter_rows(named=True):
@@ -221,8 +212,8 @@ def logprint(*args, **kwargs):
     logprint(f"saved {outputfolderpath / 'ejectapartanalysis.dat'}")
 
     logprint(f"total mass of sph particle {totmass} max dist {rmax} mean dist {rmean}")
-    logprint(f"smoothing length min {hmin} mean {hmean}")
-    logprint("ratio between vrad and vperp mean", vratiomean)
+    logprint(f"smoothing length min {dfsnapshot['h'].min()} mean {hmean}")
+    logprint("ratio between vrad and vperp mean", dfsnapshot.select(pl.col("vperp") - pl.col("vrad")).mean().item(0, 0))
 
     # check maybe cm and correct by shifting
 

From 341e4fa10367691d917393b0a971441417f057cd Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 14:53:48 +0000
Subject: [PATCH 066/150] Update solvespencerfanocmd.py

---
 artistools/nonthermal/solvespencerfanocmd.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/artistools/nonthermal/solvespencerfanocmd.py b/artistools/nonthermal/solvespencerfanocmd.py
index 5a358e3ab..1f5a44f2d 100755
--- a/artistools/nonthermal/solvespencerfanocmd.py
+++ b/artistools/nonthermal/solvespencerfanocmd.py
@@ -160,6 +160,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
     if Path(args.outputfile).is_dir():
         args.outputfile = Path(args.outputfile, defaultoutputfile)
 
+    ionpopdict: dict[tuple[int, int] | int, float]
     if args.composition == "artis":
         if args.timedays:
             args.timestep = at.get_timestep_of_timedays(modelpath, args.timedays)

From 43ed7eb88386677b0dc91a22290a921b83f68de8 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 14:56:11 +0000
Subject: [PATCH 067/150] Update mypy to 1.7.0

---
 artistools/estimators/plotestimators.py      | 2 +-
 artistools/nonthermal/solvespencerfanocmd.py | 6 +-----
 requirements.txt                             | 2 +-
 3 files changed, 3 insertions(+), 7 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index a59d2f9f5..cf5f628d8 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -988,7 +988,7 @@ def addargs(parser: argparse.ArgumentParser) -> None:
 
     parser.add_argument("--notitle", action="store_true", help="Suppress the top title from the plot")
 
-    parser.add_argument("-plotlist", type=list, default=[], help="Plot list (when calling from Python only)")  # type: ignore[arg-type]
+    parser.add_argument("-plotlist", type=list, default=[], help="Plot list (when calling from Python only)")
 
     parser.add_argument(
         "-ionpoptype",
diff --git a/artistools/nonthermal/solvespencerfanocmd.py b/artistools/nonthermal/solvespencerfanocmd.py
index 1f5a44f2d..1a76b341e 100755
--- a/artistools/nonthermal/solvespencerfanocmd.py
+++ b/artistools/nonthermal/solvespencerfanocmd.py
@@ -306,11 +306,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             ionpopdict[(compelement_atomicnumber, 2)] = nntot * x_e
 
         # keep only the ion populations, not element or total populations
-        ions = [
-            key
-            for key in ionpopdict
-            if isinstance(key, tuple) and len(key) == 2 and ionpopdict[key] / nntot >= minionfraction
-        ]
+        ions = [key for key in ionpopdict if isinstance(key, tuple) and ionpopdict[key] / nntot >= minionfraction]
         ions.sort()
 
         if args.noexcitation:
diff --git a/requirements.txt b/requirements.txt
index 0f2c5d68e..db6f4b4bf 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,7 +4,7 @@ coverage>=7.3.2
 extinction>=0.4.6
 imageio>=2.32.0
 matplotlib>=3.8.1
-mypy>=1.6.1
+mypy>=1.7.0
 numpy>=1.26.1
 pandas>=2.1.2
 polars>=0.19.12

From 646feea48484bb42ff52fedbda6456d1188b346c Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 14:58:16 +0000
Subject: [PATCH 068/150] Update maptogrid.py

---
 artistools/inputmodel/maptogrid.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/artistools/inputmodel/maptogrid.py b/artistools/inputmodel/maptogrid.py
index aed3b414d..ce2bd0a83 100755
--- a/artistools/inputmodel/maptogrid.py
+++ b/artistools/inputmodel/maptogrid.py
@@ -171,6 +171,9 @@ def logprint(*args, **kwargs):
     dfsnapshot = dfsnapshot.with_columns(dis_orig=(pl.col("x") ** 2 + pl.col("y") ** 2 + pl.col("z") ** 2).sqrt())
 
     dfsnapshot = dfsnapshot.with_columns(
+        x_orig=pl.col("x"),
+        y_orig=pl.col("y"),
+        z_orig=pl.col("z"),
         x=pl.col("x") + pl.col("vx") * dtextra,
         y=pl.col("y") + pl.col("vy") * dtextra,
         z=pl.col("z") + pl.col("vz") * dtextra,

From bdc6a4c7722da3478e16a09bcbec2da71d6166d2 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 16:04:27 +0000
Subject: [PATCH 069/150] Update estimators.py

---
 artistools/estimators/estimators.py | 13 ++-----------
 1 file changed, 2 insertions(+), 11 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 31cf3ad31..6ac486658 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -120,7 +120,6 @@ def get_units_string(variable: str) -> str:
 def read_estimators_from_file(
     estfilepath: Path | str,
     printfilename: bool = False,
-    skip_emptycells: bool = True,
 ) -> pl.DataFrame:
     if printfilename:
         estfilepath = Path(estfilepath)
@@ -139,11 +138,7 @@ def read_estimators_from_file(
 
             if row[0] == "timestep":
                 # yield the previous block before starting a new one
-                if (
-                    timestep is not None
-                    and modelgridindex is not None
-                    and (not skip_emptycells or not estimblock.get("emptycell", True))
-                ):
+                if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
                     estimblock["timestep"] = timestep
                     estimblock["modelgridindex"] = modelgridindex
                     estimblocklist.append(estimblock)
@@ -227,11 +222,7 @@ def read_estimators_from_file(
                     estimblock[f"cooling_{coolingtype}"] = float(value)
 
     # reached the end of file
-    if (
-        timestep is not None
-        and modelgridindex is not None
-        and (not skip_emptycells or not estimblock.get("emptycell", True))
-    ):
+    if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
         estimblock["timestep"] = timestep
         estimblock["modelgridindex"] = modelgridindex
         estimblocklist.append(estimblock)

From 1d19211e8620e316e6547090b4dbc78b3f83d34a Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 16:09:01 +0000
Subject: [PATCH 070/150] Update estimators.py

---
 artistools/estimators/estimators.py | 27 ++++++++++++---------------
 1 file changed, 12 insertions(+), 15 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 6ac486658..9952f6f24 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -128,8 +128,6 @@ def read_estimators_from_file(
 
     estimblocklist: list[dict[str, t.Any]] = []
     with at.zopen(estfilepath) as estimfile:
-        timestep: int | None = None
-        modelgridindex: int | None = None
         estimblock: dict[str, t.Any] = {}
         for line in estimfile:
             row: list[str] = line.split()
@@ -138,19 +136,16 @@ def read_estimators_from_file(
 
             if row[0] == "timestep":
                 # yield the previous block before starting a new one
-                if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
-                    estimblock["timestep"] = timestep
-                    estimblock["modelgridindex"] = modelgridindex
+                if (
+                    estimblock.get("timestep") is not None
+                    and estimblock.get("modelgridindex") is not None
+                    and (not estimblock.get("emptycell", True))
+                ):
                     estimblocklist.append(estimblock)
 
-                timestep = int(row[1])
-                # if timestep > itstep:
-                #     print(f"Dropping estimator data from timestep {timestep} and later (> itstep {itstep})")
-                #     # itstep in input.txt is updated by ARTIS at every timestep, so the data beyond here
-                #     # could be half-written to disk and cause parsing errors
-                #     return
+                estimblock["timestep"] = int(row[1])
 
-                modelgridindex = int(row[3])
+                estimblock["modelgridindex"] = int(row[3])
                 emptycell = row[4] == "EMPTYCELL"
                 estimblock = {"emptycell": emptycell}
                 if not emptycell:
@@ -222,9 +217,11 @@ def read_estimators_from_file(
                     estimblock[f"cooling_{coolingtype}"] = float(value)
 
     # reached the end of file
-    if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
-        estimblock["timestep"] = timestep
-        estimblock["modelgridindex"] = modelgridindex
+    if (
+        estimblock.get("timestep") is not None
+        and estimblock.get("modelgridindex") is not None
+        and not estimblock.get("emptycell", True)
+    ):
         estimblocklist.append(estimblock)
 
     return pl.DataFrame(estimblocklist).with_columns(

From 0133bb0884cbd794459c42eb5c0127e7424d8cdc Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 16:11:04 +0000
Subject: [PATCH 071/150] Update estimators.py

---
 artistools/estimators/estimators.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 9952f6f24..147050a92 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -11,7 +11,6 @@
 import sys
 import typing as t
 from collections import namedtuple
-from functools import partial
 from pathlib import Path
 
 import numpy as np
@@ -127,8 +126,8 @@ def read_estimators_from_file(
         print(f"  Reading {estfilepath.relative_to(estfilepath.parent.parent)} ({filesize:.2f} MiB)")
 
     estimblocklist: list[dict[str, t.Any]] = []
+    estimblock: dict[str, t.Any] = {}
     with at.zopen(estfilepath) as estimfile:
-        estimblock: dict[str, t.Any] = {}
         for line in estimfile:
             row: list[str] = line.split()
             if not row:
@@ -265,11 +264,9 @@ def get_rankbatch_parquetfile(
             f"  reading {len(list(estfilepaths))} estimator files from {folderpath.relative_to(Path(folderpath).parent)}"
         )
 
-        processfile = partial(read_estimators_from_file)
-
         pldf_group = None
         with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
-            for pldf_file in pool.imap(processfile, estfilepaths):
+            for pldf_file in pool.imap(read_estimators_from_file, estfilepaths):
                 if pldf_group is None:
                     pldf_group = pldf_file
                 else:

From 2f2b0afada6f00d353d743041b064dd45745167e Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 16:19:39 +0000
Subject: [PATCH 072/150] Update estimators.py

---
 artistools/estimators/estimators.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 147050a92..ab948daff 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -247,14 +247,14 @@ def batched(iterable, n):
 def get_rankbatch_parquetfile(
     modelpath: Path,
     folderpath: Path,
-    mpiranks: t.Sequence[int],
+    batch_mpiranks: t.Sequence[int],
 ) -> Path:
-    parquetfilepath = folderpath / f"estimators_{mpiranks[0]:04d}_{mpiranks[-1]:04d}.out.parquet.tmp"
+    parquetfilepath = folderpath / f"estimators_{batch_mpiranks[0]:04d}_{batch_mpiranks[-1]:04d}.out.parquet.tmp"
 
     if not parquetfilepath.exists():
         print(f"{parquetfilepath.relative_to(modelpath.parent)} does not exist")
         estfilepaths = []
-        for mpirank in mpiranks:
+        for mpirank in batch_mpiranks:
             # not worth printing an error, because ranks with no cells to update do not produce an estimator file
             with contextlib.suppress(FileNotFoundError):
                 estfilepath = at.firstexisting(f"estimators_{mpirank:04d}.out", folder=folderpath, tryzipped=True)

From b2294b3ae0d01b1de0c2e15daf0eb0f15846f660 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 16:23:15 +0000
Subject: [PATCH 073/150] Update estimators.py

---
 artistools/estimators/estimators.py | 23 ++++++++++-------------
 1 file changed, 10 insertions(+), 13 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index ab948daff..e44381c2d 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -127,6 +127,8 @@ def read_estimators_from_file(
 
     estimblocklist: list[dict[str, t.Any]] = []
     estimblock: dict[str, t.Any] = {}
+    timestep: int | None = None
+    modelgridindex: int | None = None
     with at.zopen(estfilepath) as estimfile:
         for line in estimfile:
             row: list[str] = line.split()
@@ -135,16 +137,13 @@ def read_estimators_from_file(
 
             if row[0] == "timestep":
                 # yield the previous block before starting a new one
-                if (
-                    estimblock.get("timestep") is not None
-                    and estimblock.get("modelgridindex") is not None
-                    and (not estimblock.get("emptycell", True))
-                ):
+                if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
+                    estimblock["timestep"] = timestep
+                    estimblock["modelgridindex"] = modelgridindex
                     estimblocklist.append(estimblock)
 
-                estimblock["timestep"] = int(row[1])
-
-                estimblock["modelgridindex"] = int(row[3])
+                timestep = int(row[1])
+                modelgridindex = int(row[3])
                 emptycell = row[4] == "EMPTYCELL"
                 estimblock = {"emptycell": emptycell}
                 if not emptycell:
@@ -216,11 +215,9 @@ def read_estimators_from_file(
                     estimblock[f"cooling_{coolingtype}"] = float(value)
 
     # reached the end of file
-    if (
-        estimblock.get("timestep") is not None
-        and estimblock.get("modelgridindex") is not None
-        and not estimblock.get("emptycell", True)
-    ):
+    if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
+        estimblock["timestep"] = timestep
+        estimblock["modelgridindex"] = modelgridindex
         estimblocklist.append(estimblock)
 
     return pl.DataFrame(estimblocklist).with_columns(

From f4a79bd494c90222882e33c7bec835d2403f429f Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 16:23:57 +0000
Subject: [PATCH 074/150] Update estimators.py

---
 artistools/estimators/estimators.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index e44381c2d..f618e7feb 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -138,9 +138,7 @@ def read_estimators_from_file(
             if row[0] == "timestep":
                 # yield the previous block before starting a new one
                 if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
-                    estimblock["timestep"] = timestep
-                    estimblock["modelgridindex"] = modelgridindex
-                    estimblocklist.append(estimblock)
+                    estimblocklist.append(estimblock | {"timestep": timestep, "modelgridindex": modelgridindex})
 
                 timestep = int(row[1])
                 modelgridindex = int(row[3])
@@ -216,9 +214,7 @@ def read_estimators_from_file(
 
     # reached the end of file
     if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
-        estimblock["timestep"] = timestep
-        estimblock["modelgridindex"] = modelgridindex
-        estimblocklist.append(estimblock)
+        estimblocklist.append(estimblock | {"timestep": timestep, "modelgridindex": modelgridindex})
 
     return pl.DataFrame(estimblocklist).with_columns(
         pl.col(pl.Int64).cast(pl.Int32), pl.col(pl.Float64).cast(pl.Float32)

From aaecc31a0f018ce73b53c3ab2bb71cf340a3de83 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 16:28:04 +0000
Subject: [PATCH 075/150] Update polars to 0.19.13 and pass filename to
 read_csv

---
 artistools/inputmodel/inputmodel_misc.py |  8 ++------
 artistools/lightcurve/lightcurve.py      | 12 ++----------
 artistools/misc.py                       | 13 ++-----------
 artistools/packets/packets.py            |  6 +-----
 artistools/spectra/spectra.py            | 17 +++++------------
 requirements.txt                         |  2 +-
 6 files changed, 13 insertions(+), 45 deletions(-)

diff --git a/artistools/inputmodel/inputmodel_misc.py b/artistools/inputmodel/inputmodel_misc.py
index 49ff0898e..d957c842d 100644
--- a/artistools/inputmodel/inputmodel_misc.py
+++ b/artistools/inputmodel/inputmodel_misc.py
@@ -133,7 +133,7 @@ def read_modelfile_text(
         pldtypes = {col: pl.Int32 if col in {"inputcellid"} else pl.Float32 for col in columns}
 
         dfmodel = pl.read_csv(
-            at.zopen(filename, forpolars=True),
+            filename,
             separator=" ",
             comment_char="#",
             new_columns=columns,
@@ -1006,11 +1006,7 @@ def get_initelemabundances_polars(
         dtypes = {col: pl.Float32 if col.startswith("X_") else pl.Int32 for col in colnames}
 
         abundancedata = pl.read_csv(
-            at.zopen(abundancefilepath, forpolars=True),
-            has_header=False,
-            separator=" ",
-            comment_char="#",
-            infer_schema_length=0,
+            abundancefilepath, has_header=False, separator=" ", comment_char="#", infer_schema_length=0
         )
 
         # fix up multiple spaces at beginning of lines
diff --git a/artistools/lightcurve/lightcurve.py b/artistools/lightcurve/lightcurve.py
index 4bd688e0b..cb577a8fb 100644
--- a/artistools/lightcurve/lightcurve.py
+++ b/artistools/lightcurve/lightcurve.py
@@ -21,19 +21,11 @@ def readfile(
     lcdata: dict[int, pl.DataFrame] = {}
     if "_res" in str(filepath):
         # get a dict of dfs with light curves at each viewing direction bin
-        lcdata_res = pl.read_csv(
-            at.zopen(filepath, forpolars=True),
-            separator=" ",
-            has_header=False,
-            new_columns=["time", "lum", "lum_cmf"],
-        )
+        lcdata_res = pl.read_csv(filepath, separator=" ", has_header=False, new_columns=["time", "lum", "lum_cmf"])
         lcdata = at.split_dataframe_dirbins(lcdata_res, index_of_repeated_value=0, output_polarsdf=True)
     else:
         dfsphericalaverage = pl.read_csv(
-            at.zopen(filepath, forpolars=True),
-            separator=" ",
-            has_header=False,
-            new_columns=["time", "lum", "lum_cmf"],
+            filepath, separator=" ", has_header=False, new_columns=["time", "lum", "lum_cmf"]
         )
 
         if list(dfsphericalaverage["time"].to_numpy()) != sorted(dfsphericalaverage["time"].to_numpy()):
diff --git a/artistools/misc.py b/artistools/misc.py
index a19597765..fa7e27aa0 100644
--- a/artistools/misc.py
+++ b/artistools/misc.py
@@ -767,15 +767,8 @@ def flatten_list(listin: list[t.Any]) -> list[t.Any]:
     return listout
 
 
-def zopen(filename: Path | str, mode: str = "rt", encoding: str | None = None, forpolars: bool = False) -> t.Any:
-    """Open filename, filename.gz or filename.xz.
-
-    Arguments:
-    ---------
-    forpolars: if polars.read_csv can read the file directly, return a Path object instead of a file object
-    """
-    if forpolars:
-        mode = "r"
+def zopen(filename: Path | str, mode: str = "rt", encoding: str | None = None) -> t.Any:
+    """Open filename, filename.gz or filename.xz."""
     ext_fopen: dict[str, t.Callable] = {".zst": pyzstd.open, ".gz": gzip.open, ".xz": xz.open}
 
     for ext, fopen in ext_fopen.items():
@@ -784,8 +777,6 @@ def zopen(filename: Path | str, mode: str = "rt", encoding: str | None = None, f
             return fopen(file_ext, mode=mode, encoding=encoding)
 
     # open() can raise file not found if this file doesn't exist
-    if forpolars:
-        return Path(filename)
     return Path(filename).open(mode=mode, encoding=encoding)  # noqa: SIM115
 
 
diff --git a/artistools/packets/packets.py b/artistools/packets/packets.py
index e2affe451..8df916cb1 100644
--- a/artistools/packets/packets.py
+++ b/artistools/packets/packets.py
@@ -321,11 +321,7 @@ def readfile_text(packetsfile: Path | str, modelpath: Path = Path()) -> pl.DataF
 
     try:
         dfpackets = pl.read_csv(
-            fpackets,
-            separator=" ",
-            has_header=False,
-            new_columns=column_names,
-            infer_schema_length=20000,
+            fpackets, separator=" ", has_header=False, new_columns=column_names, infer_schema_length=20000
         )
 
     except Exception:
diff --git a/artistools/spectra/spectra.py b/artistools/spectra/spectra.py
index 8a8925352..d7175a518 100644
--- a/artistools/spectra/spectra.py
+++ b/artistools/spectra/spectra.py
@@ -255,12 +255,7 @@ def read_spec(modelpath: Path) -> pl.DataFrame:
     print(f"Reading {specfilename}")
 
     return (
-        pl.read_csv(
-            at.zopen(specfilename, forpolars=True),
-            separator=" ",
-            infer_schema_length=0,
-            truncate_ragged_lines=True,
-        )
+        pl.read_csv(specfilename, separator=" ", infer_schema_length=0, truncate_ragged_lines=True)
         .with_columns(pl.all().cast(pl.Float64))
         .rename({"0": "nu"})
     )
@@ -276,9 +271,7 @@ def read_spec_res(modelpath: Path) -> dict[int, pl.DataFrame]:
     )
 
     print(f"Reading {specfilename} (in read_spec_res)")
-    res_specdata_in = pl.read_csv(
-        at.zopen(specfilename, forpolars=True), separator=" ", has_header=False, infer_schema_length=0
-    )
+    res_specdata_in = pl.read_csv(specfilename, separator=" ", has_header=False, infer_schema_length=0)
 
     # drop last column of nulls (caused by trailing space on each line)
     if res_specdata_in[res_specdata_in.columns[-1]].is_null().all():
@@ -325,9 +318,9 @@ def read_emission_absorption_file(emabsfilename: str | Path) -> pl.DataFrame:
     except AttributeError:
         print(f" Reading {emabsfilename}")
 
-    dfemabs = pl.read_csv(
-        at.zopen(emabsfilename, forpolars=True), separator=" ", has_header=False, infer_schema_length=0
-    ).with_columns(pl.all().cast(pl.Float32, strict=False))
+    dfemabs = pl.read_csv(emabsfilename, separator=" ", has_header=False, infer_schema_length=0).with_columns(
+        pl.all().cast(pl.Float32, strict=False)
+    )
 
     # drop last column of nulls (caused by trailing space on each line)
     if dfemabs[dfemabs.columns[-1]].is_null().all():
diff --git a/requirements.txt b/requirements.txt
index db6f4b4bf..11513a4cb 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,7 +7,7 @@ matplotlib>=3.8.1
 mypy>=1.7.0
 numpy>=1.26.1
 pandas>=2.1.2
-polars>=0.19.12
+polars>=0.19.13
 pre-commit>=3.5.0
 pyarrow>=14.0.1
 pynonthermal>=2021.10.12

From 2e882c6e3660a1ec8faed31f58d92a3d72d4937b Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 16:36:25 +0000
Subject: [PATCH 076/150] Revert "Update polars to 0.19.13 and pass filename to
 read_csv"

This reverts commit aaecc31a0f018ce73b53c3ab2bb71cf340a3de83.
---
 artistools/inputmodel/inputmodel_misc.py |  8 ++++++--
 artistools/lightcurve/lightcurve.py      | 12 ++++++++++--
 artistools/misc.py                       | 13 +++++++++++--
 artistools/packets/packets.py            |  6 +++++-
 artistools/spectra/spectra.py            | 17 ++++++++++++-----
 requirements.txt                         |  2 +-
 6 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/artistools/inputmodel/inputmodel_misc.py b/artistools/inputmodel/inputmodel_misc.py
index d957c842d..49ff0898e 100644
--- a/artistools/inputmodel/inputmodel_misc.py
+++ b/artistools/inputmodel/inputmodel_misc.py
@@ -133,7 +133,7 @@ def read_modelfile_text(
         pldtypes = {col: pl.Int32 if col in {"inputcellid"} else pl.Float32 for col in columns}
 
         dfmodel = pl.read_csv(
-            filename,
+            at.zopen(filename, forpolars=True),
             separator=" ",
             comment_char="#",
             new_columns=columns,
@@ -1006,7 +1006,11 @@ def get_initelemabundances_polars(
         dtypes = {col: pl.Float32 if col.startswith("X_") else pl.Int32 for col in colnames}
 
         abundancedata = pl.read_csv(
-            abundancefilepath, has_header=False, separator=" ", comment_char="#", infer_schema_length=0
+            at.zopen(abundancefilepath, forpolars=True),
+            has_header=False,
+            separator=" ",
+            comment_char="#",
+            infer_schema_length=0,
         )
 
         # fix up multiple spaces at beginning of lines
diff --git a/artistools/lightcurve/lightcurve.py b/artistools/lightcurve/lightcurve.py
index cb577a8fb..4bd688e0b 100644
--- a/artistools/lightcurve/lightcurve.py
+++ b/artistools/lightcurve/lightcurve.py
@@ -21,11 +21,19 @@ def readfile(
     lcdata: dict[int, pl.DataFrame] = {}
     if "_res" in str(filepath):
         # get a dict of dfs with light curves at each viewing direction bin
-        lcdata_res = pl.read_csv(filepath, separator=" ", has_header=False, new_columns=["time", "lum", "lum_cmf"])
+        lcdata_res = pl.read_csv(
+            at.zopen(filepath, forpolars=True),
+            separator=" ",
+            has_header=False,
+            new_columns=["time", "lum", "lum_cmf"],
+        )
         lcdata = at.split_dataframe_dirbins(lcdata_res, index_of_repeated_value=0, output_polarsdf=True)
     else:
         dfsphericalaverage = pl.read_csv(
-            filepath, separator=" ", has_header=False, new_columns=["time", "lum", "lum_cmf"]
+            at.zopen(filepath, forpolars=True),
+            separator=" ",
+            has_header=False,
+            new_columns=["time", "lum", "lum_cmf"],
         )
 
         if list(dfsphericalaverage["time"].to_numpy()) != sorted(dfsphericalaverage["time"].to_numpy()):
diff --git a/artistools/misc.py b/artistools/misc.py
index fa7e27aa0..a19597765 100644
--- a/artistools/misc.py
+++ b/artistools/misc.py
@@ -767,8 +767,15 @@ def flatten_list(listin: list[t.Any]) -> list[t.Any]:
     return listout
 
 
-def zopen(filename: Path | str, mode: str = "rt", encoding: str | None = None) -> t.Any:
-    """Open filename, filename.gz or filename.xz."""
+def zopen(filename: Path | str, mode: str = "rt", encoding: str | None = None, forpolars: bool = False) -> t.Any:
+    """Open filename, filename.gz or filename.xz.
+
+    Arguments:
+    ---------
+    forpolars: if polars.read_csv can read the file directly, return a Path object instead of a file object
+    """
+    if forpolars:
+        mode = "r"
     ext_fopen: dict[str, t.Callable] = {".zst": pyzstd.open, ".gz": gzip.open, ".xz": xz.open}
 
     for ext, fopen in ext_fopen.items():
@@ -777,6 +784,8 @@ def zopen(filename: Path | str, mode: str = "rt", encoding: str | None = None) -
             return fopen(file_ext, mode=mode, encoding=encoding)
 
     # open() can raise file not found if this file doesn't exist
+    if forpolars:
+        return Path(filename)
     return Path(filename).open(mode=mode, encoding=encoding)  # noqa: SIM115
 
 
diff --git a/artistools/packets/packets.py b/artistools/packets/packets.py
index 8df916cb1..e2affe451 100644
--- a/artistools/packets/packets.py
+++ b/artistools/packets/packets.py
@@ -321,7 +321,11 @@ def readfile_text(packetsfile: Path | str, modelpath: Path = Path()) -> pl.DataF
 
     try:
         dfpackets = pl.read_csv(
-            fpackets, separator=" ", has_header=False, new_columns=column_names, infer_schema_length=20000
+            fpackets,
+            separator=" ",
+            has_header=False,
+            new_columns=column_names,
+            infer_schema_length=20000,
         )
 
     except Exception:
diff --git a/artistools/spectra/spectra.py b/artistools/spectra/spectra.py
index d7175a518..8a8925352 100644
--- a/artistools/spectra/spectra.py
+++ b/artistools/spectra/spectra.py
@@ -255,7 +255,12 @@ def read_spec(modelpath: Path) -> pl.DataFrame:
     print(f"Reading {specfilename}")
 
     return (
-        pl.read_csv(specfilename, separator=" ", infer_schema_length=0, truncate_ragged_lines=True)
+        pl.read_csv(
+            at.zopen(specfilename, forpolars=True),
+            separator=" ",
+            infer_schema_length=0,
+            truncate_ragged_lines=True,
+        )
         .with_columns(pl.all().cast(pl.Float64))
         .rename({"0": "nu"})
     )
@@ -271,7 +276,9 @@ def read_spec_res(modelpath: Path) -> dict[int, pl.DataFrame]:
     )
 
     print(f"Reading {specfilename} (in read_spec_res)")
-    res_specdata_in = pl.read_csv(specfilename, separator=" ", has_header=False, infer_schema_length=0)
+    res_specdata_in = pl.read_csv(
+        at.zopen(specfilename, forpolars=True), separator=" ", has_header=False, infer_schema_length=0
+    )
 
     # drop last column of nulls (caused by trailing space on each line)
     if res_specdata_in[res_specdata_in.columns[-1]].is_null().all():
@@ -318,9 +325,9 @@ def read_emission_absorption_file(emabsfilename: str | Path) -> pl.DataFrame:
     except AttributeError:
         print(f" Reading {emabsfilename}")
 
-    dfemabs = pl.read_csv(emabsfilename, separator=" ", has_header=False, infer_schema_length=0).with_columns(
-        pl.all().cast(pl.Float32, strict=False)
-    )
+    dfemabs = pl.read_csv(
+        at.zopen(emabsfilename, forpolars=True), separator=" ", has_header=False, infer_schema_length=0
+    ).with_columns(pl.all().cast(pl.Float32, strict=False))
 
     # drop last column of nulls (caused by trailing space on each line)
     if dfemabs[dfemabs.columns[-1]].is_null().all():
diff --git a/requirements.txt b/requirements.txt
index 11513a4cb..db6f4b4bf 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,7 +7,7 @@ matplotlib>=3.8.1
 mypy>=1.7.0
 numpy>=1.26.1
 pandas>=2.1.2
-polars>=0.19.13
+polars>=0.19.12
 pre-commit>=3.5.0
 pyarrow>=14.0.1
 pynonthermal>=2021.10.12

From c6635bd99cb8a841693e52fc2d30c68015b0a795 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 16:37:47 +0000
Subject: [PATCH 077/150] Keep zopen for .xz files

---
 artistools/misc.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/artistools/misc.py b/artistools/misc.py
index a19597765..250f615e9 100644
--- a/artistools/misc.py
+++ b/artistools/misc.py
@@ -779,13 +779,16 @@ def zopen(filename: Path | str, mode: str = "rt", encoding: str | None = None, f
     ext_fopen: dict[str, t.Callable] = {".zst": pyzstd.open, ".gz": gzip.open, ".xz": xz.open}
 
     for ext, fopen in ext_fopen.items():
-        file_ext = str(filename) if str(filename).endswith(ext) else str(filename) + ext
-        if Path(file_ext).exists():
-            return fopen(file_ext, mode=mode, encoding=encoding)
+        file_withext = str(filename) if str(filename).endswith(ext) else str(filename) + ext
+        if Path(file_withext).exists():
+            if forpolars and ext in {".gz", ".zst"}:
+                return Path(file_withext)
+            return fopen(file_withext, mode=mode, encoding=encoding)
 
-    # open() can raise file not found if this file doesn't exist
     if forpolars:
         return Path(filename)
+
+    # open() can raise file not found if this file doesn't exist
     return Path(filename).open(mode=mode, encoding=encoding)  # noqa: SIM115
 
 

From 0192227fcb1cb5b802d8cabe02b39b0f0c0a0f37 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 21:06:01 +0000
Subject: [PATCH 078/150] Fix pylint error

---
 artistools/__init__.py                   |  1 +
 artistools/inputmodel/inputmodel_misc.py |  4 ++--
 artistools/lightcurve/lightcurve.py      | 10 ++------
 artistools/misc.py                       | 30 +++++++++++++-----------
 artistools/spectra/spectra.py            | 13 +++-------
 pyproject.toml                           |  1 -
 6 files changed, 24 insertions(+), 35 deletions(-)

diff --git a/artistools/__init__.py b/artistools/__init__.py
index f9ff09b45..bd9e9f8bb 100644
--- a/artistools/__init__.py
+++ b/artistools/__init__.py
@@ -101,6 +101,7 @@
 from artistools.misc import trim_or_pad
 from artistools.misc import vec_len
 from artistools.misc import zopen
+from artistools.misc import zopenpl
 from artistools.plottools import set_mpl_style
 
 
diff --git a/artistools/inputmodel/inputmodel_misc.py b/artistools/inputmodel/inputmodel_misc.py
index 49ff0898e..22d379f1f 100644
--- a/artistools/inputmodel/inputmodel_misc.py
+++ b/artistools/inputmodel/inputmodel_misc.py
@@ -133,7 +133,7 @@ def read_modelfile_text(
         pldtypes = {col: pl.Int32 if col in {"inputcellid"} else pl.Float32 for col in columns}
 
         dfmodel = pl.read_csv(
-            at.zopen(filename, forpolars=True),
+            at.zopenpl(filename),
             separator=" ",
             comment_char="#",
             new_columns=columns,
@@ -1006,7 +1006,7 @@ def get_initelemabundances_polars(
         dtypes = {col: pl.Float32 if col.startswith("X_") else pl.Int32 for col in colnames}
 
         abundancedata = pl.read_csv(
-            at.zopen(abundancefilepath, forpolars=True),
+            at.zopenpl(abundancefilepath),
             has_header=False,
             separator=" ",
             comment_char="#",
diff --git a/artistools/lightcurve/lightcurve.py b/artistools/lightcurve/lightcurve.py
index 4bd688e0b..5045f7bd2 100644
--- a/artistools/lightcurve/lightcurve.py
+++ b/artistools/lightcurve/lightcurve.py
@@ -22,18 +22,12 @@ def readfile(
     if "_res" in str(filepath):
         # get a dict of dfs with light curves at each viewing direction bin
         lcdata_res = pl.read_csv(
-            at.zopen(filepath, forpolars=True),
-            separator=" ",
-            has_header=False,
-            new_columns=["time", "lum", "lum_cmf"],
+            at.zopenpl(filepath), separator=" ", has_header=False, new_columns=["time", "lum", "lum_cmf"]
         )
         lcdata = at.split_dataframe_dirbins(lcdata_res, index_of_repeated_value=0, output_polarsdf=True)
     else:
         dfsphericalaverage = pl.read_csv(
-            at.zopen(filepath, forpolars=True),
-            separator=" ",
-            has_header=False,
-            new_columns=["time", "lum", "lum_cmf"],
+            at.zopenpl(filepath), separator=" ", has_header=False, new_columns=["time", "lum", "lum_cmf"]
         )
 
         if list(dfsphericalaverage["time"].to_numpy()) != sorted(dfsphericalaverage["time"].to_numpy()):
diff --git a/artistools/misc.py b/artistools/misc.py
index 250f615e9..d3d63658e 100644
--- a/artistools/misc.py
+++ b/artistools/misc.py
@@ -767,31 +767,33 @@ def flatten_list(listin: list[t.Any]) -> list[t.Any]:
     return listout
 
 
-def zopen(filename: Path | str, mode: str = "rt", encoding: str | None = None, forpolars: bool = False) -> t.Any:
-    """Open filename, filename.gz or filename.xz.
-
-    Arguments:
-    ---------
-    forpolars: if polars.read_csv can read the file directly, return a Path object instead of a file object
-    """
-    if forpolars:
-        mode = "r"
+def zopen(filename: Path | str, mode: str = "rt", encoding: str | None = None) -> t.Any:
+    """Open filename, filename.ztd, filename.gz or filename.xz."""
     ext_fopen: dict[str, t.Callable] = {".zst": pyzstd.open, ".gz": gzip.open, ".xz": xz.open}
 
     for ext, fopen in ext_fopen.items():
         file_withext = str(filename) if str(filename).endswith(ext) else str(filename) + ext
         if Path(file_withext).exists():
-            if forpolars and ext in {".gz", ".zst"}:
-                return Path(file_withext)
             return fopen(file_withext, mode=mode, encoding=encoding)
 
-    if forpolars:
-        return Path(filename)
-
     # open() can raise file not found if this file doesn't exist
     return Path(filename).open(mode=mode, encoding=encoding)  # noqa: SIM115
 
 
+def zopenpl(filename: Path | str, mode: str = "rt", encoding: str | None = None) -> t.Any | Path:
+    """Open filename, filename.ztd, filename.gz or filename.xz. If polars.read_csv can read the file directly, return a Path object instead of a file object."""
+    mode = "r"
+    ext_fopen: dict[str, t.Callable] = {".zst": pyzstd.open, ".gz": gzip.open, ".xz": xz.open}
+
+    for ext, fopen in ext_fopen.items():
+        file_withext = str(filename) if str(filename).endswith(ext) else str(filename) + ext
+        if Path(file_withext).exists():
+            if ext in {".gz", ".zst"}:
+                return Path(file_withext)
+            return fopen(file_withext, mode=mode, encoding=encoding)
+    return Path(filename)
+
+
 def firstexisting(
     filelist: t.Sequence[str | Path] | str | Path,
     folder: Path | str = Path(),
diff --git a/artistools/spectra/spectra.py b/artistools/spectra/spectra.py
index 8a8925352..bf6aa6f7a 100644
--- a/artistools/spectra/spectra.py
+++ b/artistools/spectra/spectra.py
@@ -255,12 +255,7 @@ def read_spec(modelpath: Path) -> pl.DataFrame:
     print(f"Reading {specfilename}")
 
     return (
-        pl.read_csv(
-            at.zopen(specfilename, forpolars=True),
-            separator=" ",
-            infer_schema_length=0,
-            truncate_ragged_lines=True,
-        )
+        pl.read_csv(at.zopenpl(specfilename), separator=" ", infer_schema_length=0, truncate_ragged_lines=True)
         .with_columns(pl.all().cast(pl.Float64))
         .rename({"0": "nu"})
     )
@@ -276,9 +271,7 @@ def read_spec_res(modelpath: Path) -> dict[int, pl.DataFrame]:
     )
 
     print(f"Reading {specfilename} (in read_spec_res)")
-    res_specdata_in = pl.read_csv(
-        at.zopen(specfilename, forpolars=True), separator=" ", has_header=False, infer_schema_length=0
-    )
+    res_specdata_in = pl.read_csv(at.zopenpl(specfilename), separator=" ", has_header=False, infer_schema_length=0)
 
     # drop last column of nulls (caused by trailing space on each line)
     if res_specdata_in[res_specdata_in.columns[-1]].is_null().all():
@@ -326,7 +319,7 @@ def read_emission_absorption_file(emabsfilename: str | Path) -> pl.DataFrame:
         print(f" Reading {emabsfilename}")
 
     dfemabs = pl.read_csv(
-        at.zopen(emabsfilename, forpolars=True), separator=" ", has_header=False, infer_schema_length=0
+        at.zopenpl(emabsfilename), separator=" ", has_header=False, infer_schema_length=0
     ).with_columns(pl.all().cast(pl.Float32, strict=False))
 
     # drop last column of nulls (caused by trailing space on each line)
diff --git a/pyproject.toml b/pyproject.toml
index 6757f7730..0feb4c1f3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -111,7 +111,6 @@ ignore_errors = true
 max-line-length = 120
 disable = """
     broad-exception-caught,
-    #eval-used,
     fixme,
     missing-function-docstring,
     missing-module-docstring,

From e9a6be1d46c0fb263884c3d2fa75882b133d61cc Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 21:35:33 +0000
Subject: [PATCH 079/150] add test_estimator_averaging

---
 artistools/estimators/test_estimators.py | 52 ++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/artistools/estimators/test_estimators.py b/artistools/estimators/test_estimators.py
index 7c2b850dc..55c4b52a0 100644
--- a/artistools/estimators/test_estimators.py
+++ b/artistools/estimators/test_estimators.py
@@ -75,6 +75,58 @@ def test_estimator_snapshot(mockplot) -> None:
         )
 
 
+@mock.patch.object(matplotlib.axes.Axes, "plot", side_effect=matplotlib.axes.Axes.plot, autospec=True)
+def test_estimator_averaging(mockplot) -> None:
+    at.estimators.plot(argsraw=[], modelpath=modelpath, plotlist=plotlist, outputfile=outputpath, timestep="50-54")
+    xarr = [0.0, 4000.0]
+    for x in mockplot.call_args_list:
+        assert xarr == x[0][1]
+
+    # order of keys is important
+    expectedvals = {
+        "init_fe": 0.10000000149011612,
+        "init_nistable": 0.0,
+        "init_ni56": 0.8999999761581421,
+        "nne": 811131.8125,
+        "TR": 6932.65771484375,
+        "Te": 5784.4521484375,
+        "averageionisation_Fe": 1.9466091928476605,
+        "averageionisation_Ni": 1.9673294753348698,
+        "averageexcitation_FeII": float("nan"),
+        "populations_FeI": 4.668364835386799e-05,
+        "populations_FeII": 0.35026945954378863,
+        "populations_FeIII": 0.39508678896764393,
+        "populations_FeIV": 0.21220745115264195,
+        "populations_FeV": 0.042389615364484115,
+        "populations_CoII": 0.1044248111887582,
+        "populations_CoIII": 0.4759472294613869,
+        "populations_CoIV": 0.419627959349855,
+        "gamma_NT_FeI": 7.741022037400234e-06,
+        "gamma_NT_FeII": 3.7947153292832773e-06,
+        "gamma_NT_FeIII": 2.824587987164586e-06,
+        "gamma_NT_FeIV": 1.7406694591346083e-06,
+        "heating_dep": 6.849705802558503e-10,
+        "heating_coll": 2.4779998053503505e-09,
+        "heating_bf": 1.2916119454357833e-13,
+        "heating_ff": 2.1250019797070045e-16,
+        "cooling_adiabatic": 1.000458830363593e-12,
+        "cooling_coll": 3.1562059632506134e-09,
+        "cooling_fb": 5.0357105638165756e-12,
+        "cooling_ff": 1.7027620090835638e-13,
+    }
+    assert len(expectedvals) == len(mockplot.call_args_list)
+    yvals = {varname: callargs[0][2] for varname, callargs in zip(expectedvals.keys(), mockplot.call_args_list)}
+
+    print({key: yarr[1] for key, yarr in yvals.items()})
+
+    for varname, expectedval in expectedvals.items():
+        assert np.allclose([expectedval, expectedval], yvals[varname], rtol=0.001, equal_nan=True), (
+            varname,
+            expectedval,
+            yvals[varname][1],
+        )
+
+
 @mock.patch.object(matplotlib.axes.Axes, "plot", side_effect=matplotlib.axes.Axes.plot, autospec=True)
 def test_estimator_timeevolution(mockplot) -> None:
     at.estimators.plot(

From 9b68513779aaad7289f340f39ed42393d3e5a8cb Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 21:36:17 +0000
Subject: [PATCH 080/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 29 +++++++++++++++----------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index cf5f628d8..2c2ca0608 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -504,12 +504,9 @@ def plot_series(
         ax.set_ylabel(serieslabel + units_string)
         linelabel = None
     print(f"Plotting {variablename}")
-    ylist: list[float] = []
-    for modelgridindex, timesteps in zip(mgilist, timestepslist):
-        yvalue = at.estimators.get_averaged_estimators(modelpath, estimators, timesteps, modelgridindex, variablename)[
-            variablename
-        ]
-        ylist.append(yvalue)
+
+    series = estimators.group_by("xvalue").agg(pl.col(variablename).mean()).lazy().collect().sort("xvalue")
+    ylist = series[variablename].to_list()
 
     if math.log10(max(ylist) / min(ylist)) > 2 or min(ylist) == 0:
         ax.set_yscale("log")
@@ -539,7 +536,7 @@ def get_xlist(
     timestepslist: t.Any,
     modelpath: str | Path,
     args: t.Any,
-) -> tuple[list[float | int], list[int | t.Sequence[int]], list[list[int]]]:
+) -> tuple[list[float | int], list[int | t.Sequence[int]], list[list[int]], pl.LazyFrame | pl.DataFrame]:
     xlist: t.Sequence[float | int]
     if xvariable in {"cellid", "modelgridindex"}:
         mgilist_out = [mgi for mgi in allnonemptymgilist if mgi <= args.xmax] if args.xmax >= 0 else allnonemptymgilist
@@ -550,11 +547,13 @@ def get_xlist(
         check_type(timestepslist, t.Sequence[int])
         xlist = timestepslist
         timestepslist_out = timestepslist
+        estimators = estimators.with_columns(xvalue=pl.Series(xlist))
     elif xvariable == "time":
         mgilist_out = allnonemptymgilist
         timearray = at.get_timestep_times(modelpath)
         check_type(timestepslist, t.Sequence[t.Sequence[int]])
         xlist = [np.mean([timearray[ts] for ts in tslist]) for tslist in timestepslist]
+        estimators = estimators.with_columns(xvalue=pl.Series(xlist))
         timestepslist_out = timestepslist
     elif xvariable in {"velocity", "beta"}:
         dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["vel_r_mid"])
@@ -578,6 +577,12 @@ def get_xlist(
         xlist = (dfmodelcollect["vel_r_mid"] / scalefactor).to_list()
         mgilist_out = dfmodelcollect["modelgridindex"].to_list()
         timestepslist_out = timestepslist
+        estimators = estimators.filter(pl.col("modelgridindex").is_in(mgilist_out))
+        estimators = (
+            estimators.lazy()
+            .join(dfmodel.select(["modelgridindex", "vel_r_mid"]).lazy(), on="modelgridindex")
+            .rename({"vel_r_mid": "xvalue"})
+        )
     else:
         xlist = []
         mgilist_out = []
@@ -597,7 +602,7 @@ def get_xlist(
 
     assert len(xlist) == len(mgilist_out) == len(timestepslist_out)
 
-    return list(xlist), list(mgilist_out), list(timestepslist_out)
+    return list(xlist), list(mgilist_out), list(timestepslist_out), estimators
 
 
 def plot_subplot(
@@ -751,7 +756,7 @@ def make_plot(
     if not args.hidexlabel:
         axes[-1].set_xlabel(f"{xvariable}{at.estimators.get_units_string(xvariable)}")
 
-    xlist, mgilist, timestepslist = get_xlist(
+    xlist, mgilist, timestepslist, estimators = get_xlist(
         xvariable, allnonemptymgilist, estimators, timestepslist_unfiltered, modelpath, args
     )
 
@@ -1079,15 +1084,15 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         #  ['_yscale', 'linear']],
         # [['initmasses', ['Ni_56', 'He', 'C', 'Mg']]],
         # ['heating_gamma/gamma_dep'],
-        ["nne", ["_ymin", 1e5], ["_ymax", 1e11]],
+        # ["nne", ["_ymin", 1e5], ["_ymax", 1e11]],
         # ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 26000]],
         ["Te"],
         # ["Te", "TR"],
-        [["averageionisation", ["Sr"]]],
+        # [["averageionisation", ["Sr"]]],
         # [['averageexcitation', ['Fe II', 'Fe III']]],
         # [["populations", ["Sr90", "Sr91", "Sr92", "Sr93", "Sr94"]]],
         #  ['_ymin', 1e-3], ['_ymax', 5]],
-        [["populations", ["Sr II"]]],
+        # [["populations", ["Sr II"]]],
         # [['populations', ['He I', 'He II', 'He III']]],
         # [['populations', ['C I', 'C II', 'C III', 'C IV', 'C V']]],
         # [['populations', ['O I', 'O II', 'O III', 'O IV']]],

From 599edb6f5d067f5ff973f532869a0df4aa4d5cc7 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 21:37:31 +0000
Subject: [PATCH 081/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 2c2ca0608..ad693d41b 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -542,12 +542,13 @@ def get_xlist(
         mgilist_out = [mgi for mgi in allnonemptymgilist if mgi <= args.xmax] if args.xmax >= 0 else allnonemptymgilist
         xlist = list(mgilist_out)
         timestepslist_out = timestepslist
+        estimators = estimators.with_columns(xvalue=pl.col("modelgridindex"))
     elif xvariable == "timestep":
         mgilist_out = allnonemptymgilist
         check_type(timestepslist, t.Sequence[int])
         xlist = timestepslist
         timestepslist_out = timestepslist
-        estimators = estimators.with_columns(xvalue=pl.Series(xlist))
+        estimators = estimators.with_columns(xvalue=pl.col("timestep"))
     elif xvariable == "time":
         mgilist_out = allnonemptymgilist
         timearray = at.get_timestep_times(modelpath)

From 0fc853c87c8a2c9b68eafc08b7c6ce5312a2ffe6 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 21:51:17 +0000
Subject: [PATCH 082/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 63 +++++++++----------------
 1 file changed, 21 insertions(+), 42 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index ad693d41b..ea8de7974 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -352,51 +352,30 @@ def get_iontuple(ionstr):
         elsymbol = at.get_elsymbol(atomic_number)
 
         ionstr = at.get_ionstring(atomic_number, ionstage, sep="_", style="spectral")
+        if seriestype == "populations":
+            if ionstage == "ALL":
+                key = f"nnelement_{elsymbol}"
+            elif hasattr(ionstage, "lower") and ionstage.startswith(at.get_elsymbol(atomic_number)):
+                # not really an ionstage but an isotope name
+                key = f"nniso_{ionstage}"
+            else:
+                key = f"nnion_{ionstr}"
+        else:
+            key = f"{seriestype}_{ionstr}"
 
         print(f"Plotting {seriestype} {ionstr}")
-        ylist = []
-        for modelgridindex, timesteps in zip(mgilist, timestepslist):
-            if seriestype == "populations":
-                if ionstage == "ALL":
-                    key = f"nnelement_{elsymbol}"
-                elif hasattr(ionstage, "lower") and ionstage.startswith(at.get_elsymbol(atomic_number)):
-                    # not really an ionstage but an isotope name
-                    key = f"nniso_{ionstage}"
-                else:
-                    key = f"nnion_{ionstr}"
-
-                estimpop = at.estimators.get_averaged_estimators(
-                    modelpath,
-                    estimators,
-                    timesteps,
-                    modelgridindex,
-                    [key, f"nnelement_{elsymbol}", "nntot"],
-                )
 
-                nionpop = estimpop.get(key, 0.0)
-
-                try:
-                    if args.ionpoptype == "absolute":
-                        yvalue = nionpop  # Plot as fraction of element population
-                    elif args.ionpoptype == "elpop":
-                        elpop = estimpop.get(f"nnelement_{elsymbol}", 0.0)
-                        yvalue = nionpop / elpop  # Plot as fraction of element population
-                    elif args.ionpoptype == "totalpop":
-                        totalpop = estimpop["nntot"]
-                        yvalue = nionpop / totalpop  # Plot as fraction of total population
-                    else:
-                        raise AssertionError
-                except ZeroDivisionError:
-                    yvalue = 0.0
-
-                ylist.append(yvalue)
+        if seriestype != "populations" or args.ionpoptype == "absolute":
+            scalefactor = pl.lit(1)
+        elif args.ionpoptype == "elpop":
+            scalefactor = pl.col(f"nnelement_{elsymbol}").mean()
+        elif args.ionpoptype == "totalpop":
+            scalefactor = pl.col("nntot").mean()
+        else:
+            raise AssertionError
 
-            else:
-                key = f"{seriestype}_{ionstr}"
-                yvalue = at.estimators.get_averaged_estimators(modelpath, estimators, timesteps, modelgridindex, key)[
-                    key
-                ]
-                ylist.append(yvalue)
+        series = estimators.group_by("xvalue").agg(pl.col(key).mean() / scalefactor).lazy().collect().sort("xvalue")
+        ylist = series[key].to_list()
 
         plotlabel = (
             ionstage
@@ -1093,7 +1072,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         # [['averageexcitation', ['Fe II', 'Fe III']]],
         # [["populations", ["Sr90", "Sr91", "Sr92", "Sr93", "Sr94"]]],
         #  ['_ymin', 1e-3], ['_ymax', 5]],
-        # [["populations", ["Sr II"]]],
+        [["populations", ["Sr II"]]],
         # [['populations', ['He I', 'He II', 'He III']]],
         # [['populations', ['C I', 'C II', 'C III', 'C IV', 'C V']]],
         # [['populations', ['O I', 'O II', 'O III', 'O IV']]],

From b60178e79be84d12e784c18a82e8e8365f3c7a85 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 22:24:03 +0000
Subject: [PATCH 083/150] Speed up plotting with polars

---
 artistools/estimators/estimators.py     |  8 ++-
 artistools/estimators/plotestimators.py | 83 +++++--------------------
 2 files changed, 21 insertions(+), 70 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index f618e7feb..41e6e79ff 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -221,7 +221,7 @@ def read_estimators_from_file(
     )
 
 
-def batched(iterable, n):
+def batched(iterable, n):  # -> Generator[list, Any, None]:
     """Batch data into iterators of length n. The last batch may be shorter."""
     # batched('ABCDEFG', 3) --> ABC DEF G
     if n < 1:
@@ -420,7 +420,11 @@ def get_averageionisation(estimatorstsmgi: pl.LazyFrame, atomic_number: int) ->
         cs.starts_with(f"nnion_{elsymb}_") | cs.by_name(f"nnelement_{elsymb}")
     ).collect()
 
-    nnelement = dfselected[f"nnelement_{elsymb}"].item(0)
+    dfnnelement = dfselected[f"nnelement_{elsymb}"]
+    if dfnnelement.is_empty():
+        return float("NaN")
+
+    nnelement = dfnnelement.item(0)
     if nnelement is None:
         return float("NaN")
 
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index ea8de7974..f3de09c8c 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -14,7 +14,6 @@
 import argcomplete
 import matplotlib.pyplot as plt
 import numpy as np
-import pandas as pd
 import polars as pl
 from typeguard import check_type
 
@@ -51,9 +50,7 @@ def get_ylabel(variable):
     return ""
 
 
-def plot_init_abundances(
-    ax, xlist, specieslist, mgilist, modelpath, seriestype, dfalldata=None, args=None, **plotkwargs
-):
+def plot_init_abundances(ax, xlist, specieslist, mgilist, modelpath, seriestype, args=None, **plotkwargs):
     assert len(xlist) - 1 == len(mgilist)
 
     if seriestype == "initabundances":
@@ -100,9 +97,6 @@ def plot_init_abundances(
                 yvalue = mergemodelabundata.loc[modelgridindex][f"{valuetype}{elsymbol}"]
             ylist.append(yvalue)
 
-        if dfalldata is not None:
-            dfalldata["initabundances." + speciesstr] = ylist
-
         ylist.insert(0, ylist[0])
         # or ax.step(where='pre', )
         color = get_elemcolor(atomic_number=atomic_number)
@@ -124,7 +118,6 @@ def plot_average_ionisation_excitation(
     mgilist,
     estimators,
     modelpath,
-    dfalldata=None,
     args=None,
     **plotkwargs,
 ):
@@ -178,9 +171,6 @@ def plot_average_ionisation_excitation(
 
         color = get_elemcolor(atomic_number=atomic_number)
 
-        if dfalldata is not None:
-            dfalldata[seriestype + "." + paramvalue] = ylist
-
         ylist.insert(0, ylist[0])
 
         xlist, ylist = at.estimators.apply_filters(xlist, ylist, args)
@@ -197,7 +187,6 @@ def plot_levelpop(
     mgilist: t.Sequence[int | t.Sequence[int]],
     estimators: pl.LazyFrame | pl.DataFrame,
     modelpath: str | Path,
-    dfalldata: pd.DataFrame | None,
     args: argparse.Namespace,
     **plotkwargs: t.Any,
 ):
@@ -260,15 +249,6 @@ def plot_levelpop(
             else:
                 ylist.append(valuesum / tdeltasum)
 
-        if dfalldata is not None:
-            elsym = at.get_elsymbol(atomic_number).lower()
-            colname = (
-                f"nlevel_on_dv_{elsym}_ionstage{ionstage}_level{levelindex}"
-                if seriestype == "levelpopulation_dn_on_dvel"
-                else f"nnlevel_{elsym}_ionstage{ionstage}_level{levelindex}"
-            )
-            dfalldata[colname] = ylist
-
         ylist.insert(0, ylist[0])
 
         xlist, ylist = at.estimators.apply_filters(xlist, ylist, args)
@@ -285,7 +265,6 @@ def plot_multi_ion_series(
     mgilist: t.Sequence[int | t.Sequence[int]],
     estimators: pl.LazyFrame | pl.DataFrame,
     modelpath: str | Path,
-    dfalldata: pd.DataFrame | None,
     args: argparse.Namespace,
     **plotkwargs: t.Any,
 ):
@@ -414,16 +393,6 @@ def get_iontuple(ionstr):
         # color = f'C{colorindex}'
         # or ax.step(where='pre', )
 
-        if dfalldata is not None:
-            elsym = at.get_elsymbol(atomic_number).lower()
-            if args.ionpoptype == "absolute":
-                colname = f"nnion_{elsym}_ionstage{ionstage}"
-            elif args.ionpoptype == "elpop":
-                colname = f"nnion_over_nnelem_{elsym}_ionstage{ionstage}"
-            elif args.ionpoptype == "totalpop":
-                colname = f"nnion_over_nntot_{elsym}_ionstage{ionstage}"
-            dfalldata[colname] = ylist
-
         ylist.insert(0, ylist[0])
 
         xlist, ylist = at.estimators.apply_filters(xlist, ylist, args)
@@ -466,7 +435,6 @@ def plot_series(
     estimators: pl.LazyFrame | pl.DataFrame,
     args: argparse.Namespace,
     nounits: bool = False,
-    dfalldata: pd.DataFrame | None = None,
     **plotkwargs: t.Any,
 ) -> None:
     """Plot something like Te or TR."""
@@ -484,8 +452,9 @@ def plot_series(
         linelabel = None
     print(f"Plotting {variablename}")
 
-    series = estimators.group_by("xvalue").agg(pl.col(variablename).mean()).lazy().collect().sort("xvalue")
+    series = estimators.group_by("xvalue").agg(pl.col(variablename).mean()).lazy().collect()
     ylist = series[variablename].to_list()
+    xlist2 = series["xvalue"].to_list()
 
     if math.log10(max(ylist) / min(ylist)) > 2 or min(ylist) == 0:
         ax.set_yscale("log")
@@ -496,12 +465,10 @@ def plot_series(
         # 'cooling_adiabatic': 'blue'
     }
 
-    if dfalldata is not None:
-        dfalldata[variablename] = ylist
-
+    xlist2.insert(0, xlist[0])
     ylist.insert(0, ylist[0])
 
-    xlist_filtered, ylist_filtered = at.estimators.apply_filters(xlist, ylist, args)
+    xlist_filtered, ylist_filtered = at.estimators.apply_filters(xlist2, ylist, args)
 
     ax.plot(
         xlist_filtered, ylist_filtered, linewidth=1.5, label=linelabel, color=dictcolors.get(variablename), **plotkwargs
@@ -558,11 +525,10 @@ def get_xlist(
         mgilist_out = dfmodelcollect["modelgridindex"].to_list()
         timestepslist_out = timestepslist
         estimators = estimators.filter(pl.col("modelgridindex").is_in(mgilist_out))
-        estimators = (
-            estimators.lazy()
-            .join(dfmodel.select(["modelgridindex", "vel_r_mid"]).lazy(), on="modelgridindex")
-            .rename({"vel_r_mid": "xvalue"})
-        )
+        estimators = estimators.lazy().join(dfmodel.select(["modelgridindex", "vel_r_mid"]).lazy(), on="modelgridindex")
+        estimators = estimators.with_columns(xvalue=(pl.col("vel_r_mid") / scalefactor))
+        estimators = estimators.sort("xvalue")
+        xlist = estimators.select("xvalue").collect()["xvalue"].to_list()
     else:
         xlist = []
         mgilist_out = []
@@ -593,7 +559,6 @@ def plot_subplot(
     mgilist: list[int | t.Sequence[int]],
     modelpath: str | Path,
     estimators: pl.LazyFrame | pl.DataFrame,
-    dfalldata: pd.DataFrame | None,
     args: argparse.Namespace,
     **plotkwargs: t.Any,
 ):
@@ -627,7 +592,6 @@ def plot_subplot(
                 estimators,
                 args,
                 nounits=sameylabel,
-                dfalldata=dfalldata,
                 **plotkwargs,
             )
             if showlegend and sameylabel and ylabel is not None:
@@ -637,7 +601,7 @@ def plot_subplot(
             seriestype, params = plotitem
 
             if seriestype in {"initabundances", "initmasses"}:
-                plot_init_abundances(ax, xlist, params, mgilist, modelpath, seriestype, dfalldata=dfalldata, args=args)
+                plot_init_abundances(ax, xlist, params, mgilist, modelpath, seriestype, args=args)
 
             elif seriestype == "levelpopulation" or seriestype.startswith("levelpopulation_"):
                 plot_levelpop(
@@ -649,7 +613,6 @@ def plot_subplot(
                     mgilist,
                     estimators,
                     modelpath,
-                    dfalldata=dfalldata,
                     args=args,
                 )
 
@@ -663,7 +626,6 @@ def plot_subplot(
                     mgilist,
                     estimators,
                     modelpath,
-                    dfalldata=dfalldata,
                     args=args,
                     **plotkwargs,
                 )
@@ -688,7 +650,6 @@ def plot_subplot(
                     mgilist,
                     estimators,
                     modelpath,
-                    dfalldata,
                     args,
                     **plotkwargs,
                 )
@@ -740,10 +701,6 @@ def make_plot(
         xvariable, allnonemptymgilist, estimators, timestepslist_unfiltered, modelpath, args
     )
 
-    dfalldata = pd.DataFrame()
-    # dfalldata.index.name = "modelgridindex"
-    dfalldata[xvariable] = xlist
-
     xlist = list(
         np.insert(xlist, 0, 0.0)
         if (xvariable.startswith("velocity") or xvariable == "beta")
@@ -784,7 +741,6 @@ def make_plot(
             mgilist,
             modelpath,
             estimators,
-            dfalldata=dfalldata,
             args=args,
             **plotkwargs,
         )
@@ -829,12 +785,6 @@ def make_plot(
         axes[0].set_title(figure_title, fontsize=8)
     # plt.suptitle(figure_title, fontsize=11, verticalalignment='top')
 
-    if args.write_data:
-        dfalldata = dfalldata.sort_index()
-        dataoutfilename = Path(outfilename).with_suffix(".txt")
-        dfalldata.to_csv(dataoutfilename)
-        print(f"Saved {dataoutfilename}")
-
     fig.savefig(outfilename)
     print(f"Saved {outfilename}")
 
@@ -992,8 +942,6 @@ def addargs(parser: argparse.ArgumentParser) -> None:
 
     parser.add_argument("--show", action="store_true", help="Show plot before quitting")
 
-    parser.add_argument("--write_data", action="store_true", help="Save data used to generate the plot in a CSV file")
-
     parser.add_argument(
         "-o", action="store", dest="outputfile", type=Path, default=Path(), help="Filename for PDF file"
     )
@@ -1064,15 +1012,14 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         #  ['_yscale', 'linear']],
         # [['initmasses', ['Ni_56', 'He', 'C', 'Mg']]],
         # ['heating_gamma/gamma_dep'],
-        # ["nne", ["_ymin", 1e5], ["_ymax", 1e11]],
-        # ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 26000]],
-        ["Te"],
+        ["nne", ["_ymin", 1e5], ["_ymax", 1e11]],
+        ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 26000]],
+        # ["Te"],
         # ["Te", "TR"],
-        # [["averageionisation", ["Sr"]]],
+        [["averageionisation", ["Sr"]]],
         # [['averageexcitation', ['Fe II', 'Fe III']]],
         # [["populations", ["Sr90", "Sr91", "Sr92", "Sr93", "Sr94"]]],
-        #  ['_ymin', 1e-3], ['_ymax', 5]],
-        [["populations", ["Sr II"]]],
+        # [["populations", ["Sr II"]]],
         # [['populations', ['He I', 'He II', 'He III']]],
         # [['populations', ['C I', 'C II', 'C III', 'C IV', 'C V']]],
         # [['populations', ['O I', 'O II', 'O III', 'O IV']]],

From 27d230625e513ff36c7cde827d08747fc504da14 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 22:30:08 +0000
Subject: [PATCH 084/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index f3de09c8c..26fa85a68 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -597,13 +597,14 @@ def plot_subplot(
             if showlegend and sameylabel and ylabel is not None:
                 ax.set_ylabel(ylabel)
         else:  # it's a sequence of values
-            showlegend = True
             seriestype, params = plotitem
 
             if seriestype in {"initabundances", "initmasses"}:
+                showlegend = True
                 plot_init_abundances(ax, xlist, params, mgilist, modelpath, seriestype, args=args)
 
             elif seriestype == "levelpopulation" or seriestype.startswith("levelpopulation_"):
+                showlegend = True
                 plot_levelpop(
                     ax,
                     xlist,
@@ -617,6 +618,7 @@ def plot_subplot(
                 )
 
             elif seriestype in {"averageionisation", "averageexcitation"}:
+                showlegend = True
                 plot_average_ionisation_excitation(
                     ax,
                     xlist,
@@ -640,6 +642,7 @@ def plot_subplot(
                 ax.set_yscale(params)
 
             else:
+                showlegend = True
                 seriestype, ionlist = plotitem
                 plot_multi_ion_series(
                     ax,
@@ -842,7 +845,7 @@ def plot_recombrates(modelpath, estimators, atomic_number, ionstage_list, **plot
                 dfrates.T_e,
                 dfrates.rrc_total,
                 linewidth=2,
-                label=ionstr + " (calibration)",
+                label=f"{ionstr} (calibration)",
                 markersize=6,
                 marker="s",
                 **plotkwargs,
@@ -1012,8 +1015,8 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         #  ['_yscale', 'linear']],
         # [['initmasses', ['Ni_56', 'He', 'C', 'Mg']]],
         # ['heating_gamma/gamma_dep'],
-        ["nne", ["_ymin", 1e5], ["_ymax", 1e11]],
-        ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 26000]],
+        ["nne", ["_ymin", 1e5], ["_ymax", 1e10]],
+        ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 16000]],
         # ["Te"],
         # ["Te", "TR"],
         [["averageionisation", ["Sr"]]],

From 3ad315d37420452c59085a1051d911e7b8ef35d7 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 23:19:52 +0000
Subject: [PATCH 085/150] Fix averageexcitation for averaged estimator

---
 artistools/estimators/estimators.py      |  4 ++--
 artistools/estimators/plotestimators.py  | 30 ++++++++++++++----------
 artistools/estimators/test_estimators.py |  2 +-
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 41e6e79ff..bbbc9fad8 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -445,7 +445,7 @@ def get_averageionisation(estimatorstsmgi: pl.LazyFrame, atomic_number: int) ->
 
 def get_averageexcitation(
     modelpath: Path, modelgridindex: int, timestep: int, atomic_number: int, ionstage: int, T_exc: float
-) -> float:
+) -> float | None:
     dfnltepops = at.nltepops.read_files(modelpath, modelgridindex=modelgridindex, timestep=timestep)
     adata = at.atomic.get_levels(modelpath)
     ionlevels = adata.query("Z == @atomic_number and ionstage == @ionstage").iloc[0].levels
@@ -453,7 +453,7 @@ def get_averageexcitation(
     energypopsum = 0
     ionpopsum = 0
     if dfnltepops.empty:
-        return float("NaN")
+        return None
 
     dfnltepops_ion = dfnltepops.query(
         "modelgridindex==@modelgridindex and timestep==@timestep and Z==@atomic_number & ionstage==@ionstage"
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 26fa85a68..a7eef9d83 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -136,11 +136,11 @@ def plot_average_ionisation_excitation(
             atomic_number = at.get_atomic_number(paramvalue.split(" ")[0])
             ionstage = at.decode_roman_numeral(paramvalue.split(" ")[1])
         ylist = []
-        for modelgridindex, timesteps in zip(mgilist, timestepslist):
-            valuesum = 0
-            tdeltasum = 0
-            for timestep in timesteps:
-                if seriestype == "averageionisation":
+        if seriestype == "averageionisation":
+            for modelgridindex, timesteps in zip(mgilist, timestepslist):
+                valuesum = 0
+                tdeltasum = 0
+                for timestep in timesteps:
                     valuesum += (
                         at.estimators.get_averageionisation(
                             estimators.filter(pl.col("timestep") == timestep).filter(
@@ -150,7 +150,13 @@ def plot_average_ionisation_excitation(
                         )
                         * arr_tdelta[timestep]
                     )
-                elif seriestype == "averageexcitation":
+                    tdeltasum += arr_tdelta[timestep]
+                ylist.append(valuesum / tdeltasum)
+        elif seriestype == "averageexcitation":
+            for modelgridindex, timesteps in zip(mgilist, timestepslist):
+                valuesum = 0
+                tdeltasum = 0
+                for timestep in timesteps:
                     T_exc = (
                         estimators.filter(pl.col("timestep") == timestep)
                         .filter(pl.col("modelgridindex") == modelgridindex)
@@ -159,13 +165,13 @@ def plot_average_ionisation_excitation(
                         .collect()
                         .item(0, 0)
                     )
-                    valuesum += (
-                        at.estimators.get_averageexcitation(
-                            modelpath, modelgridindex, timestep, atomic_number, ionstage, T_exc
-                        )
-                        * arr_tdelta[timestep]
+                    exc = at.estimators.get_averageexcitation(
+                        modelpath, modelgridindex, timestep, atomic_number, ionstage, T_exc
                     )
-                tdeltasum += arr_tdelta[timestep]
+                    if exc is None:
+                        continue
+                    valuesum += exc * arr_tdelta[timestep]
+                    tdeltasum += arr_tdelta[timestep]
 
             ylist.append(valuesum / tdeltasum)
 
diff --git a/artistools/estimators/test_estimators.py b/artistools/estimators/test_estimators.py
index 55c4b52a0..12c80b9d4 100644
--- a/artistools/estimators/test_estimators.py
+++ b/artistools/estimators/test_estimators.py
@@ -92,7 +92,7 @@ def test_estimator_averaging(mockplot) -> None:
         "Te": 5784.4521484375,
         "averageionisation_Fe": 1.9466091928476605,
         "averageionisation_Ni": 1.9673294753348698,
-        "averageexcitation_FeII": float("nan"),
+        "averageexcitation_FeII": 0.1975447074846265,
         "populations_FeI": 4.668364835386799e-05,
         "populations_FeII": 0.35026945954378863,
         "populations_FeIII": 0.39508678896764393,

From e4ef13182e41d3ddc5ebc84c5fe3275771778684 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Fri, 10 Nov 2023 23:21:23 +0000
Subject: [PATCH 086/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index a7eef9d83..131e9a630 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -137,6 +137,7 @@ def plot_average_ionisation_excitation(
             ionstage = at.decode_roman_numeral(paramvalue.split(" ")[1])
         ylist = []
         if seriestype == "averageionisation":
+            # TODO: replace loop with Polars groupby
             for modelgridindex, timesteps in zip(mgilist, timestepslist):
                 valuesum = 0
                 tdeltasum = 0

From dc730ab0bdc1166a0c2a8a637409c4a2a01effc7 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 12:56:05 +0000
Subject: [PATCH 087/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 39 ++++++++++++-------------
 1 file changed, 19 insertions(+), 20 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 131e9a630..fc4d77a5f 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -433,11 +433,9 @@ def get_iontuple(ionstr):
 
 def plot_series(
     ax: plt.Axes,
-    xlist: t.Sequence[int | float] | np.ndarray,
+    xvariable: str,
     variablename: str,
     showlegend: bool,
-    timestepslist: t.Sequence[t.Sequence[int]],
-    mgilist: t.Sequence[int | t.Sequence[int]],
     modelpath: str | Path,
     estimators: pl.LazyFrame | pl.DataFrame,
     args: argparse.Namespace,
@@ -445,7 +443,6 @@ def plot_series(
     **plotkwargs: t.Any,
 ) -> None:
     """Plot something like Te or TR."""
-    assert len(xlist) - 1 == len(mgilist) == len(timestepslist)
     formattedvariablename = at.estimators.get_dictlabelreplacements().get(variablename, variablename)
     serieslabel = f"{formattedvariablename}"
     units_string = at.estimators.get_units_string(variablename)
@@ -459,9 +456,9 @@ def plot_series(
         linelabel = None
     print(f"Plotting {variablename}")
 
-    series = estimators.group_by("xvalue").agg(pl.col(variablename).mean()).lazy().collect()
+    series = estimators.group_by("xvalue", maintain_order=True).agg(pl.col(variablename).mean()).lazy().collect()
     ylist = series[variablename].to_list()
-    xlist2 = series["xvalue"].to_list()
+    xlist = series["xvalue"].to_list()
 
     if math.log10(max(ylist) / min(ylist)) > 2 or min(ylist) == 0:
         ax.set_yscale("log")
@@ -472,10 +469,12 @@ def plot_series(
         # 'cooling_adiabatic': 'blue'
     }
 
-    xlist2.insert(0, xlist[0])
-    ylist.insert(0, ylist[0])
+    if xvariable.startswith("velocity") or xvariable == "beta":
+        # make a line segment from 0 velocity
+        xlist.insert(0, 0.0)
+        ylist.insert(0, ylist[0])
 
-    xlist_filtered, ylist_filtered = at.estimators.apply_filters(xlist2, ylist, args)
+    xlist_filtered, ylist_filtered = at.estimators.apply_filters(xlist, ylist, args)
 
     ax.plot(
         xlist_filtered, ylist_filtered, linewidth=1.5, label=linelabel, color=dictcolors.get(variablename), **plotkwargs
@@ -521,13 +520,13 @@ def get_xlist(
         if args.readonlymgi:
             dfmodel = dfmodel.filter(pl.col("modelgridindex").is_in(args.modelgridindex))
         dfmodel = dfmodel.select(["modelgridindex", "vel_r_mid"]).sort(by="vel_r_mid")
+        scalefactor = 1e5 if xvariable == "velocity" else 29979245800
         if args.xmax > 0:
-            dfmodel = dfmodel.filter(pl.col("vel_r_mid") / 1e5 <= args.xmax)
+            dfmodel = dfmodel.filter(pl.col("vel_r_mid") / scalefactor <= args.xmax)
         else:
             dfmodel = dfmodel.filter(pl.col("vel_r_mid") <= modelmeta["vmax_cmps"])
         dfmodelcollect = dfmodel.select(["vel_r_mid", "modelgridindex"]).collect()
 
-        scalefactor = 1e5 if xvariable == "velocity" else 29979245800
         xlist = (dfmodelcollect["vel_r_mid"] / scalefactor).to_list()
         mgilist_out = dfmodelcollect["modelgridindex"].to_list()
         timestepslist_out = timestepslist
@@ -562,6 +561,7 @@ def plot_subplot(
     ax: plt.Axes,
     timestepslist: list[list[int]],
     xlist: list[float | int],
+    xvariable: str,
     plotitems: list[t.Any],
     mgilist: list[int | t.Sequence[int]],
     modelpath: str | Path,
@@ -589,15 +589,13 @@ def plot_subplot(
         if isinstance(plotitem, str):
             showlegend = seriescount > 1 or len(plotitem) > 20 or not sameylabel
             plot_series(
-                ax,
-                xlist,
-                plotitem,
-                showlegend,
-                timestepslist,
-                mgilist,
-                modelpath,
-                estimators,
-                args,
+                ax=ax,
+                xvariable=xvariable,
+                variablename=plotitem,
+                showlegend=showlegend,
+                modelpath=modelpath,
+                estimators=estimators,
+                args=args,
                 nounits=sameylabel,
                 **plotkwargs,
             )
@@ -747,6 +745,7 @@ def make_plot(
             ax,
             timestepslist,
             xlist,
+            xvariable,
             plotitems,
             mgilist,
             modelpath,

From 57d1e1ecbc030280af62bcebb1fa69ceb5eabd4d Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 13:20:02 +0000
Subject: [PATCH 088/150] Improve warnings

---
 artistools/estimators/estimators.py     |  4 +++
 artistools/estimators/plotestimators.py | 34 ++++++++++++-------------
 artistools/nltepops/nltepops.py         |  6 ++++-
 3 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index bbbc9fad8..126acc908 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -447,6 +447,9 @@ def get_averageexcitation(
     modelpath: Path, modelgridindex: int, timestep: int, atomic_number: int, ionstage: int, T_exc: float
 ) -> float | None:
     dfnltepops = at.nltepops.read_files(modelpath, modelgridindex=modelgridindex, timestep=timestep)
+    if dfnltepops.empty:
+        print(f"WARNING: NLTE pops not found for cell {modelgridindex} at timestep {timestep}")
+
     adata = at.atomic.get_levels(modelpath)
     ionlevels = adata.query("Z == @atomic_number and ionstage == @ionstage").iloc[0].levels
 
@@ -477,6 +480,7 @@ def get_averageexcitation(
         boltzfac_sum = ionlevels.iloc[levelnumber_sl:].eval("g * exp(- energy_ev / @k_b / @T_exc)").sum()
         # adjust to the actual superlevel population from ARTIS
         energypopsum += energy_boltzfac_sum * superlevelrow.n_NLTE / boltzfac_sum
+
     return energypopsum / ionpopsum
 
 
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index fc4d77a5f..3f4d19b6b 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -130,6 +130,7 @@ def plot_average_ionisation_excitation(
 
     arr_tdelta = at.get_timestep_times(modelpath, loc="delta")
     for paramvalue in params:
+        print(f"Plotting {seriestype} {paramvalue}")
         if seriestype == "averageionisation":
             atomic_number = at.get_atomic_number(paramvalue)
         else:
@@ -152,11 +153,12 @@ def plot_average_ionisation_excitation(
                         * arr_tdelta[timestep]
                     )
                     tdeltasum += arr_tdelta[timestep]
-                ylist.append(valuesum / tdeltasum)
+                ylist.append(valuesum / tdeltasum if tdeltasum > 0 else float("nan"))
         elif seriestype == "averageexcitation":
+            print("  This will be slow!")
             for modelgridindex, timesteps in zip(mgilist, timestepslist):
-                valuesum = 0
-                tdeltasum = 0
+                exc_ev_times_tdelta_sum = 0.0
+                tdeltasum = 0.0
                 for timestep in timesteps:
                     T_exc = (
                         estimators.filter(pl.col("timestep") == timestep)
@@ -166,18 +168,18 @@ def plot_average_ionisation_excitation(
                         .collect()
                         .item(0, 0)
                     )
-                    exc = at.estimators.get_averageexcitation(
+                    exc_ev = at.estimators.get_averageexcitation(
                         modelpath, modelgridindex, timestep, atomic_number, ionstage, T_exc
                     )
-                    if exc is None:
-                        continue
-                    valuesum += exc * arr_tdelta[timestep]
-                    tdeltasum += arr_tdelta[timestep]
-
-            ylist.append(valuesum / tdeltasum)
+                    if exc_ev is not None:
+                        exc_ev_times_tdelta_sum += exc_ev * arr_tdelta[timestep]
+                        tdeltasum += arr_tdelta[timestep]
+            if tdeltasum == 0.0:
+                msg = f"ERROR: No excitation data found for {paramvalue}"
+                raise ValueError(msg)
+            ylist.append(exc_ev_times_tdelta_sum / tdeltasum if tdeltasum > 0 else float("nan"))
 
         color = get_elemcolor(atomic_number=atomic_number)
-
         ylist.insert(0, ylist[0])
 
         xlist, ylist = at.estimators.apply_filters(xlist, ylist, args)
@@ -732,11 +734,9 @@ def make_plot(
                 allts.add(ts)
 
     estimators = (
-        estimators.filter(pl.col("modelgridindex").is_in(mgilist))
-        .filter(pl.col("timestep").is_in(allts))
-        .lazy()
-        .collect()
-        .lazy()
+        estimators.filter(pl.col("modelgridindex").is_in(mgilist)).filter(pl.col("timestep").is_in(allts)).lazy()
+        # .collect()
+        # .lazy()
     )
 
     for ax, plotitems in zip(axes, plotlist):
@@ -1085,7 +1085,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
     assert estimators is not None
 
     for ts in reversed(timesteps_included):
-        tswithdata = estimators.select("timestep").unique().collect().to_numpy()
+        tswithdata = estimators.select("timestep").unique().collect().to_series()
         for ts in timesteps_included:
             if ts not in tswithdata:
                 timesteps_included.remove(ts)
diff --git a/artistools/nltepops/nltepops.py b/artistools/nltepops/nltepops.py
index 60eba59bd..6e4be53cf 100644
--- a/artistools/nltepops/nltepops.py
+++ b/artistools/nltepops/nltepops.py
@@ -177,7 +177,11 @@ def read_file_filtered(nltefilepath, strquery=None, dfqueryvars=None):
 
 @lru_cache(maxsize=2)
 def read_files(
-    modelpath, timestep=-1, modelgridindex=-1, dfquery=None, dfqueryvars: dict | None = None
+    modelpath: str | Path,
+    timestep: int = -1,
+    modelgridindex: int = -1,
+    dfquery: str | None = None,
+    dfqueryvars: dict | None = None,
 ) -> pd.DataFrame:
     """Read in NLTE populations from a model for a particular timestep and grid cell."""
     if dfqueryvars is None:

From 2308521e2491e38cb601e20b5faf43a6cfc975d5 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 13:30:36 +0000
Subject: [PATCH 089/150] Improve ionisation calculation with polars
 expressions

---
 artistools/estimators/estimators.py     | 32 ++++++++++++-------------
 artistools/estimators/plotestimators.py |  6 ++---
 2 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 126acc908..6eb893797 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -413,12 +413,13 @@ def get_averaged_estimators(
 
 
 def get_averageionisation(estimatorstsmgi: pl.LazyFrame, atomic_number: int) -> float:
-    free_electron_weighted_pop_sum = 0.0
     elsymb = at.get_elsymbol(atomic_number)
 
-    dfselected = estimatorstsmgi.select(
-        cs.starts_with(f"nnion_{elsymb}_") | cs.by_name(f"nnelement_{elsymb}")
-    ).collect()
+    dfselected = (
+        estimatorstsmgi.select(cs.starts_with(f"nnion_{elsymb}_") | cs.by_name(f"nnelement_{elsymb}"))
+        .fill_null(0.0)
+        .collect()
+    )
 
     dfnnelement = dfselected[f"nnelement_{elsymb}"]
     if dfnnelement.is_empty():
@@ -428,19 +429,16 @@ def get_averageionisation(estimatorstsmgi: pl.LazyFrame, atomic_number: int) ->
     if nnelement is None:
         return float("NaN")
 
-    found = False
-    popsum = 0.0
-    for key in dfselected.select(cs.starts_with(f"nnion_{elsymb}_")).columns:
-        found = True
-        nnion = dfselected[key].item(0)
-        if nnion is None:
-            continue
-
-        ionstage = at.decode_roman_numeral(key.removeprefix(f"nnion_{elsymb}_"))
-        free_electron_weighted_pop_sum += nnion * (ionstage - 1)
-        popsum += nnion
-
-    return free_electron_weighted_pop_sum / nnelement if found else float("NaN")
+    ioncols = [col for col in dfselected.columns if col.startswith(f"nnion_{elsymb}_")]
+    if not ioncols:
+        return float("NaN")
+    ioncharges = [at.decode_roman_numeral(col.removeprefix(f"nnion_{elsymb}_")) - 1 for col in ioncols]
+    return (
+        dfselected.select(
+            pl.sum_horizontal([pl.col(ioncol) * ioncharge for ioncol, ioncharge in zip(ioncols, ioncharges)])
+        ).item(0, 0)
+        / nnelement
+    )
 
 
 def get_averageexcitation(
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 3f4d19b6b..5ef4f3d61 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -1026,9 +1026,9 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         # ["Te"],
         # ["Te", "TR"],
         [["averageionisation", ["Sr"]]],
-        # [['averageexcitation', ['Fe II', 'Fe III']]],
-        # [["populations", ["Sr90", "Sr91", "Sr92", "Sr93", "Sr94"]]],
-        # [["populations", ["Sr II"]]],
+        # [["averageexcitation", ["Fe II", "Fe III"]]],
+        [["populations", ["Sr90", "Sr91", "Sr92", "Sr93", "Sr94"]]],
+        [["populations", ["Sr I", "Sr II", "Sr III"]]],
         # [['populations', ['He I', 'He II', 'He III']]],
         # [['populations', ['C I', 'C II', 'C III', 'C IV', 'C V']]],
         # [['populations', ['O I', 'O II', 'O III', 'O IV']]],

From d35bae95592d19ac5523c031a511bee3dc72bb27 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 13:35:36 +0000
Subject: [PATCH 090/150] Remove get_averageionisation()

---
 artistools/estimators/__init__.py       |  1 -
 artistools/estimators/estimators.py     | 30 -------------------
 artistools/estimators/plotestimators.py | 40 ++++++++++++++++++++-----
 3 files changed, 32 insertions(+), 39 deletions(-)

diff --git a/artistools/estimators/__init__.py b/artistools/estimators/__init__.py
index 724a6c6d2..9e0baae80 100644
--- a/artistools/estimators/__init__.py
+++ b/artistools/estimators/__init__.py
@@ -6,7 +6,6 @@
 from artistools.estimators.estimators import apply_filters
 from artistools.estimators.estimators import get_averaged_estimators
 from artistools.estimators.estimators import get_averageexcitation
-from artistools.estimators.estimators import get_averageionisation
 from artistools.estimators.estimators import get_dictlabelreplacements
 from artistools.estimators.estimators import get_ionrecombrates_fromfile
 from artistools.estimators.estimators import get_partiallycompletetimesteps
diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 6eb893797..da559e51b 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -16,7 +16,6 @@
 import numpy as np
 import pandas as pd
 import polars as pl
-import polars.selectors as cs
 
 import artistools as at
 
@@ -412,35 +411,6 @@ def get_averaged_estimators(
     return dictout
 
 
-def get_averageionisation(estimatorstsmgi: pl.LazyFrame, atomic_number: int) -> float:
-    elsymb = at.get_elsymbol(atomic_number)
-
-    dfselected = (
-        estimatorstsmgi.select(cs.starts_with(f"nnion_{elsymb}_") | cs.by_name(f"nnelement_{elsymb}"))
-        .fill_null(0.0)
-        .collect()
-    )
-
-    dfnnelement = dfselected[f"nnelement_{elsymb}"]
-    if dfnnelement.is_empty():
-        return float("NaN")
-
-    nnelement = dfnnelement.item(0)
-    if nnelement is None:
-        return float("NaN")
-
-    ioncols = [col for col in dfselected.columns if col.startswith(f"nnion_{elsymb}_")]
-    if not ioncols:
-        return float("NaN")
-    ioncharges = [at.decode_roman_numeral(col.removeprefix(f"nnion_{elsymb}_")) - 1 for col in ioncols]
-    return (
-        dfselected.select(
-            pl.sum_horizontal([pl.col(ioncol) * ioncharge for ioncol, ioncharge in zip(ioncols, ioncharges)])
-        ).item(0, 0)
-        / nnelement
-    )
-
-
 def get_averageexcitation(
     modelpath: Path, modelgridindex: int, timestep: int, atomic_number: int, ionstage: int, T_exc: float
 ) -> float | None:
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 5ef4f3d61..e10aa1382 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -15,6 +15,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import polars as pl
+import polars.selectors as cs
 from typeguard import check_type
 
 import artistools as at
@@ -143,15 +144,38 @@ def plot_average_ionisation_excitation(
                 valuesum = 0
                 tdeltasum = 0
                 for timestep in timesteps:
-                    valuesum += (
-                        at.estimators.get_averageionisation(
-                            estimators.filter(pl.col("timestep") == timestep).filter(
-                                pl.col("modelgridindex") == modelgridindex
-                            ),
-                            atomic_number,
-                        )
-                        * arr_tdelta[timestep]
+                    estimatorstsmgi = estimators.filter(pl.col("timestep") == timestep).filter(
+                        pl.col("modelgridindex") == modelgridindex
                     )
+                    elsymb = at.get_elsymbol(atomic_number)
+                    dfselected = (
+                        estimatorstsmgi.select(cs.starts_with(f"nnion_{elsymb}_") | cs.by_name(f"nnelement_{elsymb}"))
+                        .fill_null(0.0)
+                        .collect()
+                    )
+
+                    dfnnelement = dfselected[f"nnelement_{elsymb}"]
+                    if dfnnelement.is_empty():
+                        myval = float("NaN")
+                    else:
+                        nnelement = dfnnelement.item(0)
+                        if nnelement is None:
+                            myval = float("NaN")
+                        elif ioncols := [col for col in dfselected.columns if col.startswith(f"nnion_{elsymb}_")]:
+                            ioncharges = [
+                                at.decode_roman_numeral(col.removeprefix(f"nnion_{elsymb}_")) - 1 for col in ioncols
+                            ]
+                            myval = (
+                                dfselected.select(
+                                    pl.sum_horizontal(
+                                        [pl.col(ioncol) * ioncharge for ioncol, ioncharge in zip(ioncols, ioncharges)]
+                                    )
+                                ).item(0, 0)
+                                / nnelement
+                            )
+                        else:
+                            myval = float("NaN")
+                    valuesum += myval * arr_tdelta[timestep]
                     tdeltasum += arr_tdelta[timestep]
                 ylist.append(valuesum / tdeltasum if tdeltasum > 0 else float("nan"))
         elif seriestype == "averageexcitation":

From b7c00cfcb039663a2142634172319df3bde00804 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 14:00:47 +0000
Subject: [PATCH 091/150] Speed up average ionisation plot with polars

---
 artistools/estimators/plotestimators.py | 71 ++++++++++++-------------
 1 file changed, 33 insertions(+), 38 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index e10aa1382..2d03ddd6f 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -139,45 +139,40 @@ def plot_average_ionisation_excitation(
             ionstage = at.decode_roman_numeral(paramvalue.split(" ")[1])
         ylist = []
         if seriestype == "averageionisation":
-            # TODO: replace loop with Polars groupby
+            elsymb = at.get_elsymbol(atomic_number)
+            dfselected = (
+                estimators.select(
+                    cs.starts_with(f"nnion_{elsymb}_")
+                    | cs.by_name(f"nnelement_{elsymb}")
+                    | cs.by_name("modelgridindex")
+                    | cs.by_name("timestep")
+                    | cs.by_name("xvalue")
+                )
+                .with_columns(pl.col(pl.Float32).fill_null(0.0))
+                .collect()
+                .join(
+                    pl.DataFrame({"timestep": range(len(arr_tdelta)), "tdelta": arr_tdelta}).with_columns(
+                        pl.col("timestep").cast(pl.Int32)
+                    ),
+                    on="timestep",
+                    how="left",
+                )
+            )
+            ioncols = [col for col in dfselected.columns if col.startswith(f"nnion_{elsymb}_")]
+            ioncharges = [at.decode_roman_numeral(col.removeprefix(f"nnion_{elsymb}_")) - 1 for col in ioncols]
+            dfselected = dfselected.with_columns(
+                (
+                    pl.sum_horizontal([pl.col(ioncol) * ioncharge for ioncol, ioncharge in zip(ioncols, ioncharges)])
+                    / pl.col(f"nnelement_{elsymb}")
+                ).alias(f"averageionisation_{elsymb}")
+            )
             for modelgridindex, timesteps in zip(mgilist, timestepslist):
-                valuesum = 0
-                tdeltasum = 0
-                for timestep in timesteps:
-                    estimatorstsmgi = estimators.filter(pl.col("timestep") == timestep).filter(
-                        pl.col("modelgridindex") == modelgridindex
-                    )
-                    elsymb = at.get_elsymbol(atomic_number)
-                    dfselected = (
-                        estimatorstsmgi.select(cs.starts_with(f"nnion_{elsymb}_") | cs.by_name(f"nnelement_{elsymb}"))
-                        .fill_null(0.0)
-                        .collect()
-                    )
-
-                    dfnnelement = dfselected[f"nnelement_{elsymb}"]
-                    if dfnnelement.is_empty():
-                        myval = float("NaN")
-                    else:
-                        nnelement = dfnnelement.item(0)
-                        if nnelement is None:
-                            myval = float("NaN")
-                        elif ioncols := [col for col in dfselected.columns if col.startswith(f"nnion_{elsymb}_")]:
-                            ioncharges = [
-                                at.decode_roman_numeral(col.removeprefix(f"nnion_{elsymb}_")) - 1 for col in ioncols
-                            ]
-                            myval = (
-                                dfselected.select(
-                                    pl.sum_horizontal(
-                                        [pl.col(ioncol) * ioncharge for ioncol, ioncharge in zip(ioncols, ioncharges)]
-                                    )
-                                ).item(0, 0)
-                                / nnelement
-                            )
-                        else:
-                            myval = float("NaN")
-                    valuesum += myval * arr_tdelta[timestep]
-                    tdeltasum += arr_tdelta[timestep]
-                ylist.append(valuesum / tdeltasum if tdeltasum > 0 else float("nan"))
+                dfselected_mgi = dfselected.filter(pl.col("modelgridindex") == modelgridindex)
+                avg_ionisation_timeavg = (
+                    dfselected_mgi.select(pl.col(f"averageionisation_{elsymb}") * pl.col("tdelta")).sum().item(0, 0)
+                    / dfselected_mgi["tdelta"].sum()
+                )
+                ylist.append(avg_ionisation_timeavg)
         elif seriestype == "averageexcitation":
             print("  This will be slow!")
             for modelgridindex, timesteps in zip(mgilist, timestepslist):

From f22a79b2e8ce4e31f7e0426dfd9287223b0c1c29 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 14:09:44 +0000
Subject: [PATCH 092/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 2d03ddd6f..26c31e3f6 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -166,6 +166,7 @@ def plot_average_ionisation_excitation(
                     / pl.col(f"nnelement_{elsymb}")
                 ).alias(f"averageionisation_{elsymb}")
             )
+            # TODO: for more performance, replace with a group_by expression
             for modelgridindex, timesteps in zip(mgilist, timestepslist):
                 dfselected_mgi = dfselected.filter(pl.col("modelgridindex") == modelgridindex)
                 avg_ionisation_timeavg = (
@@ -286,7 +287,7 @@ def plot_levelpop(
 
 def plot_multi_ion_series(
     ax: plt.Axes,
-    xlist: t.Sequence[int | float] | np.ndarray,
+    xvariable: str,
     seriestype: str,
     ionlist: t.Sequence[str],
     timestepslist: t.Sequence[t.Sequence[int]],
@@ -297,7 +298,6 @@ def plot_multi_ion_series(
     **plotkwargs: t.Any,
 ):
     """Plot an ion-specific property, e.g., populations."""
-    assert len(xlist) - 1 == len(mgilist) == len(timestepslist)
     # if seriestype == 'populations':
     #     ax.yaxis.set_major_locator(ticker.MultipleLocator(base=0.10))
 
@@ -382,7 +382,12 @@ def get_iontuple(ionstr):
             raise AssertionError
 
         series = estimators.group_by("xvalue").agg(pl.col(key).mean() / scalefactor).lazy().collect().sort("xvalue")
+        xlist = series["xvalue"].to_list()
         ylist = series[key].to_list()
+        if xvariable.startswith("velocity") or xvariable == "beta":
+            # make a line segment from 0 velocity
+            xlist.insert(0, 0.0)
+            ylist.insert(0, ylist[0])
 
         plotlabel = (
             ionstage
@@ -421,8 +426,6 @@ def get_iontuple(ionstr):
         # color = f'C{colorindex}'
         # or ax.step(where='pre', )
 
-        ylist.insert(0, ylist[0])
-
         xlist, ylist = at.estimators.apply_filters(xlist, ylist, args)
         if plotkwargs.get("linestyle", "solid") != "None":
             plotkwargs["dashes"] = dashes
@@ -533,6 +536,7 @@ def get_xlist(
         dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["vel_r_mid"])
         if modelmeta["dimensions"] > 1:
             args.markersonly = True
+            args.colorbyion = True
         if modelmeta["vmax_cmps"] > 0.3 * 29979245800:
             args.x = "beta"
             xvariable = "beta"
@@ -672,7 +676,7 @@ def plot_subplot(
                 seriestype, ionlist = plotitem
                 plot_multi_ion_series(
                     ax,
-                    xlist,
+                    xvariable,
                     seriestype,
                     ionlist,
                     timestepslist,

From bdcac76e459adf54961899aac510f5003c647b39 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 14:24:23 +0000
Subject: [PATCH 093/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 34 +++++++++----------------
 1 file changed, 12 insertions(+), 22 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 26c31e3f6..ad500d72c 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -290,8 +290,6 @@ def plot_multi_ion_series(
     xvariable: str,
     seriestype: str,
     ionlist: t.Sequence[str],
-    timestepslist: t.Sequence[t.Sequence[int]],
-    mgilist: t.Sequence[int | t.Sequence[int]],
     estimators: pl.LazyFrame | pl.DataFrame,
     modelpath: str | Path,
     args: argparse.Namespace,
@@ -334,16 +332,10 @@ def get_iontuple(ionstr):
                     missingions.add((atomic_number, ionstage))
 
     except FileNotFoundError:
-        print("WARNING: Could not read an ARTIS compositiondata.txt file")
-        estim_mgits = estimators.filter(pl.col("timestep") == timestepslist[0][0]).filter(
-            pl.col("modelgridindex") == mgilist[0]
-        )
+        print("WARNING: Could not read an ARTIS compositiondata.txt file to check ion availability")
         for atomic_number, ionstage in iontuplelist:
             ionstr = at.get_ionstring(atomic_number, ionstage, sep="_", style="spectral")
-            if (
-                f"nnion_{ionstr}" not in estim_mgits.columns
-                or estim_mgits.select(f"nnion_{ionstr}").lazy().collect()[f"nnion_{ionstr}"].is_null().all()
-            ):
+            if f"nnion_{ionstr}" not in estimators.columns:
                 missingions.add((atomic_number, ionstage))
 
     if missingions:
@@ -536,7 +528,6 @@ def get_xlist(
         dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["vel_r_mid"])
         if modelmeta["dimensions"] > 1:
             args.markersonly = True
-            args.colorbyion = True
         if modelmeta["vmax_cmps"] > 0.3 * 29979245800:
             args.x = "beta"
             xvariable = "beta"
@@ -675,15 +666,13 @@ def plot_subplot(
                 showlegend = True
                 seriestype, ionlist = plotitem
                 plot_multi_ion_series(
-                    ax,
-                    xvariable,
-                    seriestype,
-                    ionlist,
-                    timestepslist,
-                    mgilist,
-                    estimators,
-                    modelpath,
-                    args,
+                    ax=ax,
+                    xvariable=xvariable,
+                    seriestype=seriestype,
+                    ionlist=ionlist,
+                    estimators=estimators,
+                    modelpath=modelpath,
+                    args=args,
                     **plotkwargs,
                 )
 
@@ -747,6 +736,9 @@ def make_plot(
         plotkwargs["linestyle"] = "None"
         plotkwargs["marker"] = "."
 
+        # with no lines, line styles cannot distringuish ions
+        args.colorbyion = True
+
     # ideally, we could start filtering the columns here, but it's still faster to collect the whole thing
     allts: set[int] = set()
     for tspoint in timestepslist:
@@ -758,8 +750,6 @@ def make_plot(
 
     estimators = (
         estimators.filter(pl.col("modelgridindex").is_in(mgilist)).filter(pl.col("timestep").is_in(allts)).lazy()
-        # .collect()
-        # .lazy()
     )
 
     for ax, plotitems in zip(axes, plotlist):

From b6fa5e872cce66f2b4c61f827f52a8cc4167942e Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 14:26:04 +0000
Subject: [PATCH 094/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index ad500d72c..df8afa7ca 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -287,14 +287,14 @@ def plot_levelpop(
 
 def plot_multi_ion_series(
     ax: plt.Axes,
-    xvariable: str,
+    startfromzero: bool,
     seriestype: str,
     ionlist: t.Sequence[str],
     estimators: pl.LazyFrame | pl.DataFrame,
     modelpath: str | Path,
     args: argparse.Namespace,
     **plotkwargs: t.Any,
-):
+) -> None:
     """Plot an ion-specific property, e.g., populations."""
     # if seriestype == 'populations':
     #     ax.yaxis.set_major_locator(ticker.MultipleLocator(base=0.10))
@@ -376,7 +376,7 @@ def get_iontuple(ionstr):
         series = estimators.group_by("xvalue").agg(pl.col(key).mean() / scalefactor).lazy().collect().sort("xvalue")
         xlist = series["xvalue"].to_list()
         ylist = series[key].to_list()
-        if xvariable.startswith("velocity") or xvariable == "beta":
+        if startfromzero:
             # make a line segment from 0 velocity
             xlist.insert(0, 0.0)
             ylist.insert(0, ylist[0])
@@ -600,7 +600,7 @@ def plot_subplot(
             elif ylabel != get_ylabel(variablename):
                 sameylabel = False
                 break
-
+    startfromzero = xvariable.startswith("velocity") or xvariable == "beta"
     for plotitem in plotitems:
         if isinstance(plotitem, str):
             showlegend = seriescount > 1 or len(plotitem) > 20 or not sameylabel
@@ -667,7 +667,7 @@ def plot_subplot(
                 seriestype, ionlist = plotitem
                 plot_multi_ion_series(
                     ax=ax,
-                    xvariable=xvariable,
+                    startfromzero=startfromzero,
                     seriestype=seriestype,
                     ionlist=ionlist,
                     estimators=estimators,

From a301f12fe973a52b2094a116bbe3d0a28305bce6 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 14:30:37 +0000
Subject: [PATCH 095/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 36 ++++++++++++-------------
 1 file changed, 17 insertions(+), 19 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index df8afa7ca..f09e25019 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -449,7 +449,7 @@ def get_iontuple(ionstr):
 
 def plot_series(
     ax: plt.Axes,
-    xvariable: str,
+    startfromzero: bool,
     variablename: str,
     showlegend: bool,
     modelpath: str | Path,
@@ -485,7 +485,7 @@ def plot_series(
         # 'cooling_adiabatic': 'blue'
     }
 
-    if xvariable.startswith("velocity") or xvariable == "beta":
+    if startfromzero:
         # make a line segment from 0 velocity
         xlist.insert(0, 0.0)
         ylist.insert(0, ylist[0])
@@ -578,6 +578,7 @@ def plot_subplot(
     timestepslist: list[list[int]],
     xlist: list[float | int],
     xvariable: str,
+    startfromzero: bool,
     plotitems: list[t.Any],
     mgilist: list[int | t.Sequence[int]],
     modelpath: str | Path,
@@ -587,7 +588,6 @@ def plot_subplot(
 ):
     """Make plot from ARTIS estimators."""
     # these three lists give the x value, modelgridex, and a list of timesteps (for averaging) for each plot of the plot
-    assert len(xlist) - 1 == len(mgilist) == len(timestepslist)
     showlegend = False
     seriescount = 0
     ylabel = None
@@ -600,13 +600,13 @@ def plot_subplot(
             elif ylabel != get_ylabel(variablename):
                 sameylabel = False
                 break
-    startfromzero = xvariable.startswith("velocity") or xvariable == "beta"
+
     for plotitem in plotitems:
         if isinstance(plotitem, str):
             showlegend = seriescount > 1 or len(plotitem) > 20 or not sameylabel
             plot_series(
                 ax=ax,
-                xvariable=xvariable,
+                startfromzero=startfromzero,
                 variablename=plotitem,
                 showlegend=showlegend,
                 modelpath=modelpath,
@@ -722,12 +722,9 @@ def make_plot(
     xlist, mgilist, timestepslist, estimators = get_xlist(
         xvariable, allnonemptymgilist, estimators, timestepslist_unfiltered, modelpath, args
     )
-
-    xlist = list(
-        np.insert(xlist, 0, 0.0)
-        if (xvariable.startswith("velocity") or xvariable == "beta")
-        else np.insert(xlist, 0, xlist[0])
-    )
+    startfromzero = xvariable.startswith("velocity") or xvariable == "beta"
+    if startfromzero:
+        xlist.insert(0, 0.0)
 
     xmin = args.xmin if args.xmin >= 0 else min(xlist)
     xmax = args.xmax if args.xmax > 0 else max(xlist)
@@ -755,14 +752,15 @@ def make_plot(
     for ax, plotitems in zip(axes, plotlist):
         ax.set_xlim(left=xmin, right=xmax)
         plot_subplot(
-            ax,
-            timestepslist,
-            xlist,
-            xvariable,
-            plotitems,
-            mgilist,
-            modelpath,
-            estimators,
+            ax=ax,
+            timestepslist=timestepslist,
+            xlist=xlist,
+            xvariable=xvariable,
+            plotitems=plotitems,
+            mgilist=mgilist,
+            modelpath=modelpath,
+            estimators=estimators,
+            startfromzero=startfromzero,
             args=args,
             **plotkwargs,
         )

From 5e418896b57a5126adbdcc9761fdaa2c3a6dc578 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 14:55:29 +0000
Subject: [PATCH 096/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index f09e25019..10bb51854 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -140,6 +140,9 @@ def plot_average_ionisation_excitation(
         ylist = []
         if seriestype == "averageionisation":
             elsymb = at.get_elsymbol(atomic_number)
+            if f"nnelement_{elsymb}" not in estimators.columns:
+                msg = f"ERROR: No element data found for {paramvalue}"
+                raise ValueError(msg)
             dfselected = (
                 estimators.select(
                     cs.starts_with(f"nnion_{elsymb}_")
@@ -541,6 +544,7 @@ def get_xlist(
             dfmodel = dfmodel.filter(pl.col("vel_r_mid") / scalefactor <= args.xmax)
         else:
             dfmodel = dfmodel.filter(pl.col("vel_r_mid") <= modelmeta["vmax_cmps"])
+
         dfmodelcollect = dfmodel.select(["vel_r_mid", "modelgridindex"]).collect()
 
         xlist = (dfmodelcollect["vel_r_mid"] / scalefactor).to_list()

From 9465d4fc5b040329d645a353a3ba7e9ac5a5ef50 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 14:55:37 +0000
Subject: [PATCH 097/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 10bb51854..6ce293c36 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -726,7 +726,7 @@ def make_plot(
     xlist, mgilist, timestepslist, estimators = get_xlist(
         xvariable, allnonemptymgilist, estimators, timestepslist_unfiltered, modelpath, args
     )
-    startfromzero = xvariable.startswith("velocity") or xvariable == "beta"
+    startfromzero = (xvariable.startswith("velocity") or xvariable == "beta") and not args.markersonly
     if startfromzero:
         xlist.insert(0, 0.0)
 

From 49af4706f4c418c67140e5356788e0c00b11a3f7 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 14:58:22 +0000
Subject: [PATCH 098/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 41 ++++++++++++++++---------
 1 file changed, 27 insertions(+), 14 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 6ce293c36..d5bd40eab 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -119,6 +119,7 @@ def plot_average_ionisation_excitation(
     mgilist,
     estimators,
     modelpath,
+    startfromzero: bool,
     args=None,
     **plotkwargs,
 ):
@@ -203,7 +204,8 @@ def plot_average_ionisation_excitation(
             ylist.append(exc_ev_times_tdelta_sum / tdeltasum if tdeltasum > 0 else float("nan"))
 
         color = get_elemcolor(atomic_number=atomic_number)
-        ylist.insert(0, ylist[0])
+        if startfromzero:
+            ylist.insert(0, ylist[0])
 
         xlist, ylist = at.estimators.apply_filters(xlist, ylist, args)
 
@@ -556,19 +558,29 @@ def get_xlist(
         estimators = estimators.sort("xvalue")
         xlist = estimators.select("xvalue").collect()["xvalue"].to_list()
     else:
-        xlist = []
-        mgilist_out = []
-        timestepslist_out = []
-        for modelgridindex, timesteps in zip(allnonemptymgilist, timestepslist):
-            xvalue = at.estimators.get_averaged_estimators(modelpath, estimators, timesteps, modelgridindex, xvariable)[
-                xvariable
-            ]
-            assert isinstance(xvalue, float | int)
-            xlist.append(xvalue)
-            mgilist_out.append(modelgridindex)
-            timestepslist_out.append(timesteps)
-            if args.xmax > 0 and xvalue > args.xmax:
-                break
+        dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["vel_r_mid"])
+        # handle xvariable is in dfmodel. TODO: handle xvariable is in estimators
+        assert xvariable in dfmodel.columns
+        if modelmeta["dimensions"] > 1:
+            args.markersonly = True
+        dfmodel = dfmodel.with_columns(pl.col("inputcellid").sub(1).alias("modelgridindex"))
+        dfmodel = dfmodel.filter(pl.col("modelgridindex").is_in(allnonemptymgilist))
+        if args.readonlymgi:
+            dfmodel = dfmodel.filter(pl.col("modelgridindex").is_in(args.modelgridindex))
+        dfmodel = dfmodel.select(["modelgridindex", xvariable]).sort(by=xvariable)
+
+        if args.xmax > 0:
+            dfmodel = dfmodel.filter(pl.col(xvariable) <= args.xmax)
+        dfmodelcollect = dfmodel.select(["modelgridindex", xvariable]).collect()
+
+        xlist = dfmodelcollect[xvariable].to_list()
+        mgilist_out = dfmodelcollect["modelgridindex"].to_list()
+        timestepslist_out = timestepslist
+        estimators = estimators.filter(pl.col("modelgridindex").is_in(mgilist_out))
+        estimators = estimators.lazy().join(dfmodel.select(["modelgridindex", xvariable]).lazy(), on="modelgridindex")
+        estimators = estimators.with_columns(xvalue=pl.col(xvariable))
+        estimators = estimators.sort("xvalue")
+        xlist = estimators.select("xvalue").collect()["xvalue"].to_list()
 
     xlist, mgilist_out, timestepslist_out = zip(*sorted(zip(xlist, mgilist_out, timestepslist_out)))
 
@@ -653,6 +665,7 @@ def plot_subplot(
                     mgilist,
                     estimators,
                     modelpath,
+                    startfromzero=startfromzero,
                     args=args,
                     **plotkwargs,
                 )

From d68fba986fe87e1e6377a83b2c5c7fd3e3312c08 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 18:14:34 +0000
Subject: [PATCH 099/150] Speed up averageionisation series

---
 artistools/estimators/plotestimators.py | 22 ++++++++++++----------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index d5bd40eab..2501d7097 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -170,14 +170,16 @@ def plot_average_ionisation_excitation(
                     / pl.col(f"nnelement_{elsymb}")
                 ).alias(f"averageionisation_{elsymb}")
             )
-            # TODO: for more performance, replace with a group_by expression
-            for modelgridindex, timesteps in zip(mgilist, timestepslist):
-                dfselected_mgi = dfselected.filter(pl.col("modelgridindex") == modelgridindex)
-                avg_ionisation_timeavg = (
-                    dfselected_mgi.select(pl.col(f"averageionisation_{elsymb}") * pl.col("tdelta")).sum().item(0, 0)
-                    / dfselected_mgi["tdelta"].sum()
-                )
-                ylist.append(avg_ionisation_timeavg)
+
+            series = dfselected.group_by("xvalue").agg(pl.col(f"averageionisation_{elsymb}").mean()).lazy().collect()
+            xlist = series["xvalue"].to_list()
+            ylist = series[f"averageionisation_{elsymb}"].to_list()
+
+            if startfromzero:
+                # make a line segment from 0 velocity
+                xlist.insert(0, 0.0)
+                ylist.insert(0, ylist[0])
+
         elif seriestype == "averageexcitation":
             print("  This will be slow!")
             for modelgridindex, timesteps in zip(mgilist, timestepslist):
@@ -202,10 +204,10 @@ def plot_average_ionisation_excitation(
                 msg = f"ERROR: No excitation data found for {paramvalue}"
                 raise ValueError(msg)
             ylist.append(exc_ev_times_tdelta_sum / tdeltasum if tdeltasum > 0 else float("nan"))
+            if startfromzero:
+                ylist.insert(0, ylist[0])
 
         color = get_elemcolor(atomic_number=atomic_number)
-        if startfromzero:
-            ylist.insert(0, ylist[0])
 
         xlist, ylist = at.estimators.apply_filters(xlist, ylist, args)
 

From 6185dc3c9c9a788298d7a5e8b74cc149993d3469 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 18:34:45 +0000
Subject: [PATCH 100/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 39 ++++++++++++++++++-------
 1 file changed, 29 insertions(+), 10 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 2501d7097..882181d8d 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -151,6 +151,7 @@ def plot_average_ionisation_excitation(
                     | cs.by_name("modelgridindex")
                     | cs.by_name("timestep")
                     | cs.by_name("xvalue")
+                    | cs.by_name("plotpointid")
                 )
                 .with_columns(pl.col(pl.Float32).fill_null(0.0))
                 .collect()
@@ -171,7 +172,12 @@ def plot_average_ionisation_excitation(
                 ).alias(f"averageionisation_{elsymb}")
             )
 
-            series = dfselected.group_by("xvalue").agg(pl.col(f"averageionisation_{elsymb}").mean()).lazy().collect()
+            series = (
+                dfselected.group_by("plotpointid", maintain_order=True)
+                .agg(pl.col(f"averageionisation_{elsymb}").mean(), pl.col("xvalue").mean())
+                .lazy()
+                .collect()
+            )
             xlist = series["xvalue"].to_list()
             ylist = series[f"averageionisation_{elsymb}"].to_list()
 
@@ -380,7 +386,13 @@ def get_iontuple(ionstr):
         else:
             raise AssertionError
 
-        series = estimators.group_by("xvalue").agg(pl.col(key).mean() / scalefactor).lazy().collect().sort("xvalue")
+        series = (
+            estimators.group_by("plotpointid")
+            .agg(pl.col(key).mean() / scalefactor, pl.col("xvalue").mean())
+            .lazy()
+            .collect()
+            .sort("xvalue")
+        )
         xlist = series["xvalue"].to_list()
         ylist = series[key].to_list()
         if startfromzero:
@@ -479,7 +491,12 @@ def plot_series(
         linelabel = None
     print(f"Plotting {variablename}")
 
-    series = estimators.group_by("xvalue", maintain_order=True).agg(pl.col(variablename).mean()).lazy().collect()
+    series = (
+        estimators.group_by("plotpointid", maintain_order=True)
+        .agg(pl.col(variablename).mean(), pl.col("xvalue").mean())
+        .lazy()
+        .collect()
+    )
     ylist = series[variablename].to_list()
     xlist = series["xvalue"].to_list()
 
@@ -517,19 +534,19 @@ def get_xlist(
         mgilist_out = [mgi for mgi in allnonemptymgilist if mgi <= args.xmax] if args.xmax >= 0 else allnonemptymgilist
         xlist = list(mgilist_out)
         timestepslist_out = timestepslist
-        estimators = estimators.with_columns(xvalue=pl.col("modelgridindex"))
+        estimators = estimators.with_columns(xvalue=pl.col("modelgridindex"), plotpointid=pl.col("modelgridindex"))
     elif xvariable == "timestep":
         mgilist_out = allnonemptymgilist
         check_type(timestepslist, t.Sequence[int])
         xlist = timestepslist
         timestepslist_out = timestepslist
-        estimators = estimators.with_columns(xvalue=pl.col("timestep"))
+        estimators = estimators.with_columns(xvalue=pl.col("timestep"), plotpointid=pl.col("timestep"))
     elif xvariable == "time":
         mgilist_out = allnonemptymgilist
         timearray = at.get_timestep_times(modelpath)
         check_type(timestepslist, t.Sequence[t.Sequence[int]])
         xlist = [np.mean([timearray[ts] for ts in tslist]) for tslist in timestepslist]
-        estimators = estimators.with_columns(xvalue=pl.Series(xlist))
+        estimators = estimators.with_columns(xvalue=pl.Series(xlist), plotpointid=pl.col("timestep"))
         timestepslist_out = timestepslist
     elif xvariable in {"velocity", "beta"}:
         dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["vel_r_mid"])
@@ -556,9 +573,11 @@ def get_xlist(
         timestepslist_out = timestepslist
         estimators = estimators.filter(pl.col("modelgridindex").is_in(mgilist_out))
         estimators = estimators.lazy().join(dfmodel.select(["modelgridindex", "vel_r_mid"]).lazy(), on="modelgridindex")
-        estimators = estimators.with_columns(xvalue=(pl.col("vel_r_mid") / scalefactor))
+        estimators = estimators.with_columns(
+            xvalue=(pl.col("vel_r_mid") / scalefactor), plotpointid=pl.col("modelgridindex")
+        )
         estimators = estimators.sort("xvalue")
-        xlist = estimators.select("xvalue").collect()["xvalue"].to_list()
+        xlist = estimators.group_by(pl.col("plotpointid")).agg(pl.col("xvalue").mean()).collect()["xvalue"].to_list()
     else:
         dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["vel_r_mid"])
         # handle xvariable is in dfmodel. TODO: handle xvariable is in estimators
@@ -579,8 +598,8 @@ def get_xlist(
         mgilist_out = dfmodelcollect["modelgridindex"].to_list()
         timestepslist_out = timestepslist
         estimators = estimators.filter(pl.col("modelgridindex").is_in(mgilist_out))
-        estimators = estimators.lazy().join(dfmodel.select(["modelgridindex", xvariable]).lazy(), on="modelgridindex")
-        estimators = estimators.with_columns(xvalue=pl.col(xvariable))
+        estimators = estimators.lazy().join(dfmodel, on="modelgridindex")
+        estimators = estimators.with_columns(xvalue=pl.col(xvariable), plotpointid=pl.col("modelgridindex"))
         estimators = estimators.sort("xvalue")
         xlist = estimators.select("xvalue").collect()["xvalue"].to_list()
 

From d458bb897d67797f49bfbda9c4b1bd3df98d3906 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 18:47:20 +0000
Subject: [PATCH 101/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 882181d8d..15077ff7d 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -579,12 +579,12 @@ def get_xlist(
         estimators = estimators.sort("xvalue")
         xlist = estimators.group_by(pl.col("plotpointid")).agg(pl.col("xvalue").mean()).collect()["xvalue"].to_list()
     else:
-        dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["vel_r_mid"])
+        dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=[xvariable])
         # handle xvariable is in dfmodel. TODO: handle xvariable is in estimators
         assert xvariable in dfmodel.columns
         if modelmeta["dimensions"] > 1:
             args.markersonly = True
-        dfmodel = dfmodel.with_columns(pl.col("inputcellid").sub(1).alias("modelgridindex"))
+
         dfmodel = dfmodel.filter(pl.col("modelgridindex").is_in(allnonemptymgilist))
         if args.readonlymgi:
             dfmodel = dfmodel.filter(pl.col("modelgridindex").is_in(args.modelgridindex))

From d3ae157d01569284d4e8fa61f1051e2c4f877fbf Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 19:28:06 +0000
Subject: [PATCH 102/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 68 +++++++++++++++----------
 1 file changed, 40 insertions(+), 28 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 15077ff7d..1e3f0825a 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -527,19 +527,18 @@ def get_xlist(
     estimators: pl.LazyFrame | pl.DataFrame,
     timestepslist: t.Any,
     modelpath: str | Path,
+    groupbyxvalue: bool,
     args: t.Any,
 ) -> tuple[list[float | int], list[int | t.Sequence[int]], list[list[int]], pl.LazyFrame | pl.DataFrame]:
     xlist: t.Sequence[float | int]
     if xvariable in {"cellid", "modelgridindex"}:
         mgilist_out = [mgi for mgi in allnonemptymgilist if mgi <= args.xmax] if args.xmax >= 0 else allnonemptymgilist
         xlist = list(mgilist_out)
-        timestepslist_out = timestepslist
         estimators = estimators.with_columns(xvalue=pl.col("modelgridindex"), plotpointid=pl.col("modelgridindex"))
     elif xvariable == "timestep":
         mgilist_out = allnonemptymgilist
         check_type(timestepslist, t.Sequence[int])
         xlist = timestepslist
-        timestepslist_out = timestepslist
         estimators = estimators.with_columns(xvalue=pl.col("timestep"), plotpointid=pl.col("timestep"))
     elif xvariable == "time":
         mgilist_out = allnonemptymgilist
@@ -547,11 +546,8 @@ def get_xlist(
         check_type(timestepslist, t.Sequence[t.Sequence[int]])
         xlist = [np.mean([timearray[ts] for ts in tslist]) for tslist in timestepslist]
         estimators = estimators.with_columns(xvalue=pl.Series(xlist), plotpointid=pl.col("timestep"))
-        timestepslist_out = timestepslist
     elif xvariable in {"velocity", "beta"}:
         dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["vel_r_mid"])
-        if modelmeta["dimensions"] > 1:
-            args.markersonly = True
         if modelmeta["vmax_cmps"] > 0.3 * 29979245800:
             args.x = "beta"
             xvariable = "beta"
@@ -570,14 +566,12 @@ def get_xlist(
 
         xlist = (dfmodelcollect["vel_r_mid"] / scalefactor).to_list()
         mgilist_out = dfmodelcollect["modelgridindex"].to_list()
-        timestepslist_out = timestepslist
         estimators = estimators.filter(pl.col("modelgridindex").is_in(mgilist_out))
         estimators = estimators.lazy().join(dfmodel.select(["modelgridindex", "vel_r_mid"]).lazy(), on="modelgridindex")
         estimators = estimators.with_columns(
             xvalue=(pl.col("vel_r_mid") / scalefactor), plotpointid=pl.col("modelgridindex")
         )
         estimators = estimators.sort("xvalue")
-        xlist = estimators.group_by(pl.col("plotpointid")).agg(pl.col("xvalue").mean()).collect()["xvalue"].to_list()
     else:
         dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=[xvariable])
         # handle xvariable is in dfmodel. TODO: handle xvariable is in estimators
@@ -596,18 +590,43 @@ def get_xlist(
 
         xlist = dfmodelcollect[xvariable].to_list()
         mgilist_out = dfmodelcollect["modelgridindex"].to_list()
-        timestepslist_out = timestepslist
         estimators = estimators.filter(pl.col("modelgridindex").is_in(mgilist_out))
         estimators = estimators.lazy().join(dfmodel, on="modelgridindex")
         estimators = estimators.with_columns(xvalue=pl.col(xvariable), plotpointid=pl.col("modelgridindex"))
-        estimators = estimators.sort("xvalue")
-        xlist = estimators.select("xvalue").collect()["xvalue"].to_list()
 
-    xlist, mgilist_out, timestepslist_out = zip(*sorted(zip(xlist, mgilist_out, timestepslist_out)))
+    allts: set[int] = set()
+    for tspoint in timestepslist:
+        if isinstance(tspoint, int):
+            allts.add(tspoint)
+        else:
+            for ts in tspoint:
+                allts.add(ts)
 
-    assert len(xlist) == len(mgilist_out) == len(timestepslist_out)
+    estimators = estimators.filter(pl.col("modelgridindex").is_in(allnonemptymgilist)).filter(
+        pl.col("timestep").is_in(allts)
+    )
 
-    return list(xlist), list(mgilist_out), list(timestepslist_out), estimators
+    # single valued line plot
+    if groupbyxvalue:
+        estimators = estimators.with_columns(plotpointid=pl.col("xvalue"))
+
+    estimators = estimators.sort("plotpointid")
+    pointgroups = (
+        (
+            estimators.select(["plotpointid", "xvalue", "modelgridindex", "timestep"])
+            .group_by("plotpointid", maintain_order=True)
+            .agg(pl.col("xvalue").first(), pl.col("modelgridindex").first(), pl.col("timestep").unique())
+        )
+        .lazy()
+        .collect()
+    )
+
+    return (
+        pointgroups["xvalue"].to_list(),
+        pointgroups["modelgridindex"].to_list(),
+        pointgroups["timestep"].to_list(),
+        estimators,
+    )
 
 
 def plot_subplot(
@@ -758,7 +777,13 @@ def make_plot(
         axes[-1].set_xlabel(f"{xvariable}{at.estimators.get_units_string(xvariable)}")
 
     xlist, mgilist, timestepslist, estimators = get_xlist(
-        xvariable, allnonemptymgilist, estimators, timestepslist_unfiltered, modelpath, args
+        xvariable=xvariable,
+        allnonemptymgilist=allnonemptymgilist,
+        estimators=estimators,
+        timestepslist=timestepslist_unfiltered,
+        modelpath=modelpath,
+        groupbyxvalue=not args.markersonly,
+        args=args,
     )
     startfromzero = (xvariable.startswith("velocity") or xvariable == "beta") and not args.markersonly
     if startfromzero:
@@ -774,19 +799,6 @@ def make_plot(
         # with no lines, line styles cannot distringuish ions
         args.colorbyion = True
 
-    # ideally, we could start filtering the columns here, but it's still faster to collect the whole thing
-    allts: set[int] = set()
-    for tspoint in timestepslist:
-        if isinstance(tspoint, int):
-            allts.add(tspoint)
-        else:
-            for ts in tspoint:
-                allts.add(ts)
-
-    estimators = (
-        estimators.filter(pl.col("modelgridindex").is_in(mgilist)).filter(pl.col("timestep").is_in(allts)).lazy()
-    )
-
     for ax, plotitems in zip(axes, plotlist):
         ax.set_xlim(left=xmin, right=xmax)
         plot_subplot(
@@ -843,8 +855,8 @@ def make_plot(
         axes[0].set_title(figure_title, fontsize=8)
     # plt.suptitle(figure_title, fontsize=11, verticalalignment='top')
 
+    print(f"Saving {outfilename} ...")
     fig.savefig(outfilename)
-    print(f"Saved {outfilename}")
 
     if args.show:
         plt.show()

From 2a9eb6411b12aa83b3469b226c1fa00db88abd2a Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 19:37:05 +0000
Subject: [PATCH 103/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 63 ++++++++++---------------
 1 file changed, 24 insertions(+), 39 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 1e3f0825a..2d2c13424 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -16,7 +16,6 @@
 import numpy as np
 import polars as pl
 import polars.selectors as cs
-from typeguard import check_type
 
 import artistools as at
 
@@ -530,22 +529,32 @@ def get_xlist(
     groupbyxvalue: bool,
     args: t.Any,
 ) -> tuple[list[float | int], list[int | t.Sequence[int]], list[list[int]], pl.LazyFrame | pl.DataFrame]:
-    xlist: t.Sequence[float | int]
+    allts: set[int] = set()
+    for tspoint in timestepslist:
+        if isinstance(tspoint, int):
+            allts.add(tspoint)
+        else:
+            for ts in tspoint:
+                allts.add(ts)
+
+    estimators = estimators.filter(pl.col("modelgridindex").is_in(allnonemptymgilist)).filter(
+        pl.col("timestep").is_in(allts)
+    )
+
     if xvariable in {"cellid", "modelgridindex"}:
-        mgilist_out = [mgi for mgi in allnonemptymgilist if mgi <= args.xmax] if args.xmax >= 0 else allnonemptymgilist
-        xlist = list(mgilist_out)
         estimators = estimators.with_columns(xvalue=pl.col("modelgridindex"), plotpointid=pl.col("modelgridindex"))
     elif xvariable == "timestep":
-        mgilist_out = allnonemptymgilist
-        check_type(timestepslist, t.Sequence[int])
-        xlist = timestepslist
         estimators = estimators.with_columns(xvalue=pl.col("timestep"), plotpointid=pl.col("timestep"))
     elif xvariable == "time":
-        mgilist_out = allnonemptymgilist
         timearray = at.get_timestep_times(modelpath)
-        check_type(timestepslist, t.Sequence[t.Sequence[int]])
-        xlist = [np.mean([timearray[ts] for ts in tslist]) for tslist in timestepslist]
-        estimators = estimators.with_columns(xvalue=pl.Series(xlist), plotpointid=pl.col("timestep"))
+        estimators = estimators.lazy().join(
+            pl.DataFrame({"timestep": range(len(timearray)), "time_mid": timearray})
+            .with_columns(pl.col("timestep").cast(pl.Int32))
+            .lazy(),
+            on="timestep",
+            how="left",
+        )
+        estimators = estimators.with_columns(xvalue=pl.col("time_mid"), plotpointid=pl.col("timestep"))
     elif xvariable in {"velocity", "beta"}:
         dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["vel_r_mid"])
         if modelmeta["vmax_cmps"] > 0.3 * 29979245800:
@@ -557,16 +566,8 @@ def get_xlist(
             dfmodel = dfmodel.filter(pl.col("modelgridindex").is_in(args.modelgridindex))
         dfmodel = dfmodel.select(["modelgridindex", "vel_r_mid"]).sort(by="vel_r_mid")
         scalefactor = 1e5 if xvariable == "velocity" else 29979245800
-        if args.xmax > 0:
-            dfmodel = dfmodel.filter(pl.col("vel_r_mid") / scalefactor <= args.xmax)
-        else:
-            dfmodel = dfmodel.filter(pl.col("vel_r_mid") <= modelmeta["vmax_cmps"])
-
-        dfmodelcollect = dfmodel.select(["vel_r_mid", "modelgridindex"]).collect()
+        dfmodel = dfmodel.filter(pl.col("vel_r_mid") <= modelmeta["vmax_cmps"])
 
-        xlist = (dfmodelcollect["vel_r_mid"] / scalefactor).to_list()
-        mgilist_out = dfmodelcollect["modelgridindex"].to_list()
-        estimators = estimators.filter(pl.col("modelgridindex").is_in(mgilist_out))
         estimators = estimators.lazy().join(dfmodel.select(["modelgridindex", "vel_r_mid"]).lazy(), on="modelgridindex")
         estimators = estimators.with_columns(
             xvalue=(pl.col("vel_r_mid") / scalefactor), plotpointid=pl.col("modelgridindex")
@@ -584,32 +585,16 @@ def get_xlist(
             dfmodel = dfmodel.filter(pl.col("modelgridindex").is_in(args.modelgridindex))
         dfmodel = dfmodel.select(["modelgridindex", xvariable]).sort(by=xvariable)
 
-        if args.xmax > 0:
-            dfmodel = dfmodel.filter(pl.col(xvariable) <= args.xmax)
-        dfmodelcollect = dfmodel.select(["modelgridindex", xvariable]).collect()
-
-        xlist = dfmodelcollect[xvariable].to_list()
-        mgilist_out = dfmodelcollect["modelgridindex"].to_list()
-        estimators = estimators.filter(pl.col("modelgridindex").is_in(mgilist_out))
         estimators = estimators.lazy().join(dfmodel, on="modelgridindex")
         estimators = estimators.with_columns(xvalue=pl.col(xvariable), plotpointid=pl.col("modelgridindex"))
 
-    allts: set[int] = set()
-    for tspoint in timestepslist:
-        if isinstance(tspoint, int):
-            allts.add(tspoint)
-        else:
-            for ts in tspoint:
-                allts.add(ts)
-
-    estimators = estimators.filter(pl.col("modelgridindex").is_in(allnonemptymgilist)).filter(
-        pl.col("timestep").is_in(allts)
-    )
-
     # single valued line plot
     if groupbyxvalue:
         estimators = estimators.with_columns(plotpointid=pl.col("xvalue"))
 
+    if args.xmax > 0:
+        dfmodel = dfmodel.filter(pl.col("xvalue") <= args.xmax)
+
     estimators = estimators.sort("plotpointid")
     pointgroups = (
         (

From d15127af9d1150af4de788b3e84c44f24870bf8f Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 19:38:12 +0000
Subject: [PATCH 104/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 2d2c13424..f10565644 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -560,10 +560,10 @@ def get_xlist(
         if modelmeta["vmax_cmps"] > 0.3 * 29979245800:
             args.x = "beta"
             xvariable = "beta"
-        dfmodel = dfmodel.with_columns(pl.col("inputcellid").sub(1).alias("modelgridindex"))
-        dfmodel = dfmodel.filter(pl.col("modelgridindex").is_in(allnonemptymgilist))
+
         if args.readonlymgi:
             dfmodel = dfmodel.filter(pl.col("modelgridindex").is_in(args.modelgridindex))
+
         dfmodel = dfmodel.select(["modelgridindex", "vel_r_mid"]).sort(by="vel_r_mid")
         scalefactor = 1e5 if xvariable == "velocity" else 29979245800
         dfmodel = dfmodel.filter(pl.col("vel_r_mid") <= modelmeta["vmax_cmps"])

From a3fc75a19aa11aeb4d7967ddd27426bfe10536d7 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 21:17:26 +0000
Subject: [PATCH 105/150] Merge model data with estimators

---
 artistools/estimators/plotestimators.py  | 55 +++++++-----------------
 artistools/inputmodel/inputmodel_misc.py | 11 +++--
 2 files changed, 23 insertions(+), 43 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index f10565644..cace483fd 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -529,18 +529,6 @@ def get_xlist(
     groupbyxvalue: bool,
     args: t.Any,
 ) -> tuple[list[float | int], list[int | t.Sequence[int]], list[list[int]], pl.LazyFrame | pl.DataFrame]:
-    allts: set[int] = set()
-    for tspoint in timestepslist:
-        if isinstance(tspoint, int):
-            allts.add(tspoint)
-        else:
-            for ts in tspoint:
-                allts.add(ts)
-
-    estimators = estimators.filter(pl.col("modelgridindex").is_in(allnonemptymgilist)).filter(
-        pl.col("timestep").is_in(allts)
-    )
-
     if xvariable in {"cellid", "modelgridindex"}:
         estimators = estimators.with_columns(xvalue=pl.col("modelgridindex"), plotpointid=pl.col("modelgridindex"))
     elif xvariable == "timestep":
@@ -556,36 +544,11 @@ def get_xlist(
         )
         estimators = estimators.with_columns(xvalue=pl.col("time_mid"), plotpointid=pl.col("timestep"))
     elif xvariable in {"velocity", "beta"}:
-        dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["vel_r_mid"])
-        if modelmeta["vmax_cmps"] > 0.3 * 29979245800:
-            args.x = "beta"
-            xvariable = "beta"
-
-        if args.readonlymgi:
-            dfmodel = dfmodel.filter(pl.col("modelgridindex").is_in(args.modelgridindex))
-
-        dfmodel = dfmodel.select(["modelgridindex", "vel_r_mid"]).sort(by="vel_r_mid")
         scalefactor = 1e5 if xvariable == "velocity" else 29979245800
-        dfmodel = dfmodel.filter(pl.col("vel_r_mid") <= modelmeta["vmax_cmps"])
-
-        estimators = estimators.lazy().join(dfmodel.select(["modelgridindex", "vel_r_mid"]).lazy(), on="modelgridindex")
         estimators = estimators.with_columns(
             xvalue=(pl.col("vel_r_mid") / scalefactor), plotpointid=pl.col("modelgridindex")
         )
-        estimators = estimators.sort("xvalue")
     else:
-        dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=[xvariable])
-        # handle xvariable is in dfmodel. TODO: handle xvariable is in estimators
-        assert xvariable in dfmodel.columns
-        if modelmeta["dimensions"] > 1:
-            args.markersonly = True
-
-        dfmodel = dfmodel.filter(pl.col("modelgridindex").is_in(allnonemptymgilist))
-        if args.readonlymgi:
-            dfmodel = dfmodel.filter(pl.col("modelgridindex").is_in(args.modelgridindex))
-        dfmodel = dfmodel.select(["modelgridindex", xvariable]).sort(by=xvariable)
-
-        estimators = estimators.lazy().join(dfmodel, on="modelgridindex")
         estimators = estimators.with_columns(xvalue=pl.col(xvariable), plotpointid=pl.col("modelgridindex"))
 
     # single valued line plot
@@ -593,7 +556,7 @@ def get_xlist(
         estimators = estimators.with_columns(plotpointid=pl.col("xvalue"))
 
     if args.xmax > 0:
-        dfmodel = dfmodel.filter(pl.col("xvalue") <= args.xmax)
+        estimators = estimators.filter(pl.col("xvalue") <= args.xmax)
 
     estimators = estimators.sort("plotpointid")
     pointgroups = (
@@ -1073,7 +1036,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         # ["Te", "TR"],
         [["averageionisation", ["Sr"]]],
         # [["averageexcitation", ["Fe II", "Fe III"]]],
-        [["populations", ["Sr90", "Sr91", "Sr92", "Sr93", "Sr94"]]],
+        [["populations", ["Sr90", "Sr91", "Sr92", "Sr94"]]],
         [["populations", ["Sr I", "Sr II", "Sr III"]]],
         # [['populations', ['He I', 'He II', 'He III']]],
         # [['populations', ['C I', 'C II', 'C III', 'C IV', 'C V']]],
@@ -1166,6 +1129,16 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         if not args.x:
             args.x = "velocity"
 
+        dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["ALL"])
+        if args.x == "velocity" and modelmeta["vmax_cmps"] > 0.3 * 29979245800:
+            args.x = "beta"
+
+        dfmodel = dfmodel.filter(pl.col("vel_r_mid") <= modelmeta["vmax_cmps"])
+        estimators = estimators.join(dfmodel, on="modelgridindex")
+
+        if args.readonlymgi:
+            estimators = estimators.filter(pl.col("modelgridindex").is_in(args.modelgridindex))
+
         if args.classicartis:
             modeldata, _ = at.inputmodel.get_modeldata(modelpath)
             allnonemptymgilist = [
@@ -1180,6 +1153,10 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         else:
             allnonemptymgilist = [mgi for mgi, assocpropcells in assoc_cells.items() if assocpropcells]
 
+        estimators = estimators.filter(pl.col("modelgridindex").is_in(allnonemptymgilist)).filter(
+            pl.col("timestep").is_in(timesteps_included)
+        )
+
         if args.multiplot:
             pdf_list = []
             modelpath_list = []
diff --git a/artistools/inputmodel/inputmodel_misc.py b/artistools/inputmodel/inputmodel_misc.py
index 22d379f1f..067e8d0d8 100644
--- a/artistools/inputmodel/inputmodel_misc.py
+++ b/artistools/inputmodel/inputmodel_misc.py
@@ -308,6 +308,7 @@ def get_modeldata_polars(
         - modelpath: either a path to model.txt file, or a folder containing model.txt
         - get_elemabundances: also read elemental abundances (abundances.txt) and
             merge with the output DataFrame
+        - derived_cols: list of derived columns to add to the model data, or "ALL" to add all possible derived columns
 
     return dfmodel, modelmeta
         - dfmodel: a pandas DataFrame with a row for each model grid cell
@@ -489,6 +490,7 @@ def add_derived_cols_to_modeldata(
     original_cols = dfmodel.columns
 
     t_model_init_seconds = modelmeta["t_model_init_days"] * 86400.0
+    keep_all = "ALL" in derived_cols
 
     dimensions = modelmeta["dimensions"]
     match dimensions:
@@ -644,17 +646,18 @@ def add_derived_cols_to_modeldata(
     dfmodel = dfmodel.with_columns([(pl.col("rho") * pl.col("volume")).alias("mass_g")])
 
     if unknown_cols := [
-        col for col in derived_cols if col not in dfmodel.columns and col not in {"pos_min", "pos_max"}
+        col for col in derived_cols if col not in dfmodel.columns and col not in {"pos_min", "pos_max", "ALL"}
     ]:
         print(f"WARNING: Unknown derived columns: {unknown_cols}")
 
-    if "pos_min" in derived_cols:
+    if "pos_min" in derived_cols or keep_all:
         derived_cols.extend([f"pos_{ax}_min" for ax in axes])
 
-    if "pos_max" in derived_cols:
+    if "pos_max" in derived_cols or keep_all:
         derived_cols.extend([f"pos_{ax}_max" for ax in axes])
 
-    dfmodel = dfmodel.drop([col for col in dfmodel.columns if col not in original_cols and col not in derived_cols])
+    if not keep_all:
+        dfmodel = dfmodel.drop([col for col in dfmodel.columns if col not in original_cols and col not in derived_cols])
 
     if "angle_bin" in derived_cols:
         assert modelpath is not None

From 8b6b41f1a0393b1629107175d43d967bedaed143 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 21:52:04 +0000
Subject: [PATCH 106/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 96 ++++++++++++++++++-------
 1 file changed, 72 insertions(+), 24 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index cace483fd..6f2e7155c 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -50,8 +50,18 @@ def get_ylabel(variable):
     return ""
 
 
-def plot_init_abundances(ax, xlist, specieslist, mgilist, modelpath, seriestype, args=None, **plotkwargs):
-    assert len(xlist) - 1 == len(mgilist)
+def plot_init_abundances(
+    ax: plt.Axes,
+    xlist: list[float],
+    specieslist: list[str],
+    mgilist: t.Sequence[float],
+    modelpath: Path,
+    seriestype: str,
+    startfromzero: bool,
+    args: argparse.Namespace,
+    **plotkwargs,
+) -> None:
+    assert len(xlist) == len(mgilist)
 
     if seriestype == "initabundances":
         mergemodelabundata, _ = at.inputmodel.get_modeldata(modelpath, get_elemabundances=True)
@@ -60,6 +70,10 @@ def plot_init_abundances(ax, xlist, specieslist, mgilist, modelpath, seriestype,
     else:
         raise AssertionError
 
+    if startfromzero:
+        xlist = xlist.copy()
+        xlist.insert(0, 0.0)
+
     for speciesstr in specieslist:
         splitvariablename = speciesstr.split("_")
         elsymbol = splitvariablename[0].strip("0123456789")
@@ -97,12 +111,13 @@ def plot_init_abundances(ax, xlist, specieslist, mgilist, modelpath, seriestype,
                 yvalue = mergemodelabundata.loc[modelgridindex][f"{valuetype}{elsymbol}"]
             ylist.append(yvalue)
 
-        ylist.insert(0, ylist[0])
-        # or ax.step(where='pre', )
         color = get_elemcolor(atomic_number=atomic_number)
 
         xlist, ylist = at.estimators.apply_filters(xlist, ylist, args)
 
+        if startfromzero:
+            ylist.insert(0, ylist[0])
+
         ax.plot(xlist, ylist, linewidth=1.5, label=linelabel, linestyle=linestyle, color=color, **plotkwargs)
 
         # if args.yscale == 'log':
@@ -129,6 +144,10 @@ def plot_average_ionisation_excitation(
     else:
         raise ValueError
 
+    if startfromzero:
+        xlist = xlist.copy()
+        xlist.insert(0, 0.0)
+
     arr_tdelta = at.get_timestep_times(modelpath, loc="delta")
     for paramvalue in params:
         print(f"Plotting {seriestype} {paramvalue}")
@@ -178,12 +197,10 @@ def plot_average_ionisation_excitation(
                 .collect()
             )
             xlist = series["xvalue"].to_list()
-            ylist = series[f"averageionisation_{elsymb}"].to_list()
-
             if startfromzero:
-                # make a line segment from 0 velocity
                 xlist.insert(0, 0.0)
-                ylist.insert(0, ylist[0])
+
+            ylist = series[f"averageionisation_{elsymb}"].to_list()
 
         elif seriestype == "averageexcitation":
             print("  This will be slow!")
@@ -209,13 +226,14 @@ def plot_average_ionisation_excitation(
                 msg = f"ERROR: No excitation data found for {paramvalue}"
                 raise ValueError(msg)
             ylist.append(exc_ev_times_tdelta_sum / tdeltasum if tdeltasum > 0 else float("nan"))
-            if startfromzero:
-                ylist.insert(0, ylist[0])
 
         color = get_elemcolor(atomic_number=atomic_number)
 
         xlist, ylist = at.estimators.apply_filters(xlist, ylist, args)
-
+        if startfromzero:
+            ylist.insert(0, ylist[0])
+        print(f"  Plotting {seriestype} {paramvalue}")
+        print(xlist, ylist, startfromzero)
         ax.plot(xlist, ylist, label=paramvalue, color=color, **plotkwargs)
 
 
@@ -528,7 +546,7 @@ def get_xlist(
     modelpath: str | Path,
     groupbyxvalue: bool,
     args: t.Any,
-) -> tuple[list[float | int], list[int | t.Sequence[int]], list[list[int]], pl.LazyFrame | pl.DataFrame]:
+) -> tuple[list[float | int], list[int], list[list[int]], pl.LazyFrame | pl.DataFrame]:
     if xvariable in {"cellid", "modelgridindex"}:
         estimators = estimators.with_columns(xvalue=pl.col("modelgridindex"), plotpointid=pl.col("modelgridindex"))
     elif xvariable == "timestep":
@@ -544,9 +562,10 @@ def get_xlist(
         )
         estimators = estimators.with_columns(xvalue=pl.col("time_mid"), plotpointid=pl.col("timestep"))
     elif xvariable in {"velocity", "beta"}:
+        velcolumn = "vel_r_mid"
         scalefactor = 1e5 if xvariable == "velocity" else 29979245800
         estimators = estimators.with_columns(
-            xvalue=(pl.col("vel_r_mid") / scalefactor), plotpointid=pl.col("modelgridindex")
+            xvalue=(pl.col(velcolumn) / scalefactor), plotpointid=pl.col("modelgridindex")
         )
     else:
         estimators = estimators.with_columns(xvalue=pl.col(xvariable), plotpointid=pl.col("modelgridindex"))
@@ -584,7 +603,7 @@ def plot_subplot(
     xvariable: str,
     startfromzero: bool,
     plotitems: list[t.Any],
-    mgilist: list[int | t.Sequence[int]],
+    mgilist: list[int],
     modelpath: str | Path,
     estimators: pl.LazyFrame | pl.DataFrame,
     args: argparse.Namespace,
@@ -626,7 +645,16 @@ def plot_subplot(
 
             if seriestype in {"initabundances", "initmasses"}:
                 showlegend = True
-                plot_init_abundances(ax, xlist, params, mgilist, modelpath, seriestype, args=args)
+                plot_init_abundances(
+                    ax=ax,
+                    xlist=xlist,
+                    specieslist=params,
+                    mgilist=mgilist,
+                    modelpath=Path(modelpath),
+                    seriestype=seriestype,
+                    startfromzero=startfromzero,
+                    args=args,
+                )
 
             elif seriestype == "levelpopulation" or seriestype.startswith("levelpopulation_"):
                 showlegend = True
@@ -734,8 +762,6 @@ def make_plot(
         args=args,
     )
     startfromzero = (xvariable.startswith("velocity") or xvariable == "beta") and not args.markersonly
-    if startfromzero:
-        xlist.insert(0, 0.0)
 
     xmin = args.xmin if args.xmin >= 0 else min(xlist)
     xmax = args.xmax if args.xmax > 0 else max(xlist)
@@ -1023,7 +1049,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
     )
 
     plotlist = args.plotlist or [
-        # [['initabundances', ['Fe', 'Ni_stable', 'Ni_56']]],
+        [["initabundances", ["Fe", "Ni_stable", "Ni_56"]]],
         # ['heating_dep', 'heating_coll', 'heating_bf', 'heating_ff',
         #  ['_yscale', 'linear']],
         # ['cooling_adiabatic', 'cooling_coll', 'cooling_fb', 'cooling_ff',
@@ -1118,11 +1144,19 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         if not args.x:
             args.x = "time"
         mgilist = [args.modelgridindex] * len(timesteps_included)
-        timesteplist_unfiltered = [[ts] for ts in timesteps_included]
+        timestepslist_unfiltered = [[ts] for ts in timesteps_included]
         if not assoc_cells.get(args.modelgridindex, []):
             msg = f"cell {args.modelgridindex} is empty. no estimators available"
             raise ValueError(msg)
-        make_plot(modelpath, timesteplist_unfiltered, mgilist, estimators, args.x, plotlist, args)
+        make_plot(
+            modelpath=modelpath,
+            timestepslist_unfiltered=timestepslist_unfiltered,
+            allnonemptymgilist=mgilist,
+            estimators=estimators,
+            xvariable=args.x,
+            plotlist=plotlist,
+            args=args,
+        )
     else:
         # plot a range of cells in a time snapshot showing internal structure
 
@@ -1161,9 +1195,15 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             pdf_list = []
             modelpath_list = []
             for timestep in range(timestepmin, timestepmax + 1):
-                timesteplist_unfiltered = [[timestep]] * len(allnonemptymgilist)
+                timestepslist_unfiltered = [[timestep]] * len(allnonemptymgilist)
                 outfilename = make_plot(
-                    modelpath, timesteplist_unfiltered, allnonemptymgilist, estimators, args.x, plotlist, args
+                    modelpath=modelpath,
+                    timestepslist_unfiltered=timestepslist_unfiltered,
+                    allnonemptymgilist=allnonemptymgilist,
+                    estimators=estimators,
+                    xvariable=args.x,
+                    plotlist=plotlist,
+                    args=args,
                 )
 
                 if "/" in outfilename:
@@ -1176,8 +1216,16 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
                 at.join_pdf_files(pdf_list, modelpath_list)
 
         else:
-            timesteplist_unfiltered = [timesteps_included] * len(allnonemptymgilist)
-            make_plot(modelpath, timesteplist_unfiltered, allnonemptymgilist, estimators, args.x, plotlist, args)
+            timestepslist_unfiltered = [timesteps_included] * len(allnonemptymgilist)
+            make_plot(
+                modelpath=modelpath,
+                timestepslist_unfiltered=timestepslist_unfiltered,
+                allnonemptymgilist=allnonemptymgilist,
+                estimators=estimators,
+                xvariable=args.x,
+                plotlist=plotlist,
+                args=args,
+            )
 
 
 if __name__ == "__main__":

From 1fbc1a357dae51b0ad8549034e97778f3353f69d Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 21:57:14 +0000
Subject: [PATCH 107/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 6f2e7155c..f00518827 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -233,7 +233,7 @@ def plot_average_ionisation_excitation(
         if startfromzero:
             ylist.insert(0, ylist[0])
         print(f"  Plotting {seriestype} {paramvalue}")
-        print(xlist, ylist, startfromzero)
+
         ax.plot(xlist, ylist, label=paramvalue, color=color, **plotkwargs)
 
 

From 459f80903e6a7512f8c7536ec498ce9352dd9e95 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 22:01:43 +0000
Subject: [PATCH 108/150] Update estimators.py

---
 artistools/estimators/estimators.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index da559e51b..bd02e176f 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -29,6 +29,9 @@ def get_variableunits(key: str | None = None) -> str | dict[str, str]:
         "Te": "K",
         "TJ": "K",
         "nne": "e^-/cm3",
+        "nniso": "cm$^{-3}$",
+        "nnion": "cm$^{-3}$",
+        "nnelement": "cm$^{-3}$",
         "heating": "erg/s/cm3",
         "heating_dep/total_dep": "Ratio",
         "cooling": "erg/s/cm3",
@@ -340,7 +343,7 @@ def read_estimators_polars(
     if match_timestep is not None:
         pldflazy = pldflazy.filter(pl.col("timestep").is_in(match_timestep))
 
-    return pldflazy
+    return pldflazy.fill_null(0)
 
 
 def read_estimators(

From a4b9fc0f96fcc6b606f312d5265481aafc58b8f7 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 22:02:00 +0000
Subject: [PATCH 109/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index f00518827..789d5598d 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -517,7 +517,7 @@ def plot_series(
     ylist = series[variablename].to_list()
     xlist = series["xvalue"].to_list()
 
-    if math.log10(max(ylist) / min(ylist)) > 2 or min(ylist) == 0:
+    if min(ylist) == 0 or math.log10(max(ylist) / min(ylist)) > 2:
         ax.set_yscale("log")
 
     dictcolors = {
@@ -1049,7 +1049,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
     )
 
     plotlist = args.plotlist or [
-        [["initabundances", ["Fe", "Ni_stable", "Ni_56"]]],
+        # [["initabundances", ["Fe", "Ni_stable", "Ni_56"]]],
         # ['heating_dep', 'heating_coll', 'heating_bf', 'heating_ff',
         #  ['_yscale', 'linear']],
         # ['cooling_adiabatic', 'cooling_coll', 'cooling_fb', 'cooling_ff',

From 40413208ad747c9fe2e4b0fb0d2a3036b9a99a9f Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 22:14:32 +0000
Subject: [PATCH 110/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 789d5598d..843e6ac80 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -184,12 +184,15 @@ def plot_average_ionisation_excitation(
             ioncols = [col for col in dfselected.columns if col.startswith(f"nnion_{elsymb}_")]
             ioncharges = [at.decode_roman_numeral(col.removeprefix(f"nnion_{elsymb}_")) - 1 for col in ioncols]
             dfselected = dfselected.with_columns(
-                (
+                pl.when(pl.col(f"nnelement_{elsymb}") > 0.0)
+                .then(
                     pl.sum_horizontal([pl.col(ioncol) * ioncharge for ioncol, ioncharge in zip(ioncols, ioncharges)])
                     / pl.col(f"nnelement_{elsymb}")
-                ).alias(f"averageionisation_{elsymb}")
+                )
+                .otherwise(pl.lit(None))
+                .alias(f"averageionisation_{elsymb}")
             )
-
+            dfselected.drop_nulls(f"averageionisation_{elsymb}")
             series = (
                 dfselected.group_by("plotpointid", maintain_order=True)
                 .agg(pl.col(f"averageionisation_{elsymb}").mean(), pl.col("xvalue").mean())

From ceaa0a14ca64309e78a7a83394c46c62f60f5b76 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 22:20:06 +0000
Subject: [PATCH 111/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 843e6ac80..50436a9ac 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -181,18 +181,16 @@ def plot_average_ionisation_excitation(
                     how="left",
                 )
             )
+            dfselected = dfselected.filter(pl.col(f"nnelement_{elsymb}") > 0.0)
+
             ioncols = [col for col in dfselected.columns if col.startswith(f"nnion_{elsymb}_")]
             ioncharges = [at.decode_roman_numeral(col.removeprefix(f"nnion_{elsymb}_")) - 1 for col in ioncols]
             dfselected = dfselected.with_columns(
-                pl.when(pl.col(f"nnelement_{elsymb}") > 0.0)
-                .then(
+                (
                     pl.sum_horizontal([pl.col(ioncol) * ioncharge for ioncol, ioncharge in zip(ioncols, ioncharges)])
                     / pl.col(f"nnelement_{elsymb}")
-                )
-                .otherwise(pl.lit(None))
-                .alias(f"averageionisation_{elsymb}")
+                ).alias(f"averageionisation_{elsymb}")
             )
-            dfselected.drop_nulls(f"averageionisation_{elsymb}")
             series = (
                 dfselected.group_by("plotpointid", maintain_order=True)
                 .agg(pl.col(f"averageionisation_{elsymb}").mean(), pl.col("xvalue").mean())

From 35b532d636d5f231a5b8f550ece940f9cccf3a9b Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sat, 11 Nov 2023 22:25:51 +0000
Subject: [PATCH 112/150] Add type hints to plot_average_ionisation_excitation

---
 artistools/estimators/estimators.py     |  2 +-
 artistools/estimators/plotestimators.py | 27 ++++++++++++++-----------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index bd02e176f..928649811 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -415,7 +415,7 @@ def get_averaged_estimators(
 
 
 def get_averageexcitation(
-    modelpath: Path, modelgridindex: int, timestep: int, atomic_number: int, ionstage: int, T_exc: float
+    modelpath: Path | str, modelgridindex: int, timestep: int, atomic_number: int, ionstage: int, T_exc: float
 ) -> float | None:
     dfnltepops = at.nltepops.read_files(modelpath, modelgridindex=modelgridindex, timestep=timestep)
     if dfnltepops.empty:
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 50436a9ac..655da5f13 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -125,14 +125,14 @@ def plot_init_abundances(
 
 
 def plot_average_ionisation_excitation(
-    ax,
-    xlist,
-    seriestype,
-    params,
-    timestepslist,
-    mgilist,
-    estimators,
-    modelpath,
+    ax: plt.Axes,
+    xlist: list[float],
+    seriestype: str,
+    params: t.Sequence[str],
+    timestepslist: t.Sequence[t.Sequence[int]],
+    mgilist: t.Sequence[int],
+    estimators: pl.LazyFrame,
+    modelpath: Path | str,
     startfromzero: bool,
     args=None,
     **plotkwargs,
@@ -185,18 +185,21 @@ def plot_average_ionisation_excitation(
 
             ioncols = [col for col in dfselected.columns if col.startswith(f"nnion_{elsymb}_")]
             ioncharges = [at.decode_roman_numeral(col.removeprefix(f"nnion_{elsymb}_")) - 1 for col in ioncols]
+
             dfselected = dfselected.with_columns(
                 (
                     pl.sum_horizontal([pl.col(ioncol) * ioncharge for ioncol, ioncharge in zip(ioncols, ioncharges)])
                     / pl.col(f"nnelement_{elsymb}")
                 ).alias(f"averageionisation_{elsymb}")
             )
+
             series = (
                 dfselected.group_by("plotpointid", maintain_order=True)
                 .agg(pl.col(f"averageionisation_{elsymb}").mean(), pl.col("xvalue").mean())
                 .lazy()
                 .collect()
             )
+
             xlist = series["xvalue"].to_list()
             if startfromzero:
                 xlist.insert(0, 0.0)
@@ -542,12 +545,12 @@ def plot_series(
 def get_xlist(
     xvariable: str,
     allnonemptymgilist: t.Sequence[int],
-    estimators: pl.LazyFrame | pl.DataFrame,
+    estimators: pl.LazyFrame,
     timestepslist: t.Any,
     modelpath: str | Path,
     groupbyxvalue: bool,
     args: t.Any,
-) -> tuple[list[float | int], list[int], list[list[int]], pl.LazyFrame | pl.DataFrame]:
+) -> tuple[list[float | int], list[int], list[list[int]], pl.LazyFrame]:
     if xvariable in {"cellid", "modelgridindex"}:
         estimators = estimators.with_columns(xvalue=pl.col("modelgridindex"), plotpointid=pl.col("modelgridindex"))
     elif xvariable == "timestep":
@@ -606,7 +609,7 @@ def plot_subplot(
     plotitems: list[t.Any],
     mgilist: list[int],
     modelpath: str | Path,
-    estimators: pl.LazyFrame | pl.DataFrame,
+    estimators: pl.LazyFrame,
     args: argparse.Namespace,
     **plotkwargs: t.Any,
 ):
@@ -729,7 +732,7 @@ def make_plot(
     modelpath: Path | str,
     timestepslist_unfiltered: list[list[int]],
     allnonemptymgilist: list[int],
-    estimators: pl.LazyFrame | pl.DataFrame,
+    estimators: pl.LazyFrame,
     xvariable: str,
     plotlist,
     args: t.Any,

From b8e5d7857b72f186bcbb4d6d9c76090d9da234b7 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 04:30:54 +0000
Subject: [PATCH 113/150] Update estimators.py

---
 artistools/estimators/estimators.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 928649811..420dcb6fc 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -9,6 +9,7 @@
 import math
 import multiprocessing
 import sys
+import time
 import typing as t
 from collections import namedtuple
 from pathlib import Path
@@ -243,8 +244,11 @@ def get_rankbatch_parquetfile(
     modelpath: Path,
     folderpath: Path,
     batch_mpiranks: t.Sequence[int],
+    batchindex: int,
 ) -> Path:
-    parquetfilepath = folderpath / f"estimators_{batch_mpiranks[0]:04d}_{batch_mpiranks[-1]:04d}.out.parquet.tmp"
+    parquetfilepath = (
+        folderpath / f"estimbatch{batchindex:02d}_{batch_mpiranks[0]:04d}_{batch_mpiranks[-1]:04d}.out.parquet.tmp"
+    )
 
     if not parquetfilepath.exists():
         print(f"{parquetfilepath.relative_to(modelpath.parent)} does not exist")
@@ -259,6 +263,8 @@ def get_rankbatch_parquetfile(
             f"  reading {len(list(estfilepaths))} estimator files from {folderpath.relative_to(Path(folderpath).parent)}"
         )
 
+        time_start = time.perf_counter()
+
         pldf_group = None
         with multiprocessing.get_context("spawn").Pool(processes=at.get_config()["num_processes"]) as pool:
             for pldf_file in pool.imap(read_estimators_from_file, estfilepaths):
@@ -271,9 +277,11 @@ def get_rankbatch_parquetfile(
             pool.join()
             pool.terminate()
 
+        print(f"    took {time.perf_counter() - time_start:.1f} s")
+
         assert pldf_group is not None
         print(f"  writing {parquetfilepath.relative_to(modelpath.parent)}")
-        pldf_group.write_parquet(parquetfilepath, compression="zstd")
+        pldf_group.write_parquet(parquetfilepath, compression="zstd", statistics=True, compression_level=8)
 
     print(f"Scanning {parquetfilepath.relative_to(modelpath.parent)}")
 
@@ -329,9 +337,9 @@ def read_estimators_polars(
     runfolders = at.get_runfolders(modelpath, timesteps=match_timestep)
 
     parquetfiles = (
-        get_rankbatch_parquetfile(modelpath, runfolder, mpiranks)
+        get_rankbatch_parquetfile(modelpath, runfolder, mpiranks, batchindex=batchindex)
         for runfolder in runfolders
-        for mpiranks in mpirank_groups
+        for batchindex, mpiranks in enumerate(mpirank_groups)
     )
 
     pldflazy = pl.concat([pl.scan_parquet(pfile) for pfile in parquetfiles], how="diagonal_relaxed")

From abffbf2f7cd1240e7627e6ebf949a2a6817c97dd Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 04:32:32 +0000
Subject: [PATCH 114/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 62 ++++++++++++-------------
 1 file changed, 31 insertions(+), 31 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 655da5f13..2f8ada2cc 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -137,10 +137,10 @@ def plot_average_ionisation_excitation(
     args=None,
     **plotkwargs,
 ):
-    if seriestype == "averageionisation":
-        ax.set_ylabel("Average ion charge")
-    elif seriestype == "averageexcitation":
+    if seriestype == "averageexcitation":
         ax.set_ylabel("Average excitation [eV]")
+    elif seriestype == "averageionisation":
+        ax.set_ylabel("Average ion charge")
     else:
         raise ValueError
 
@@ -157,7 +157,32 @@ def plot_average_ionisation_excitation(
             atomic_number = at.get_atomic_number(paramvalue.split(" ")[0])
             ionstage = at.decode_roman_numeral(paramvalue.split(" ")[1])
         ylist = []
-        if seriestype == "averageionisation":
+        if seriestype == "averageexcitation":
+            print("  This will be slow!")
+            for modelgridindex, timesteps in zip(mgilist, timestepslist):
+                exc_ev_times_tdelta_sum = 0.0
+                tdeltasum = 0.0
+                for timestep in timesteps:
+                    T_exc = (
+                        estimators.filter(pl.col("timestep") == timestep)
+                        .filter(pl.col("modelgridindex") == modelgridindex)
+                        .select("Te")
+                        .lazy()
+                        .collect()
+                        .item(0, 0)
+                    )
+                    exc_ev = at.estimators.get_averageexcitation(
+                        modelpath, modelgridindex, timestep, atomic_number, ionstage, T_exc
+                    )
+                    if exc_ev is not None:
+                        exc_ev_times_tdelta_sum += exc_ev * arr_tdelta[timestep]
+                        tdeltasum += arr_tdelta[timestep]
+            if tdeltasum == 0.0:
+                msg = f"ERROR: No excitation data found for {paramvalue}"
+                raise ValueError(msg)
+            ylist.append(exc_ev_times_tdelta_sum / tdeltasum if tdeltasum > 0 else float("nan"))
+
+        elif seriestype == "averageionisation":
             elsymb = at.get_elsymbol(atomic_number)
             if f"nnelement_{elsymb}" not in estimators.columns:
                 msg = f"ERROR: No element data found for {paramvalue}"
@@ -206,31 +231,6 @@ def plot_average_ionisation_excitation(
 
             ylist = series[f"averageionisation_{elsymb}"].to_list()
 
-        elif seriestype == "averageexcitation":
-            print("  This will be slow!")
-            for modelgridindex, timesteps in zip(mgilist, timestepslist):
-                exc_ev_times_tdelta_sum = 0.0
-                tdeltasum = 0.0
-                for timestep in timesteps:
-                    T_exc = (
-                        estimators.filter(pl.col("timestep") == timestep)
-                        .filter(pl.col("modelgridindex") == modelgridindex)
-                        .select("Te")
-                        .lazy()
-                        .collect()
-                        .item(0, 0)
-                    )
-                    exc_ev = at.estimators.get_averageexcitation(
-                        modelpath, modelgridindex, timestep, atomic_number, ionstage, T_exc
-                    )
-                    if exc_ev is not None:
-                        exc_ev_times_tdelta_sum += exc_ev * arr_tdelta[timestep]
-                        tdeltasum += arr_tdelta[timestep]
-            if tdeltasum == 0.0:
-                msg = f"ERROR: No excitation data found for {paramvalue}"
-                raise ValueError(msg)
-            ylist.append(exc_ev_times_tdelta_sum / tdeltasum if tdeltasum > 0 else float("nan"))
-
         color = get_elemcolor(atomic_number=atomic_number)
 
         xlist, ylist = at.estimators.apply_filters(xlist, ylist, args)
@@ -1061,12 +1061,12 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         # [['initmasses', ['Ni_56', 'He', 'C', 'Mg']]],
         # ['heating_gamma/gamma_dep'],
         ["nne", ["_ymin", 1e5], ["_ymax", 1e10]],
-        ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 16000]],
+        ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 10000]],
         # ["Te"],
         # ["Te", "TR"],
         [["averageionisation", ["Sr"]]],
         # [["averageexcitation", ["Fe II", "Fe III"]]],
-        [["populations", ["Sr90", "Sr91", "Sr92", "Sr94"]]],
+        # [["populations", ["Sr90", "Sr91", "Sr92", "Sr94"]]],
         [["populations", ["Sr I", "Sr II", "Sr III"]]],
         # [['populations', ['He I', 'He II', 'He III']]],
         # [['populations', ['C I', 'C II', 'C III', 'C IV', 'C V']]],

From e9b1e5e0b53a336c1ae629116afed66eda4b8f18 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 04:37:38 +0000
Subject: [PATCH 115/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 2f8ada2cc..206668dcf 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -408,7 +408,7 @@ def get_iontuple(ionstr):
             raise AssertionError
 
         series = (
-            estimators.group_by("plotpointid")
+            estimators.group_by("plotpointid", maintain_order=True)
             .agg(pl.col(key).mean() / scalefactor, pl.col("xvalue").mean())
             .lazy()
             .collect()

From a025ab2259bd2f95a54aae7d420952d9cacab9c0 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 13:11:22 +0000
Subject: [PATCH 116/150] Rename read_estimators_polars to scan_estimators

---
 artistools/estimators/__init__.py       | 2 +-
 artistools/estimators/estimators.py     | 4 ++--
 artistools/estimators/plotestimators.py | 2 +-
 artistools/plotspherical.py             | 2 +-
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/artistools/estimators/__init__.py b/artistools/estimators/__init__.py
index 9e0baae80..df2b7d88c 100644
--- a/artistools/estimators/__init__.py
+++ b/artistools/estimators/__init__.py
@@ -14,6 +14,6 @@
 from artistools.estimators.estimators import get_variableunits
 from artistools.estimators.estimators import read_estimators
 from artistools.estimators.estimators import read_estimators_from_file
-from artistools.estimators.estimators import read_estimators_polars
+from artistools.estimators.estimators import scan_estimators
 from artistools.estimators.plotestimators import addargs
 from artistools.estimators.plotestimators import main as plot
diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 420dcb6fc..58c70e596 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -288,7 +288,7 @@ def get_rankbatch_parquetfile(
     return parquetfilepath
 
 
-def read_estimators_polars(
+def scan_estimators(
     modelpath: Path | str = Path(),
     modelgridindex: None | int | t.Sequence[int] = None,
     timestep: None | int | t.Sequence[int] = None,
@@ -362,7 +362,7 @@ def read_estimators(
 ) -> dict[tuple[int, int], dict[str, t.Any]]:
     if isinstance(keys, str):
         keys = {keys}
-    pldflazy = read_estimators_polars(modelpath, modelgridindex, timestep)
+    pldflazy = scan_estimators(modelpath, modelgridindex, timestep)
     estimators: dict[tuple[int, int], dict[str, t.Any]] = {}
     for estimtsmgi in pldflazy.collect().iter_rows(named=True):
         ts, mgi = estimtsmgi["timestep"], estimtsmgi["modelgridindex"]
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 206668dcf..4e6d3e583 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -1116,7 +1116,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             ]
         ).lazy()
     else:
-        estimators = at.estimators.read_estimators_polars(
+        estimators = at.estimators.scan_estimators(
             modelpath=modelpath,
             modelgridindex=args.modelgridindex,
             timestep=tuple(timesteps_included),
diff --git a/artistools/plotspherical.py b/artistools/plotspherical.py
index 10f46adc7..e87fac5b3 100755
--- a/artistools/plotspherical.py
+++ b/artistools/plotspherical.py
@@ -120,7 +120,7 @@ def plot_spherical(
         )
 
         df_estimators = (
-            at.estimators.read_estimators_polars(modelpath=modelpath)
+            at.estimators.scan_estimators(modelpath=modelpath)
             .select(["timestep", "modelgridindex", "TR"])
             .drop_nulls()
             .rename({"timestep": "em_timestep", "modelgridindex": "em_modelgridindex", "TR": "em_TR"})

From 7d460dab9c7747d7ca9da916efb2dd22983a6fba Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 13:16:05 +0000
Subject: [PATCH 117/150] Print parquet file size to stdout

---
 artistools/estimators/estimators.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 58c70e596..826e27db2 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -283,7 +283,8 @@ def get_rankbatch_parquetfile(
         print(f"  writing {parquetfilepath.relative_to(modelpath.parent)}")
         pldf_group.write_parquet(parquetfilepath, compression="zstd", statistics=True, compression_level=8)
 
-    print(f"Scanning {parquetfilepath.relative_to(modelpath.parent)}")
+    filesize = parquetfilepath.stat().st_size / 1024 / 1024
+    print(f"Scanning {parquetfilepath.relative_to(modelpath.parent)} ({filesize:.2f} MiB)")
 
     return parquetfilepath
 

From 2b6ba8ce91140edeb1b96e80fecaea177f13f4dc Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 13:38:37 +0000
Subject: [PATCH 118/150] Simplify get_runfolder_timesteps

---
 artistools/misc.py | 25 +++++++++----------------
 1 file changed, 9 insertions(+), 16 deletions(-)

diff --git a/artistools/misc.py b/artistools/misc.py
index d3d63658e..3f68bd0dc 100644
--- a/artistools/misc.py
+++ b/artistools/misc.py
@@ -1,5 +1,4 @@
 import argparse
-import contextlib
 import gzip
 import io
 import math
@@ -1123,21 +1122,15 @@ def get_inputparams(modelpath: Path) -> dict[str, t.Any]:
 @lru_cache(maxsize=16)
 def get_runfolder_timesteps(folderpath: Path | str) -> tuple[int, ...]:
     """Get the set of timesteps covered by the output files in an ARTIS run folder."""
-    folder_timesteps = set()
-    with contextlib.suppress(FileNotFoundError), zopen(Path(folderpath, "estimators_0000.out")) as estfile:
-        restart_timestep = -1
-        for line in estfile:
-            if line.startswith("timestep "):
-                timestep = int(line.split()[1])
-
-                if restart_timestep < 0 and timestep != 0 and 0 not in folder_timesteps:
-                    # the first timestep of a restarted run is duplicate and should be ignored
-                    restart_timestep = timestep
-
-                if timestep != restart_timestep:
-                    folder_timesteps.add(timestep)
-
-    return tuple(folder_timesteps)
+    estimfiles = sorted(Path(folderpath).glob("estimators_*.out*"))
+    if not estimfiles:
+        return ()
+
+    with zopen(estimfiles[0]) as estfile:
+        timesteps_contained = sorted({int(line.split()[1]) for line in estfile if line.startswith("timestep ")})
+        # the first timestep of a restarted run is duplicate and should be ignored
+        restart_timestep = None if 0 in timesteps_contained else timesteps_contained[0]
+        return tuple(ts for ts in timesteps_contained if ts != restart_timestep)
 
 
 def get_runfolders(

From 027a2b8615c34a2149de0553386107c5cc657d23 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 13:21:00 +0000
Subject: [PATCH 119/150] Remove emptycell and lognne variables from estimators
 (all nonempty and lognne is redundant)

---
 artistools/codecomparison.py            |  3 ---
 artistools/estimators/estimators.py     |  9 +++------
 artistools/estimators/plotestimators.py |  6 +-----
 artistools/gsinetwork.py                |  2 +-
 artistools/linefluxes.py                |  4 +---
 artistools/nltepops/plotnltepops.py     | 15 +++++----------
 artistools/spectra/plotspectra.py       |  5 ++---
 artistools/transitions.py               |  3 ---
 artistools/writecomparisondata.py       |  6 +-----
 9 files changed, 14 insertions(+), 39 deletions(-)

diff --git a/artistools/codecomparison.py b/artistools/codecomparison.py
index 96717232c..1668d0d71 100644
--- a/artistools/codecomparison.py
+++ b/artistools/codecomparison.py
@@ -86,9 +86,6 @@ def read_reference_estimators(
 
                 key = (cur_timestep, cur_modelgridindex)
 
-                if key not in estimators:
-                    estimators[key] = {"emptycell": False}
-
                 estimators[key]["vel_mid"] = float(row[0])
                 estimators[key]["Te"] = float(row[1])
                 estimators[key]["rho"] = float(row[2])
diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 826e27db2..310d7c256 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -6,7 +6,6 @@
 import argparse
 import contextlib
 import itertools
-import math
 import multiprocessing
 import sys
 import time
@@ -132,6 +131,7 @@ def read_estimators_from_file(
     estimblock: dict[str, t.Any] = {}
     timestep: int | None = None
     modelgridindex: int | None = None
+    emptycell: bool | None = None
     with at.zopen(estfilepath) as estimfile:
         for line in estimfile:
             row: list[str] = line.split()
@@ -140,18 +140,16 @@ def read_estimators_from_file(
 
             if row[0] == "timestep":
                 # yield the previous block before starting a new one
-                if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
+                if timestep is not None and modelgridindex is not None and not emptycell:
                     estimblocklist.append(estimblock | {"timestep": timestep, "modelgridindex": modelgridindex})
 
                 timestep = int(row[1])
                 modelgridindex = int(row[3])
                 emptycell = row[4] == "EMPTYCELL"
-                estimblock = {"emptycell": emptycell}
                 if not emptycell:
                     # will be TR, Te, W, TJ, nne
                     for variablename, value in zip(row[4::2], row[5::2]):
                         estimblock[variablename] = float(value)
-                    estimblock["lognne"] = math.log10(estimblock["nne"]) if estimblock["nne"] > 0 else float("-inf")
 
             elif row[1].startswith("Z="):
                 variablename = row[0]
@@ -216,7 +214,7 @@ def read_estimators_from_file(
                     estimblock[f"cooling_{coolingtype}"] = float(value)
 
     # reached the end of file
-    if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
+    if timestep is not None and modelgridindex is not None and not emptycell:
         estimblocklist.append(estimblock | {"timestep": timestep, "modelgridindex": modelgridindex})
 
     return pl.DataFrame(estimblocklist).with_columns(
@@ -327,7 +325,6 @@ def scan_estimators(
                     **estimvals,
                 }
                 for (ts, mgi), estimvals in estimators.items()
-                if not estimvals.get("emptycell", True)
             ]
         ).lazy()
 
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 4e6d3e583..fc3854050 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -864,10 +864,7 @@ def plot_recombrates(modelpath, estimators, atomic_number, ionstage_list, **plot
         list_rrc = []
         list_rrc2 = []
         for dicttimestepmodelgrid in estimators.values():
-            if (
-                not dicttimestepmodelgrid["emptycell"]
-                and (atomic_number, ionstage) in dicttimestepmodelgrid["RRC_LTE_Nahar"]
-            ):
+            if (atomic_number, ionstage) in dicttimestepmodelgrid["RRC_LTE_Nahar"]:
                 listT_e.append(dicttimestepmodelgrid["Te"])
                 list_rrc.append(dicttimestepmodelgrid["RRC_LTE_Nahar"][(atomic_number, ionstage)])
                 list_rrc2.append(dicttimestepmodelgrid["Alpha_R"][(atomic_number, ionstage)])
@@ -1112,7 +1109,6 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
                     **estimvals,
                 }
                 for (ts, mgi), estimvals in estimatorsdict.items()
-                if not estimvals.get("emptycell", True)
             ]
         ).lazy()
     else:
diff --git a/artistools/gsinetwork.py b/artistools/gsinetwork.py
index c56be8dc8..bd8643f2f 100755
--- a/artistools/gsinetwork.py
+++ b/artistools/gsinetwork.py
@@ -487,7 +487,7 @@ def plot_qdot_abund_modelcells(
         for nts, mgi in sorted(estimators.keys()):
             if nts in partiallycomplete_timesteps:
                 continue
-            if mgi not in mgiplotlist and not get_global_Ye or estimators[(nts, mgi)]["emptycell"]:
+            if mgi not in mgiplotlist and not get_global_Ye:
                 continue
 
             if first_mgi is None:
diff --git a/artistools/linefluxes.py b/artistools/linefluxes.py
index 0f40bc942..649368690 100755
--- a/artistools/linefluxes.py
+++ b/artistools/linefluxes.py
@@ -459,9 +459,7 @@ def get_packets_with_emission_conditions(
 
     modeldata, _ = at.inputmodel.get_modeldata(modelpath)
     ts = at.get_timestep_of_timedays(modelpath, tend)
-    allnonemptymgilist = [
-        modelgridindex for modelgridindex in modeldata.index if not estimators[(ts, modelgridindex)]["emptycell"]
-    ]
+    allnonemptymgilist = list({modelgridindex for estimts, modelgridindex in estimators if estimts == ts})
 
     # model_tmids = at.get_timestep_times(modelpath, loc='mid')
     # arr_velocity_mid = tuple(list([(float(v1) + float(v2)) * 0.5 for v1, v2 in zip(
diff --git a/artistools/nltepops/plotnltepops.py b/artistools/nltepops/plotnltepops.py
index fc193630e..8b5c5da4c 100755
--- a/artistools/nltepops/plotnltepops.py
+++ b/artistools/nltepops/plotnltepops.py
@@ -572,16 +572,11 @@ def make_plot(modelpath, atomic_number, ionstages_displayed, mgilist, timestep,
         print(f"Z={atomic_number} {elsymbol}")
 
         if estimators:
-            if not estimators[(timestep, modelgridindex)]["emptycell"]:
-                T_e = estimators[(timestep, modelgridindex)]["Te"]
-                T_R = estimators[(timestep, modelgridindex)]["TR"]
-                W = estimators[(timestep, modelgridindex)]["W"]
-                nne = estimators[(timestep, modelgridindex)]["nne"]
-                print(f"nne = {nne} cm^-3, T_e = {T_e} K, T_R = {T_R} K, W = {W}")
-            else:
-                print(f"ERROR: cell {modelgridindex} is empty. Setting T_e = T_R = {args.exc_temperature} K")
-                T_e = args.exc_temperature
-                T_R = args.exc_temperature
+            T_e = estimators[(timestep, modelgridindex)]["Te"]
+            T_R = estimators[(timestep, modelgridindex)]["TR"]
+            W = estimators[(timestep, modelgridindex)]["W"]
+            nne = estimators[(timestep, modelgridindex)]["nne"]
+            print(f"nne = {nne} cm^-3, T_e = {T_e} K, T_R = {T_R} K, W = {W}")
         else:
             print("WARNING: No estimator data. Setting T_e = T_R =  6000 K")
             T_e = args.exc_temperature
diff --git a/artistools/spectra/plotspectra.py b/artistools/spectra/plotspectra.py
index 234b79e96..38752ab36 100755
--- a/artistools/spectra/plotspectra.py
+++ b/artistools/spectra/plotspectra.py
@@ -809,9 +809,8 @@ def make_contrib_plot(axes: t.Iterable[plt.Axes], modelpath: Path, densityplotyv
 
     else:
         estimators = at.estimators.read_estimators(modelpath=modelpath)
-        allnonemptymgilist = [
-            modelgridindex for modelgridindex in modeldata.index if not estimators[(0, modelgridindex)]["emptycell"]
-        ]
+        allnonemptymgilist = list({modelgridindex for ts, modelgridindex in estimators})
+
     assert estimators is not None
     packetsfiles = at.packets.get_packetsfilepaths(modelpath, args.maxpacketfiles)
     assert args.timemin is not None
diff --git a/artistools/transitions.py b/artistools/transitions.py
index f12a8b1a5..77c60ac70 100755
--- a/artistools/transitions.py
+++ b/artistools/transitions.py
@@ -258,9 +258,6 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             sys.exit(1)
 
         estimators = estimators_all[(timestep, modelgridindex)]
-        if estimators["emptycell"]:
-            print(f"ERROR: cell {modelgridindex} is marked as empty")
-            sys.exit(1)
 
     # also calculate wavelengths outside the plot range to include lines whose
     # edges pass through the plot range
diff --git a/artistools/writecomparisondata.py b/artistools/writecomparisondata.py
index 283b46fa0..1f624162a 100755
--- a/artistools/writecomparisondata.py
+++ b/artistools/writecomparisondata.py
@@ -214,11 +214,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
 
         modeldata, t_model_init_days, _ = at.inputmodel.get_modeldata_tuple(modelpath)
         estimators = at.estimators.read_estimators(modelpath=modelpath)
-        allnonemptymgilist = [
-            modelgridindex
-            for modelgridindex in modeldata.index
-            if not estimators[(selected_timesteps[0], modelgridindex)]["emptycell"]
-        ]
+        allnonemptymgilist = list({modelgridindex for ts, modelgridindex in estimators if ts == selected_timesteps[0]})
 
         try:
             write_lbol_edep(

From 2c2f6d2daad2593269e6b52d5d16c0323da0b237 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 14:01:54 +0000
Subject: [PATCH 120/150] Update estimators.py

---
 artistools/estimators/estimators.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 310d7c256..57a26b9bd 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -131,7 +131,6 @@ def read_estimators_from_file(
     estimblock: dict[str, t.Any] = {}
     timestep: int | None = None
     modelgridindex: int | None = None
-    emptycell: bool | None = None
     with at.zopen(estfilepath) as estimfile:
         for line in estimfile:
             row: list[str] = line.split()
@@ -140,12 +139,13 @@ def read_estimators_from_file(
 
             if row[0] == "timestep":
                 # yield the previous block before starting a new one
-                if timestep is not None and modelgridindex is not None and not emptycell:
+                if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
                     estimblocklist.append(estimblock | {"timestep": timestep, "modelgridindex": modelgridindex})
 
                 timestep = int(row[1])
                 modelgridindex = int(row[3])
                 emptycell = row[4] == "EMPTYCELL"
+                estimblock = {"emptycell": emptycell}
                 if not emptycell:
                     # will be TR, Te, W, TJ, nne
                     for variablename, value in zip(row[4::2], row[5::2]):
@@ -214,7 +214,7 @@ def read_estimators_from_file(
                     estimblock[f"cooling_{coolingtype}"] = float(value)
 
     # reached the end of file
-    if timestep is not None and modelgridindex is not None and not emptycell:
+    if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
         estimblocklist.append(estimblock | {"timestep": timestep, "modelgridindex": modelgridindex})
 
     return pl.DataFrame(estimblocklist).with_columns(
@@ -325,6 +325,7 @@ def scan_estimators(
                     **estimvals,
                 }
                 for (ts, mgi), estimvals in estimators.items()
+                if not estimvals.get("emptycell", True)
             ]
         ).lazy()
 

From 386ad1771d93d50cddc3924bb87f25be8f73b6e6 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 21:11:09 +0000
Subject: [PATCH 121/150] Revert "Update estimators.py"

This reverts commit 2c2f6d2daad2593269e6b52d5d16c0323da0b237.
---
 artistools/estimators/estimators.py | 19 ++++++++-----------
 1 file changed, 8 insertions(+), 11 deletions(-)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 57a26b9bd..7a98a7a3e 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -129,8 +129,6 @@ def read_estimators_from_file(
 
     estimblocklist: list[dict[str, t.Any]] = []
     estimblock: dict[str, t.Any] = {}
-    timestep: int | None = None
-    modelgridindex: int | None = None
     with at.zopen(estfilepath) as estimfile:
         for line in estimfile:
             row: list[str] = line.split()
@@ -139,15 +137,15 @@ def read_estimators_from_file(
 
             if row[0] == "timestep":
                 # yield the previous block before starting a new one
-                if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
-                    estimblocklist.append(estimblock | {"timestep": timestep, "modelgridindex": modelgridindex})
+                if estimblock:
+                    estimblocklist.append(estimblock)
 
-                timestep = int(row[1])
-                modelgridindex = int(row[3])
                 emptycell = row[4] == "EMPTYCELL"
-                estimblock = {"emptycell": emptycell}
-                if not emptycell:
+                if emptycell:
+                    estimblock = {}
+                else:
                     # will be TR, Te, W, TJ, nne
+                    estimblock = {"timestep": int(row[1]), "modelgridindex": int(row[3])}
                     for variablename, value in zip(row[4::2], row[5::2]):
                         estimblock[variablename] = float(value)
 
@@ -214,8 +212,8 @@ def read_estimators_from_file(
                     estimblock[f"cooling_{coolingtype}"] = float(value)
 
     # reached the end of file
-    if timestep is not None and modelgridindex is not None and (not estimblock.get("emptycell", True)):
-        estimblocklist.append(estimblock | {"timestep": timestep, "modelgridindex": modelgridindex})
+    if estimblock:
+        estimblocklist.append(estimblock)
 
     return pl.DataFrame(estimblocklist).with_columns(
         pl.col(pl.Int64).cast(pl.Int32), pl.col(pl.Float64).cast(pl.Float32)
@@ -325,7 +323,6 @@ def scan_estimators(
                     **estimvals,
                 }
                 for (ts, mgi), estimvals in estimators.items()
-                if not estimvals.get("emptycell", True)
             ]
         ).lazy()
 

From 99e9c75dc81134b571801074d3db51a9ff9d6ce6 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 21:32:27 +0000
Subject: [PATCH 122/150] Fix multiplot

---
 artistools/estimators/plotestimators.py | 20 ++++++--------------
 1 file changed, 6 insertions(+), 14 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index fc3854050..511b23d19 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -551,6 +551,8 @@ def get_xlist(
     groupbyxvalue: bool,
     args: t.Any,
 ) -> tuple[list[float | int], list[int], list[list[int]], pl.LazyFrame]:
+    estimators = estimators.filter(pl.col("timestep").is_in([ts for tssublist in timestepslist for ts in tssublist]))
+
     if xvariable in {"cellid", "modelgridindex"}:
         estimators = estimators.with_columns(xvalue=pl.col("modelgridindex"), plotpointid=pl.col("modelgridindex"))
     elif xvariable == "timestep":
@@ -739,6 +741,7 @@ def make_plot(
     **plotkwargs: t.Any,
 ):
     modelname = at.get_model_name(modelpath)
+
     fig, axes = plt.subplots(
         nrows=len(plotlist),
         ncols=1,
@@ -804,23 +807,12 @@ def make_plot(
 
     else:
         timeavg = (args.timemin + args.timemax) / 2.0
-        if args.multiplot and not args.classicartis:
-            assert isinstance(timestepslist[0], list)
-            tdays = (
-                estimators.filter(pl.col("timestep") == timestepslist[0][0])
-                .filter(pl.col("modelgridindex") == mgilist[0])
-                .select("tdays")
-                .lazy()
-                .collect()
-                .item(0)
-            )
-            figure_title = f"{modelname}\nTimestep {timestepslist[0][0]} ({tdays:.2f}d)"
-        elif args.multiplot:
-            assert isinstance(timestepslist[0], int)
-            timedays = float(at.get_timestep_time(modelpath, timestepslist[0]))
+        if args.multiplot:
+            timedays = float(at.get_timestep_time(modelpath, timestepslist[0][0]))
             figure_title = f"{modelname}\nTimestep {timestepslist[0][0]} ({timedays:.2f}d)"
         else:
             figure_title = f"{modelname}\nTimestep {timestepslist[0][0]} ({timeavg:.2f}d)"
+        print(f"Plotting {figure_title.replace('\n', ' ')}")
 
         defaultoutputfile = Path("plotestimators_ts{timestep:02d}_{timeavg:.2f}d.pdf")
         if Path(args.outputfile).is_dir():

From e59a4f4f96312391d7b9c410895593481e304981 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 21:44:53 +0000
Subject: [PATCH 123/150] Refactor join_pdf_files

---
 artistools/estimators/plotestimators.py | 15 +++++----------
 artistools/misc.py                      | 12 ++++++------
 artistools/radfield.py                  |  8 ++------
 3 files changed, 13 insertions(+), 22 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 511b23d19..78d52de76 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -812,7 +812,7 @@ def make_plot(
             figure_title = f"{modelname}\nTimestep {timestepslist[0][0]} ({timedays:.2f}d)"
         else:
             figure_title = f"{modelname}\nTimestep {timestepslist[0][0]} ({timeavg:.2f}d)"
-        print(f"Plotting {figure_title.replace('\n', ' ')}")
+        print("Plotting ", figure_title.replace("\n", " "))
 
         defaultoutputfile = Path("plotestimators_ts{timestep:02d}_{timeavg:.2f}d.pdf")
         if Path(args.outputfile).is_dir():
@@ -1184,8 +1184,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         )
 
         if args.multiplot:
-            pdf_list = []
-            modelpath_list = []
+            pdf_files = []
             for timestep in range(timestepmin, timestepmax + 1):
                 timestepslist_unfiltered = [[timestep]] * len(allnonemptymgilist)
                 outfilename = make_plot(
@@ -1198,14 +1197,10 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
                     args=args,
                 )
 
-                if "/" in outfilename:
-                    outfilename = outfilename.split("/")[1]
+                pdf_files.append(outfilename)
 
-                pdf_list.append(outfilename)
-                modelpath_list.append(modelpath)
-
-            if len(pdf_list) > 1:
-                at.join_pdf_files(pdf_list, modelpath_list)
+            if len(pdf_files) > 1:
+                at.join_pdf_files(pdf_files)
 
         else:
             timestepslist_unfiltered = [timesteps_included] * len(allnonemptymgilist)
diff --git a/artistools/misc.py b/artistools/misc.py
index 3f68bd0dc..1a4824703 100644
--- a/artistools/misc.py
+++ b/artistools/misc.py
@@ -925,18 +925,18 @@ def savgolfilterfunc(ylist: t.Any) -> t.Any:
     return filterfunc
 
 
-def join_pdf_files(pdf_list: list[str], modelpath_list: list[Path]) -> None:
+def join_pdf_files(pdf_files: list[str]) -> None:
     """Merge a list of PDF files into a single PDF file."""
     from PyPDF2 import PdfMerger
 
     merger = PdfMerger()
 
-    for pdf, modelpath in zip(pdf_list, modelpath_list):
-        fullpath = firstexisting([pdf], folder=modelpath)
-        merger.append(fullpath.open("rb"))
-        fullpath.unlink()
+    for pdfpath in pdf_files:
+        with Path(pdfpath).open("rb") as pdffile:
+            merger.append(pdffile)
+        Path(pdfpath).unlink()
 
-    resultfilename = f'{pdf_list[0].split(".")[0]}-{pdf_list[-1].split(".")[0]}'
+    resultfilename = f'{pdf_files[0].split(".")[0]}-{pdf_files[-1].split(".")[0]}'
     with Path(f"{resultfilename}.pdf").open("wb") as resultfile:
         merger.write(resultfile)
 
diff --git a/artistools/radfield.py b/artistools/radfield.py
index a15b0ffe7..63f2bc4d7 100755
--- a/artistools/radfield.py
+++ b/artistools/radfield.py
@@ -1008,7 +1008,6 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
     modelpath = args.modelpath
 
     pdf_list = []
-    modelpath_list = []
     modelgridindexlist = []
 
     if args.velocity >= 0.0:
@@ -1042,7 +1041,6 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
                     normalised=args.normalised,
                 ):
                     pdf_list.append(outputfile)
-                    modelpath_list.append(args.modelpath)
         elif args.xaxis == "timestep":
             outputfile = args.outputfile.format(modelgridindex=modelgridindex)
             plot_timeevolution(modelpath, outputfile, modelgridindex, args)
@@ -1051,10 +1049,8 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             raise AssertionError
 
     if len(pdf_list) > 1:
-        print(pdf_list, modelpath_list)
-        at.join_pdf_files(pdf_list, modelpath_list)
-        return
-    return
+        print(pdf_list)
+        at.join_pdf_files(pdf_list)
 
 
 if __name__ == "__main__":

From c259239dae7f153238aaad4835735e4d9e6f0863 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 21:50:20 +0000
Subject: [PATCH 124/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 32 +++++++++----------------
 1 file changed, 11 insertions(+), 21 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 78d52de76..46eb3e7b8 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -1183,28 +1183,13 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             pl.col("timestep").is_in(timesteps_included)
         )
 
-        if args.multiplot:
-            pdf_files = []
-            for timestep in range(timestepmin, timestepmax + 1):
-                timestepslist_unfiltered = [[timestep]] * len(allnonemptymgilist)
-                outfilename = make_plot(
-                    modelpath=modelpath,
-                    timestepslist_unfiltered=timestepslist_unfiltered,
-                    allnonemptymgilist=allnonemptymgilist,
-                    estimators=estimators,
-                    xvariable=args.x,
-                    plotlist=plotlist,
-                    args=args,
-                )
-
-                pdf_files.append(outfilename)
-
-            if len(pdf_files) > 1:
-                at.join_pdf_files(pdf_files)
-
-        else:
+        frames_timesteps_included = (
+            [[ts] for ts in range(timestepmin, timestepmax + 1)] if args.multiplot else [timesteps_included]
+        )
+        pdf_files = []
+        for timesteps_included in frames_timesteps_included:
             timestepslist_unfiltered = [timesteps_included] * len(allnonemptymgilist)
-            make_plot(
+            outfilename = make_plot(
                 modelpath=modelpath,
                 timestepslist_unfiltered=timestepslist_unfiltered,
                 allnonemptymgilist=allnonemptymgilist,
@@ -1214,6 +1199,11 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
                 args=args,
             )
 
+            pdf_files.append(outfilename)
+
+        if len(pdf_files) > 1:
+            at.join_pdf_files(pdf_files)
+
 
 if __name__ == "__main__":
     main()

From 6e39a7b3bd48e6812cbdb63e11d0f3aad62d3676 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 21:58:56 +0000
Subject: [PATCH 125/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 25 ++++++++++++++-----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 46eb3e7b8..a17a1d031 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -717,17 +717,18 @@ def plot_subplot(
 
     ax.tick_params(right=True)
     if showlegend and not args.nolegend:
-        if plotitems[0][0] == "populations" and args.yscale == "log":
-            ax.legend(
-                loc="best", handlelength=2, ncol=math.ceil(len(plotitems[0][1]) / 2.0), frameon=False, numpoints=1
-            )
-        else:
-            ax.legend(
-                loc="best",
-                handlelength=2,
-                frameon=False,
-                numpoints=1,
-            )  # prop={'size': 9})
+        ax.legend(
+            loc="best",
+            handlelength=2,
+            frameon=False,
+            numpoints=1,
+            **(
+                {"ncol": math.ceil(len(plotitems[0][1]) / 2.0)}
+                if (plotitems[0][0] == "populations" and args.yscale == "log")
+                else {}
+            ),
+            markerscale=10,
+        )
 
 
 def make_plot(
@@ -776,6 +777,8 @@ def make_plot(
     if args.markersonly:
         plotkwargs["linestyle"] = "None"
         plotkwargs["marker"] = "."
+        plotkwargs["markersize"] = 1
+        plotkwargs["alpha"] = 0.5
 
         # with no lines, line styles cannot distringuish ions
         args.colorbyion = True

From f1da2acc5bdbf4c1944dacb71100fd13752a4340 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 22:07:18 +0000
Subject: [PATCH 126/150] Update inputmodel_misc.py

---
 artistools/inputmodel/inputmodel_misc.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/artistools/inputmodel/inputmodel_misc.py b/artistools/inputmodel/inputmodel_misc.py
index 067e8d0d8..5b3289900 100644
--- a/artistools/inputmodel/inputmodel_misc.py
+++ b/artistools/inputmodel/inputmodel_misc.py
@@ -645,6 +645,10 @@ def add_derived_cols_to_modeldata(
 
     dfmodel = dfmodel.with_columns([(pl.col("rho") * pl.col("volume")).alias("mass_g")])
 
+    dfmodel = dfmodel.with_columns(
+        [(pl.col(colname) / 29979245800.0).alias(f"{colname}_on_c") for colname in dfmodel.columns]
+    )
+
     if unknown_cols := [
         col for col in derived_cols if col not in dfmodel.columns and col not in {"pos_min", "pos_max", "ALL"}
     ]:

From b1bcf2babd6c5afe379ae85bf65bdb4eaa1516bb Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 22:18:41 +0000
Subject: [PATCH 127/150] Add formatted names for cell velocities

---
 artistools/estimators/__init__.py       | 2 +-
 artistools/estimators/estimators.py     | 8 ++++++--
 artistools/estimators/plotestimators.py | 8 +++++---
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/artistools/estimators/__init__.py b/artistools/estimators/__init__.py
index df2b7d88c..33bd77b35 100644
--- a/artistools/estimators/__init__.py
+++ b/artistools/estimators/__init__.py
@@ -6,12 +6,12 @@
 from artistools.estimators.estimators import apply_filters
 from artistools.estimators.estimators import get_averaged_estimators
 from artistools.estimators.estimators import get_averageexcitation
-from artistools.estimators.estimators import get_dictlabelreplacements
 from artistools.estimators.estimators import get_ionrecombrates_fromfile
 from artistools.estimators.estimators import get_partiallycompletetimesteps
 from artistools.estimators.estimators import get_units_string
 from artistools.estimators.estimators import get_variablelongunits
 from artistools.estimators.estimators import get_variableunits
+from artistools.estimators.estimators import get_varname_formatted
 from artistools.estimators.estimators import read_estimators
 from artistools.estimators.estimators import read_estimators_from_file
 from artistools.estimators.estimators import scan_estimators
diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index 7a98a7a3e..fc1f87541 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -38,6 +38,8 @@ def get_variableunits(key: str | None = None) -> str | dict[str, str]:
         "velocity": "km/s",
         "beta": "v/c",
         "vel_r_max_kmps": "km/s",
+        **{f"vel_{ax}_mid": "cm/s" for ax in ["x", "y", "z", "r", "rcyl"]},
+        **{f"vel_{ax}_mid_on_c": "c" for ax in ["x", "y", "z", "r", "rcyl"]},
     }
     return variableunits[key] if key else variableunits
 
@@ -52,8 +54,8 @@ def get_variablelongunits(key: str | None = None) -> str | dict[str, str]:
     return variablelongunits[key] if key else variablelongunits
 
 
-def get_dictlabelreplacements() -> dict[str, str]:
-    return {
+def get_varname_formatted(varname: str) -> str:
+    replacements = {
         "nne": r"n$_{\rm e}$",
         "lognne": r"Log n$_{\rm e}$",
         "Te": r"T$_{\rm e}$",
@@ -62,7 +64,9 @@ def get_dictlabelreplacements() -> dict[str, str]:
         "gamma_NT": r"$\Gamma_{\rm non-thermal}$ [s$^{-1}$]",
         "gamma_R_bfest": r"$\Gamma_{\rm phot}$ [s$^{-1}$]",
         "heating_dep/total_dep": "Heating fraction",
+        **{f"vel_{ax}_mid_on_c": f"$v_{{{ax}}}$" for ax in ["x", "y", "z", "r", "rcyl"]},
     }
+    return replacements.get(varname, varname)
 
 
 def apply_filters(
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index a17a1d031..f692d5b6c 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -476,7 +476,7 @@ def get_iontuple(ionstr):
         else:
             raise AssertionError
     else:
-        ax.set_ylabel(at.estimators.get_dictlabelreplacements().get(seriestype, seriestype))
+        ax.set_ylabel(at.estimators.get_varname_formatted(seriestype))
 
     if plotted_something:
         ax.set_yscale(args.yscale)
@@ -499,7 +499,7 @@ def plot_series(
     **plotkwargs: t.Any,
 ) -> None:
     """Plot something like Te or TR."""
-    formattedvariablename = at.estimators.get_dictlabelreplacements().get(variablename, variablename)
+    formattedvariablename = at.estimators.get_varname_formatted(variablename)
     serieslabel = f"{formattedvariablename}"
     units_string = at.estimators.get_units_string(variablename)
 
@@ -758,7 +758,9 @@ def make_plot(
 
     # ax.xaxis.set_minor_locator(ticker.MultipleLocator(base=5))
     if not args.hidexlabel:
-        axes[-1].set_xlabel(f"{xvariable}{at.estimators.get_units_string(xvariable)}")
+        axes[-1].set_xlabel(
+            f"{at.estimators.get_varname_formatted(xvariable)}{at.estimators.get_units_string(xvariable)}"
+        )
 
     xlist, mgilist, timestepslist, estimators = get_xlist(
         xvariable=xvariable,

From 4987df14c1bf6a81f79dc03ee57988062f676f4c Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 22:23:20 +0000
Subject: [PATCH 128/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index f692d5b6c..c63fd01ed 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -1055,7 +1055,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         # [['initmasses', ['Ni_56', 'He', 'C', 'Mg']]],
         # ['heating_gamma/gamma_dep'],
         ["nne", ["_ymin", 1e5], ["_ymax", 1e10]],
-        ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 10000]],
+        ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 15000]],
         # ["Te"],
         # ["Te", "TR"],
         [["averageionisation", ["Sr"]]],

From 501b6116df36c5f0be739e032cdebcc867ee8a02 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 22:58:27 +0000
Subject: [PATCH 129/150] Add --makegif option to plotestimators

---
 artistools/estimators/plotestimators.py | 33 +++++++++++++++++++------
 artistools/plotspherical.py             |  2 +-
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index c63fd01ed..fc72c83c1 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -210,6 +210,7 @@ def plot_average_ionisation_excitation(
 
             ioncols = [col for col in dfselected.columns if col.startswith(f"nnion_{elsymb}_")]
             ioncharges = [at.decode_roman_numeral(col.removeprefix(f"nnion_{elsymb}_")) - 1 for col in ioncols]
+            ax.set_ylim(0.0, max(ioncharges) + 0.1)
 
             dfselected = dfselected.with_columns(
                 (
@@ -482,9 +483,12 @@ def get_iontuple(ionstr):
         ax.set_yscale(args.yscale)
         if args.yscale == "log":
             ymin, ymax = ax.get_ylim()
+            ymin = max(ymin, ymax / 1e10)
+            ax.set_ylim(bottom=ymin)
+            # make space for the legend
             new_ymax = ymax * 10 ** (0.3 * math.log10(ymax / ymin))
             if ymin > 0 and new_ymax > ymin and np.isfinite(new_ymax):
-                ax.set_ylim(ymin, new_ymax)
+                ax.set_ylim(top=new_ymax)
 
 
 def plot_series(
@@ -718,7 +722,7 @@ def plot_subplot(
     ax.tick_params(right=True)
     if showlegend and not args.nolegend:
         ax.legend(
-            loc="best",
+            loc="upper right",
             handlelength=2,
             frameon=False,
             numpoints=1,
@@ -830,8 +834,10 @@ def make_plot(
         axes[0].set_title(figure_title, fontsize=8)
     # plt.suptitle(figure_title, fontsize=11, verticalalignment='top')
 
+    if args.makegif:
+        outfilename = outfilename.replace(".pdf", ".png")
     print(f"Saving {outfilename} ...")
-    fig.savefig(outfilename)
+    fig.savefig(outfilename, dpi=300)
 
     if args.show:
         plt.show()
@@ -963,6 +969,8 @@ def addargs(parser: argparse.ArgumentParser) -> None:
         help="Savitzky-Golay filter. Specify the window_length and polyorder.e.g. -filtersavgol 5 3",
     )
 
+    parser.add_argument("--makegif", action="store_true", help="Make a gif with time evolution (requires --multiplot)")
+
     parser.add_argument("--notitle", action="store_true", help="Suppress the top title from the plot")
 
     parser.add_argument("-plotlist", type=list, default=[], help="Plot list (when calling from Python only)")
@@ -1136,6 +1144,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
 
     assoc_cells, mgi_of_propcells = at.get_grid_mapping(modelpath)
 
+    outdir = args.outputfile if (args.outputfile).is_dir() else Path()
     if not args.readonlymgi and (args.modelgridindex is not None or args.x in {"time", "timestep"}):
         # plot time evolution in specific cell
         if not args.x:
@@ -1191,7 +1200,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         frames_timesteps_included = (
             [[ts] for ts in range(timestepmin, timestepmax + 1)] if args.multiplot else [timesteps_included]
         )
-        pdf_files = []
+        outputfiles = []
         for timesteps_included in frames_timesteps_included:
             timestepslist_unfiltered = [timesteps_included] * len(allnonemptymgilist)
             outfilename = make_plot(
@@ -1204,10 +1213,20 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
                 args=args,
             )
 
-            pdf_files.append(outfilename)
+            outputfiles.append(outfilename)
+
+        if args.makegif:
+            assert args.multiplot
+            import imageio.v2 as iio
 
-        if len(pdf_files) > 1:
-            at.join_pdf_files(pdf_files)
+            gifname = outdir / f"plotestim_evolution_ts{timestepmin:03d}_ts{timestepmax:03d}.gif"
+            with iio.get_writer(gifname, mode="I", duration=500) as writer:
+                for filename in outputfiles:
+                    image = iio.imread(filename)
+                    writer.append_data(image)  # type: ignore[attr-defined]
+            print(f"Created gif: {gifname}")
+        elif len(outputfiles) > 1:
+            at.join_pdf_files(outputfiles)
 
 
 if __name__ == "__main__":
diff --git a/artistools/plotspherical.py b/artistools/plotspherical.py
index e87fac5b3..9c643de20 100755
--- a/artistools/plotspherical.py
+++ b/artistools/plotspherical.py
@@ -352,7 +352,7 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non
         import imageio.v2 as iio
 
         gifname = outdir / "sphericalplot.gif" if (args.outputfile).is_dir() else args.outputfile
-        with iio.get_writer(gifname, mode="I", fps=1.5) as writer:
+        with iio.get_writer(gifname, mode="I", duration=(1000 * 1 / 1.5)) as writer:
             for filename in outputfilenames:
                 image = iio.imread(filename)
                 writer.append_data(image)  # type: ignore[attr-defined]

From 36b9870233788ef908a3e701c807c0c307a5ad6a Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 23:22:03 +0000
Subject: [PATCH 130/150] Add png multiplot

---
 artistools/__init__.py                  |  2 +-
 artistools/estimators/plotestimators.py | 45 ++++++++++++++-----------
 artistools/misc.py                      |  2 +-
 artistools/radfield.py                  |  2 +-
 4 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/artistools/__init__.py b/artistools/__init__.py
index bd9e9f8bb..a36247620 100644
--- a/artistools/__init__.py
+++ b/artistools/__init__.py
@@ -85,10 +85,10 @@
 from artistools.misc import get_wid_init_at_tmin
 from artistools.misc import get_wid_init_at_tmodel
 from artistools.misc import get_z_a_nucname
-from artistools.misc import join_pdf_files
 from artistools.misc import linetuple
 from artistools.misc import makelist
 from artistools.misc import match_closest_time
+from artistools.misc import merge_pdf_files
 from artistools.misc import namedtuple
 from artistools.misc import parse_range
 from artistools.misc import parse_range_list
diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index fc72c83c1..16b2e3aff 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -808,7 +808,7 @@ def make_plot(
     if len(set(mgilist)) == 1 and len(timestepslist[0]) > 1:  # single grid cell versus time plot
         figure_title = f"{modelname}\nCell {mgilist[0]}"
 
-        defaultoutputfile = Path("plotestimators_cell{modelgridindex:03d}.pdf")
+        defaultoutputfile = Path("plotestimators_cell{modelgridindex:03d}.{format}")
         if Path(args.outputfile).is_dir():
             args.outputfile = str(Path(args.outputfile, defaultoutputfile))
 
@@ -823,19 +823,17 @@ def make_plot(
             figure_title = f"{modelname}\nTimestep {timestepslist[0][0]} ({timeavg:.2f}d)"
         print("Plotting ", figure_title.replace("\n", " "))
 
-        defaultoutputfile = Path("plotestimators_ts{timestep:02d}_{timeavg:.2f}d.pdf")
+        defaultoutputfile = Path("plotestimators_ts{timestep:02d}_{timeavg:.2f}d.{format}")
         if Path(args.outputfile).is_dir():
             args.outputfile = str(Path(args.outputfile, defaultoutputfile))
 
         assert isinstance(timestepslist[0], list)
-        outfilename = str(args.outputfile).format(timestep=timestepslist[0][0], timeavg=timeavg)
+        outfilename = str(args.outputfile).format(timestep=timestepslist[0][0], timeavg=timeavg, format=args.format)
 
     if not args.notitle:
         axes[0].set_title(figure_title, fontsize=8)
     # plt.suptitle(figure_title, fontsize=11, verticalalignment='top')
 
-    if args.makegif:
-        outfilename = outfilename.replace(".pdf", ".png")
     print(f"Saving {outfilename} ...")
     fig.savefig(outfilename, dpi=300)
 
@@ -969,6 +967,8 @@ def addargs(parser: argparse.ArgumentParser) -> None:
         help="Savitzky-Golay filter. Specify the window_length and polyorder.e.g. -filtersavgol 5 3",
     )
 
+    parser.add_argument("-format", "-f", default="pdf", choices=["pdf", "png"], help="Set format of output plot files")
+
     parser.add_argument("--makegif", action="store_true", help="Make a gif with time evolution (requires --multiplot)")
 
     parser.add_argument("--notitle", action="store_true", help="Suppress the top title from the plot")
@@ -1062,11 +1062,11 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         #  ['_yscale', 'linear']],
         # [['initmasses', ['Ni_56', 'He', 'C', 'Mg']]],
         # ['heating_gamma/gamma_dep'],
-        ["nne", ["_ymin", 1e5], ["_ymax", 1e10]],
+        # ["nne", ["_ymin", 1e5], ["_ymax", 1e10]],
         ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 15000]],
         # ["Te"],
         # ["Te", "TR"],
-        [["averageionisation", ["Sr"]]],
+        [["averageionisation", ["Sr", "Y", "Zr"]]],
         # [["averageexcitation", ["Fe II", "Fe III"]]],
         # [["populations", ["Sr90", "Sr91", "Sr92", "Sr94"]]],
         [["populations", ["Sr I", "Sr II", "Sr III"]]],
@@ -1200,6 +1200,11 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         frames_timesteps_included = (
             [[ts] for ts in range(timestepmin, timestepmax + 1)] if args.multiplot else [timesteps_included]
         )
+
+        if args.makegif:
+            args.multiplot = True
+            args.format = "png"
+
         outputfiles = []
         for timesteps_included in frames_timesteps_included:
             timestepslist_unfiltered = [timesteps_included] * len(allnonemptymgilist)
@@ -1215,18 +1220,20 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
 
             outputfiles.append(outfilename)
 
-        if args.makegif:
-            assert args.multiplot
-            import imageio.v2 as iio
-
-            gifname = outdir / f"plotestim_evolution_ts{timestepmin:03d}_ts{timestepmax:03d}.gif"
-            with iio.get_writer(gifname, mode="I", duration=500) as writer:
-                for filename in outputfiles:
-                    image = iio.imread(filename)
-                    writer.append_data(image)  # type: ignore[attr-defined]
-            print(f"Created gif: {gifname}")
-        elif len(outputfiles) > 1:
-            at.join_pdf_files(outputfiles)
+        if len(outputfiles) > 1:
+            if args.makegif:
+                assert args.multiplot
+                assert args.format == "png"
+                import imageio.v2 as iio
+
+                gifname = outdir / f"plotestim_evolution_ts{timestepmin:03d}_ts{timestepmax:03d}.gif"
+                with iio.get_writer(gifname, mode="I", duration=1000) as writer:
+                    for filename in outputfiles:
+                        image = iio.imread(filename)
+                        writer.append_data(image)  # type: ignore[attr-defined]
+                print(f"Created gif: {gifname}")
+            elif args.format == "pdf":
+                at.merge_pdf_files(outputfiles)
 
 
 if __name__ == "__main__":
diff --git a/artistools/misc.py b/artistools/misc.py
index 1a4824703..8ab02116f 100644
--- a/artistools/misc.py
+++ b/artistools/misc.py
@@ -925,7 +925,7 @@ def savgolfilterfunc(ylist: t.Any) -> t.Any:
     return filterfunc
 
 
-def join_pdf_files(pdf_files: list[str]) -> None:
+def merge_pdf_files(pdf_files: list[str]) -> None:
     """Merge a list of PDF files into a single PDF file."""
     from PyPDF2 import PdfMerger
 
diff --git a/artistools/radfield.py b/artistools/radfield.py
index 63f2bc4d7..3751d110b 100755
--- a/artistools/radfield.py
+++ b/artistools/radfield.py
@@ -1050,7 +1050,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
 
     if len(pdf_list) > 1:
         print(pdf_list)
-        at.join_pdf_files(pdf_list)
+        at.merge_pdf_files(pdf_list)
 
 
 if __name__ == "__main__":

From 7a9c2bbf158a1957235780b597647d22ccfbe8d1 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 23:23:28 +0000
Subject: [PATCH 131/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 16b2e3aff..603bb2366 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -237,7 +237,6 @@ def plot_average_ionisation_excitation(
         xlist, ylist = at.estimators.apply_filters(xlist, ylist, args)
         if startfromzero:
             ylist.insert(0, ylist[0])
-        print(f"  Plotting {seriestype} {paramvalue}")
 
         ax.plot(xlist, ylist, label=paramvalue, color=color, **plotkwargs)
 

From 97e1a36fc7040f106311df732c5f14b2724d7fea Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 23:40:27 +0000
Subject: [PATCH 132/150] Fix estimator output file names

---
 artistools/estimators/plotestimators.py | 23 +++++++++++++----------
 artistools/misc.py                      |  2 +-
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 603bb2366..3b6283b2c 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -807,31 +807,34 @@ def make_plot(
     if len(set(mgilist)) == 1 and len(timestepslist[0]) > 1:  # single grid cell versus time plot
         figure_title = f"{modelname}\nCell {mgilist[0]}"
 
-        defaultoutputfile = Path("plotestimators_cell{modelgridindex:03d}.{format}")
+        defaultoutputfile = "plotestimators_cell{modelgridindex:03d}.{format}"
         if Path(args.outputfile).is_dir():
-            args.outputfile = str(Path(args.outputfile, defaultoutputfile))
+            args.outputfile = str(Path(args.outputfile) / defaultoutputfile)
 
         outfilename = str(args.outputfile).format(modelgridindex=mgilist[0])
 
     else:
-        timeavg = (args.timemin + args.timemax) / 2.0
         if args.multiplot:
-            timedays = float(at.get_timestep_time(modelpath, timestepslist[0][0]))
-            figure_title = f"{modelname}\nTimestep {timestepslist[0][0]} ({timedays:.2f}d)"
+            timestep = f"ts{timestepslist[0][0]:02d}"
+            timedays = f"{at.get_timestep_time(modelpath, timestepslist[0][0]):.2f}d"
         else:
-            figure_title = f"{modelname}\nTimestep {timestepslist[0][0]} ({timeavg:.2f}d)"
+            timestepmin = min(timestepslist[0])
+            timestepmax = max(timestepslist[0])
+            timestep = f"ts{timestepmin:02d}-ts{timestepmax:02d}"
+            timedays = f"{at.get_timestep_time(modelpath, timestepmin):.2f}d-{at.get_timestep_time(modelpath, timestepmax):.2f}d"
+
+        figure_title = f"{modelname}\nTimestep {timestep} ({timedays})"
         print("Plotting ", figure_title.replace("\n", " "))
 
-        defaultoutputfile = Path("plotestimators_ts{timestep:02d}_{timeavg:.2f}d.{format}")
+        defaultoutputfile = "plotestimators_{timestep}_{timedays}.{format}"
         if Path(args.outputfile).is_dir():
-            args.outputfile = str(Path(args.outputfile, defaultoutputfile))
+            args.outputfile = str(Path(args.outputfile) / defaultoutputfile)
 
         assert isinstance(timestepslist[0], list)
-        outfilename = str(args.outputfile).format(timestep=timestepslist[0][0], timeavg=timeavg, format=args.format)
+        outfilename = str(args.outputfile).format(timestep=timestep, timedays=timedays, format=args.format)
 
     if not args.notitle:
         axes[0].set_title(figure_title, fontsize=8)
-    # plt.suptitle(figure_title, fontsize=11, verticalalignment='top')
 
     print(f"Saving {outfilename} ...")
     fig.savefig(outfilename, dpi=300)
diff --git a/artistools/misc.py b/artistools/misc.py
index 8ab02116f..3d2a9e4a5 100644
--- a/artistools/misc.py
+++ b/artistools/misc.py
@@ -936,7 +936,7 @@ def merge_pdf_files(pdf_files: list[str]) -> None:
             merger.append(pdffile)
         Path(pdfpath).unlink()
 
-    resultfilename = f'{pdf_files[0].split(".")[0]}-{pdf_files[-1].split(".")[0]}'
+    resultfilename = f'{pdf_files[0].replace(".pdf","")}-{pdf_files[-1].replace(".pdf","")}'
     with Path(f"{resultfilename}.pdf").open("wb") as resultfile:
         merger.write(resultfile)
 

From 159eab0a9ac774b281e60a79de1070910a7fc686 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 23:41:38 +0000
Subject: [PATCH 133/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 3b6283b2c..151dccad2 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -811,7 +811,7 @@ def make_plot(
         if Path(args.outputfile).is_dir():
             args.outputfile = str(Path(args.outputfile) / defaultoutputfile)
 
-        outfilename = str(args.outputfile).format(modelgridindex=mgilist[0])
+        outfilename = str(args.outputfile).format(modelgridindex=mgilist[0], format=args.format)
 
     else:
         if args.multiplot:

From c74b6f9aa3e9fe83fb1d03f1713a07334fe96150 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Sun, 12 Nov 2023 23:46:00 +0000
Subject: [PATCH 134/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 151dccad2..520c920b0 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -1068,7 +1068,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 15000]],
         # ["Te"],
         # ["Te", "TR"],
-        [["averageionisation", ["Sr", "Y", "Zr"]]],
+        [["averageionisation", ["Sr"]]],
         # [["averageexcitation", ["Fe II", "Fe III"]]],
         # [["populations", ["Sr90", "Sr91", "Sr92", "Sr94"]]],
         [["populations", ["Sr I", "Sr II", "Sr III"]]],

From b875aaa1dc234981af40bbf6bf72d7e339662ad9 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Mon, 13 Nov 2023 00:45:25 +0000
Subject: [PATCH 135/150] Update

---
 artistools/estimators/plotestimators.py  |  2 +-
 artistools/inputmodel/inputmodel_misc.py |  3 ++-
 artistools/packets/packets.py            |  6 +++--
 artistools/plotspherical.py              | 30 +++++++++++++++---------
 4 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 520c920b0..26cf7c2b8 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -1071,7 +1071,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         [["averageionisation", ["Sr"]]],
         # [["averageexcitation", ["Fe II", "Fe III"]]],
         # [["populations", ["Sr90", "Sr91", "Sr92", "Sr94"]]],
-        [["populations", ["Sr I", "Sr II", "Sr III"]]],
+        [["populations", ["Sr I", "Sr II", "Sr III", "Sr IV"]]],
         # [['populations', ['He I', 'He II', 'He III']]],
         # [['populations', ['C I', 'C II', 'C III', 'C IV', 'C V']]],
         # [['populations', ['O I', 'O II', 'O III', 'O IV']]],
diff --git a/artistools/inputmodel/inputmodel_misc.py b/artistools/inputmodel/inputmodel_misc.py
index 5b3289900..a34f00aa0 100644
--- a/artistools/inputmodel/inputmodel_misc.py
+++ b/artistools/inputmodel/inputmodel_misc.py
@@ -373,7 +373,8 @@ def get_modeldata_polars(
         gc.collect()
         dfmodel = pl.scan_parquet(filenameparquet)
 
-    print(f"  model is {modelmeta['dimensions']}D with {modelmeta['npts_model']} cells")
+    if not printwarningsonly:
+        print(f"  model is {modelmeta['dimensions']}D with {modelmeta['npts_model']} cells")
     dfmodel = dfmodel.lazy()
 
     if get_elemabundances:
diff --git a/artistools/packets/packets.py b/artistools/packets/packets.py
index e2affe451..cf3e4755e 100644
--- a/artistools/packets/packets.py
+++ b/artistools/packets/packets.py
@@ -199,7 +199,9 @@ def emtrue_timestep(packet) -> int:
     return dfpackets
 
 
-def add_derived_columns_lazy(dfpackets: pl.LazyFrame, modelmeta: dict, dfmodel: pd.DataFrame | None) -> pl.LazyFrame:
+def add_derived_columns_lazy(
+    dfpackets: pl.LazyFrame, modelmeta: dict[str, t.Any], dfmodel: pd.DataFrame | pl.LazyFrame | None
+) -> pl.LazyFrame:
     """Add columns to a packets DataFrame that are derived from the values that are stored in the packets files.
 
     We might as well add everything, since the columns only get calculated when they are actually used (polars LazyFrame).
@@ -269,7 +271,7 @@ def add_derived_columns_lazy(dfpackets: pl.LazyFrame, modelmeta: dict, dfmodel:
             )
     elif modelmeta["dimensions"] == 1:
         assert dfmodel is not None, "dfmodel must be provided for 1D models to set em_modelgridindex"
-        velbins = (dfmodel["vel_r_max_kmps"] * 1000).to_list()
+        velbins = (dfmodel.select("vel_r_max_kmps").lazy().collect()["vel_r_max_kmps"] * 1000.0).to_list()
         dfpackets = dfpackets.with_columns(
             (
                 pl.col("emission_velocity")
diff --git a/artistools/plotspherical.py b/artistools/plotspherical.py
index 9c643de20..51b3c9811 100755
--- a/artistools/plotspherical.py
+++ b/artistools/plotspherical.py
@@ -25,6 +25,9 @@ def plot_spherical(
     timemaxdays: float | None,
     nphibins: int,
     ncosthetabins: int,
+    dfmodel: pl.LazyFrame | None = None,
+    modelmeta: dict[str, t.Any] | None = None,
+    dfestimators: pl.LazyFrame | None = None,
     maxpacketfiles: int | None = None,
     atomic_number: int | None = None,
     ionstage: int | None = None,
@@ -36,8 +39,6 @@ def plot_spherical(
     if plotvars is None:
         plotvars = ["luminosity", "emvelocityoverc", "emlosvelocityoverc"]
 
-    dfmodel, modelmeta = at.get_modeldata(modelpath=modelpath, getheadersonly=True, printwarningsonly=True)
-
     _, tmin_d_valid, tmax_d_valid = at.get_escaped_arrivalrange(modelpath)
     if tmin_d_valid is None or tmax_d_valid is None:
         print("WARNING! The observer never gets light from the entire ejecta. Plotting all packets anyway")
@@ -81,7 +82,6 @@ def plot_spherical(
     # dfpackets = dfpackets.filter(pl.col("dirz") > 0.9)
 
     aggs = []
-    dfpackets = at.packets.add_derived_columns_lazy(dfpackets, modelmeta=modelmeta, dfmodel=dfmodel)
 
     if "emvelocityoverc" in plotvars:
         aggs.append(
@@ -119,13 +119,13 @@ def plot_spherical(
             ).alias("em_timestep")
         )
 
-        df_estimators = (
-            at.estimators.scan_estimators(modelpath=modelpath)
-            .select(["timestep", "modelgridindex", "TR"])
+        assert dfestimators is not None
+        dfestimators = (
+            dfestimators.select(["timestep", "modelgridindex", "TR"])
             .drop_nulls()
             .rename({"timestep": "em_timestep", "modelgridindex": "em_modelgridindex", "TR": "em_TR"})
         )
-        dfpackets = dfpackets.join(df_estimators, on=["em_timestep", "em_modelgridindex"], how="left")
+        dfpackets = dfpackets.join(dfestimators, on=["em_timestep", "em_modelgridindex"], how="left")
         aggs.append(((pl.col("em_TR") * pl.col("e_rf")).mean() / pl.col("e_rf").mean()).alias("temperature"))
 
     if atomic_number is not None or ionstage is not None:
@@ -141,7 +141,6 @@ def plot_spherical(
         dfpackets = dfpackets.filter(pl.col("emissiontype").is_in(selected_emtypes))
 
     aggs.append(pl.count())
-
     dfpackets = dfpackets.group_by(["costhetabin", "phibin"]).agg(aggs)
     dfpackets = dfpackets.select(["costhetabin", "phibin", "count", *plotvars])
 
@@ -176,7 +175,7 @@ def plot_spherical(
         1,
         figsize=(figscale * at.get_config()["figwidth"], 3.7 * len(plotvars)),
         subplot_kw={"projection": "mollweide"},
-        tight_layout={"pad": 0.1, "w_pad": 0.0, "h_pad": 0.0},
+        tight_layout={"pad": 0.5, "w_pad": 0.5, "h_pad": 0.5},
     )
 
     if len(plotvars) == 1:
@@ -295,9 +294,13 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non
         assert args.atomic_number is None
         args.atomic_number = at.get_atomic_number(args.elem)
 
+    dfmodel, modelmeta = at.get_modeldata_polars(modelpath=args.modelpath, getheadersonly=True, printwarningsonly=True)
+    dfestimators = at.estimators.scan_estimators(modelpath=args.modelpath) if "temperature" in args.plotvars else None
+
     nprocs_read, dfpackets = at.packets.get_packets_pl(
         args.modelpath, args.maxpacketfiles, packet_type="TYPE_ESCAPE", escape_type="TYPE_RPKT"
     )
+    dfpackets = at.packets.add_derived_columns_lazy(dfpackets, modelmeta=modelmeta, dfmodel=dfmodel)
 
     if args.makegif:
         tstarts = at.get_timestep_times(args.modelpath, loc="start")
@@ -319,6 +322,9 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non
         fig, axes, timemindays, timemaxdays = plot_spherical(
             modelpath=args.modelpath,
             dfpackets=dfpackets,
+            dfestimators=dfestimators,
+            dfmodel=dfmodel,
+            modelmeta=modelmeta,
             nprocs_read=nprocs_read,
             timemindays=tstart,
             timemaxdays=tend,
@@ -341,7 +347,7 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non
             else outdir / args.outputfile
         )
 
-        fig.savefig(outfilename, format=outformat)
+        fig.savefig(outfilename, format=outformat, dpi=300)
         print(f"Saved {outfilename}")
         plt.close()
         plt.clf()
@@ -351,7 +357,9 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non
     if args.makegif:
         import imageio.v2 as iio
 
-        gifname = outdir / "sphericalplot.gif" if (args.outputfile).is_dir() else args.outputfile
+        gifname = (
+            outdir / "sphericalplot.gif" if (args.outputfile).is_dir() else args.outputfile.replace(".pdf", ".gif")
+        )
         with iio.get_writer(gifname, mode="I", duration=(1000 * 1 / 1.5)) as writer:
             for filename in outputfilenames:
                 image = iio.imread(filename)

From 027b46453cc2e9f75b8afad926615f0b04b8b74f Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Mon, 13 Nov 2023 00:45:32 +0000
Subject: [PATCH 136/150] Update requirements.txt

---
 requirements.txt | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index db6f4b4bf..ad3978064 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,13 +1,13 @@
-argcomplete>=3.1.4
+argcomplete>=3.1.6
 astropy>=5.3.4
 coverage>=7.3.2
 extinction>=0.4.6
 imageio>=2.32.0
 matplotlib>=3.8.1
 mypy>=1.7.0
-numpy>=1.26.1
-pandas>=2.1.2
-polars>=0.19.12
+numpy>=1.26.2
+pandas>=2.1.3
+polars>=0.19.13
 pre-commit>=3.5.0
 pyarrow>=14.0.1
 pynonthermal>=2021.10.12

From 0ddd2739c357ef0c8eb596cf5df6613a9c84564b Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Mon, 13 Nov 2023 09:22:44 +0000
Subject: [PATCH 137/150] Update plotspherical.py

---
 artistools/plotspherical.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/artistools/plotspherical.py b/artistools/plotspherical.py
index 51b3c9811..e965fc7a6 100755
--- a/artistools/plotspherical.py
+++ b/artistools/plotspherical.py
@@ -240,6 +240,7 @@ def addargs(parser: argparse.ArgumentParser) -> None:
         default=Path(),
         help="Path to ARTIS folder",
     )
+    parser.add_argument("-timestep", "-ts", action="store", type=int, default=None, help="Timestep index")
     parser.add_argument("-timemin", "-tmin", action="store", type=float, default=None, help="Time minimum [d]")
     parser.add_argument("-timemax", "-tmax", action="store", type=float, default=None, help="Time maximum [d]")
     parser.add_argument("-nphibins", action="store", type=int, default=64, help="Number of azimuthal bins")
@@ -273,7 +274,7 @@ def addargs(parser: argparse.ArgumentParser) -> None:
         action="store",
         dest="outputfile",
         type=Path,
-        default=Path("plotspherical.pdf"),
+        default=Path("plotspherical.{outformat}"),
         help="Filename for PDF file",
     )
 
@@ -302,15 +303,18 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non
     )
     dfpackets = at.packets.add_derived_columns_lazy(dfpackets, modelmeta=modelmeta, dfmodel=dfmodel)
 
+    tstarts = at.get_timestep_times(args.modelpath, loc="start")
+    tends = at.get_timestep_times(args.modelpath, loc="end")
     if args.makegif:
-        tstarts = at.get_timestep_times(args.modelpath, loc="start")
-        tends = at.get_timestep_times(args.modelpath, loc="end")
         time_ranges = [
             (tstart, tend)
             for tstart, tend in zip(tstarts, tends)
             if ((args.timemin is None or tstart >= args.timemin) and (args.timemax is None or tend <= args.timemax))
         ]
         outformat = "png"
+    elif args.timestep is not None:
+        time_ranges = [(tstarts[args.timestep], tends[args.timestep])]
+        outformat = "pdf"
     else:
         time_ranges = [(args.timemin, args.timemax)]
         outformat = "pdf"
@@ -341,11 +345,11 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non
 
         axes[0].set_title(f"{timemindays:.2f}-{timemaxdays:.2f} days")
 
-        outfilename = (
-            outdir / f"plotspherical_{timemindays:.2f}-{timemaxdays:.2f}d.{outformat}"
+        outfilename = str(
+            outdir / "plotspherical_{timemindays:.2f}-{timemaxdays:.2f}d.{outformat}"
             if args.makegif or (args.outputfile).is_dir()
             else outdir / args.outputfile
-        )
+        ).format(timemindays=timemindays, timemaxdays=timemaxdays, outformat=outformat)
 
         fig.savefig(outfilename, format=outformat, dpi=300)
         print(f"Saved {outfilename}")
@@ -358,7 +362,7 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non
         import imageio.v2 as iio
 
         gifname = (
-            outdir / "sphericalplot.gif" if (args.outputfile).is_dir() else args.outputfile.replace(".pdf", ".gif")
+            outdir / "sphericalplot.gif" if (args.outputfile).is_dir() else args.outputfile.format(outformat=outformat)
         )
         with iio.get_writer(gifname, mode="I", duration=(1000 * 1 / 1.5)) as writer:
             for filename in outputfilenames:

From fbcd24cc17426a337e30b416618ec44756cbfc76 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Mon, 13 Nov 2023 09:47:11 +0000
Subject: [PATCH 138/150] Update plotspherical.py

---
 artistools/plotspherical.py | 39 ++++++++++++++++++++++---------------
 1 file changed, 23 insertions(+), 16 deletions(-)

diff --git a/artistools/plotspherical.py b/artistools/plotspherical.py
index e965fc7a6..eb7e5547b 100755
--- a/artistools/plotspherical.py
+++ b/artistools/plotspherical.py
@@ -240,7 +240,7 @@ def addargs(parser: argparse.ArgumentParser) -> None:
         default=Path(),
         help="Path to ARTIS folder",
     )
-    parser.add_argument("-timestep", "-ts", action="store", type=int, default=None, help="Timestep index")
+    parser.add_argument("-timestep", "-ts", action="store", type=str, default=None, help="Timestep index")
     parser.add_argument("-timemin", "-tmin", action="store", type=float, default=None, help="Time minimum [d]")
     parser.add_argument("-timemax", "-tmax", action="store", type=float, default=None, help="Time maximum [d]")
     parser.add_argument("-nphibins", action="store", type=int, default=64, help="Number of azimuthal bins")
@@ -273,11 +273,13 @@ def addargs(parser: argparse.ArgumentParser) -> None:
         "-o",
         action="store",
         dest="outputfile",
-        type=Path,
-        default=Path("plotspherical.{outformat}"),
-        help="Filename for PDF file",
+        type=str,
+        default="",
+        help="Filename for plot output file",
     )
 
+    parser.add_argument("-format", "-f", default="", choices=["pdf", "png"], help="Set format of output plot files")
+
 
 def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = None, **kwargs: t.Any) -> None:
     """Plot direction maps based on escaped packets."""
@@ -307,21 +309,23 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non
     tends = at.get_timestep_times(args.modelpath, loc="end")
     if args.makegif:
         time_ranges = [
-            (tstart, tend)
-            for tstart, tend in zip(tstarts, tends)
+            (tstart, tend, f"timestep {ts}")
+            for ts, (tstart, tend) in enumerate(zip(tstarts, tends))
             if ((args.timemin is None or tstart >= args.timemin) and (args.timemax is None or tend <= args.timemax))
         ]
         outformat = "png"
     elif args.timestep is not None:
-        time_ranges = [(tstarts[args.timestep], tends[args.timestep])]
-        outformat = "pdf"
+        time_ranges = [
+            (tstarts[int(ts)], tends[int(ts)], f"timestep {ts}") for ts in at.parse_range_list(args.timestep)
+        ]
+        outformat = args.format or "pdf"
     else:
-        time_ranges = [(args.timemin, args.timemax)]
-        outformat = "pdf"
+        time_ranges = [(args.timemin, args.timemax, "")]
+        outformat = args.format or "pdf"
 
-    outdir = args.outputfile if (args.outputfile).is_dir() else Path()
     outputfilenames = []
-    for tstart, tend in time_ranges:
+    for tstart, tend, label in time_ranges:
+        print(f"Plotting spherical map for {tstart:.2f}-{tend:.2f} days {label}")
         # tstart and tend are requested, but the actual plotted time range may be different
         fig, axes, timemindays, timemaxdays = plot_spherical(
             modelpath=args.modelpath,
@@ -345,10 +349,11 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non
 
         axes[0].set_title(f"{timemindays:.2f}-{timemaxdays:.2f} days")
 
+        defaultfilename = "plotspherical_{timemindays:.2f}-{timemaxdays:.2f}d.{outformat}"
         outfilename = str(
-            outdir / "plotspherical_{timemindays:.2f}-{timemaxdays:.2f}d.{outformat}"
-            if args.makegif or (args.outputfile).is_dir()
-            else outdir / args.outputfile
+            args.outputfile
+            if (args.outputfile and not Path(args.outputfile).is_dir() and not args.makegif)
+            else Path(args.outputfile) / defaultfilename
         ).format(timemindays=timemindays, timemaxdays=timemaxdays, outformat=outformat)
 
         fig.savefig(outfilename, format=outformat, dpi=300)
@@ -362,7 +367,9 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non
         import imageio.v2 as iio
 
         gifname = (
-            outdir / "sphericalplot.gif" if (args.outputfile).is_dir() else args.outputfile.format(outformat=outformat)
+            Path(args.outputfile) / "sphericalplot.gif"
+            if Path(args.outputfile).is_dir()
+            else args.outputfile.format(outformat=outformat)
         )
         with iio.get_writer(gifname, mode="I", duration=(1000 * 1 / 1.5)) as writer:
             for filename in outputfilenames:

From d43c090bb13156135243264443752236ebef7854 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Mon, 13 Nov 2023 10:11:07 +0000
Subject: [PATCH 139/150] Update plotspherical.py

---
 artistools/plotspherical.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/artistools/plotspherical.py b/artistools/plotspherical.py
index eb7e5547b..d56a08550 100755
--- a/artistools/plotspherical.py
+++ b/artistools/plotspherical.py
@@ -325,7 +325,8 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non
 
     outputfilenames = []
     for tstart, tend, label in time_ranges:
-        print(f"Plotting spherical map for {tstart:.2f}-{tend:.2f} days {label}")
+        if tstart is not None and tend is not None:
+            print(f"Plotting spherical map for {tstart:.2f}-{tend:.2f} days {label}")
         # tstart and tend are requested, but the actual plotted time range may be different
         fig, axes, timemindays, timemaxdays = plot_spherical(
             modelpath=args.modelpath,

From 483bb046c81e08e86b43c54bfd8790afb0fbafcb Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Mon, 13 Nov 2023 10:16:40 +0000
Subject: [PATCH 140/150] Update plotspherical.py

---
 artistools/plotspherical.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/artistools/plotspherical.py b/artistools/plotspherical.py
index d56a08550..6329518cf 100755
--- a/artistools/plotspherical.py
+++ b/artistools/plotspherical.py
@@ -173,9 +173,10 @@ def plot_spherical(
     fig, axes = plt.subplots(
         len(plotvars),
         1,
-        figsize=(figscale * at.get_config()["figwidth"], 3.7 * len(plotvars)),
+        figsize=(figscale * at.get_config()["figwidth"], 3.2 * len(plotvars)),
         subplot_kw={"projection": "mollweide"},
-        tight_layout={"pad": 0.5, "w_pad": 0.5, "h_pad": 0.5},
+        # tight_layout={"pad": 0, "w_pad": 0, "h_pad": 5.0},
+        gridspec_kw={"wspace": 0.0, "hspace": 0.0},
     )
 
     if len(plotvars) == 1:
@@ -204,7 +205,7 @@ def plot_spherical(
             case _:
                 raise AssertionError
 
-        cbar = fig.colorbar(colormesh, ax=ax, location="bottom")
+        cbar = fig.colorbar(colormesh, ax=ax, location="bottom", pad=0.2)
         cbar.outline.set_linewidth(0)  # type: ignore[operator]
         cbar.ax.tick_params(axis="both", direction="out")
         cbar.ax.xaxis.set_ticks_position("top")

From 0bf797577601fcd907695b48e667d9ec58c59643 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Mon, 13 Nov 2023 20:04:20 +0100
Subject: [PATCH 141/150] Update plotinitialcomposition.py

---
 artistools/inputmodel/plotinitialcomposition.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/artistools/inputmodel/plotinitialcomposition.py b/artistools/inputmodel/plotinitialcomposition.py
index d8f42531b..2151505b3 100755
--- a/artistools/inputmodel/plotinitialcomposition.py
+++ b/artistools/inputmodel/plotinitialcomposition.py
@@ -196,12 +196,12 @@ def plot_2d_initial_abundances(modelpath, args=None) -> None:
         cbar.set_label(r"log10($\rho$) [g/cm3]" if args.logcolorscale else r"$\rho$ [g/cm3]")
 
     defaultfilename = Path(modelpath) / f"plotcomposition_{','.join(args.plotvars)}.pdf"
-    if args.outputfile.is_dir():
-        outfilename = defaultfilename
+    if args.outputfile and Path(args.outputfile).is_dir():
+        outfilename = Path(modelpath) / defaultfilename
     elif args.outputfile:
         outfilename = args.outputfile
     else:
-        outfilename = Path(modelpath) / defaultfilename
+        outfilename = defaultfilename
 
     plt.savefig(outfilename, format="pdf")
 
@@ -266,7 +266,7 @@ def make_3d_plot(modelpath, args):
     model = merge_dfs
 
     # choose what surface will be coloured by
-    if args.rho:
+    if "rho" in args.plotvars:
         coloursurfaceby = "rho"
     elif args.opacity:
         model["opacity"] = at.inputmodel.opacityinputfile.get_opacity_from_file(modelpath)

From 087515ff41f9755e64c518c04c57140348c6f872 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Mon, 13 Nov 2023 20:04:41 +0100
Subject: [PATCH 142/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 26cf7c2b8..b28ebcee4 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -730,7 +730,7 @@ def plot_subplot(
                 if (plotitems[0][0] == "populations" and args.yscale == "log")
                 else {}
             ),
-            markerscale=10,
+            markerscale=3,
         )
 
 
@@ -782,7 +782,7 @@ def make_plot(
     if args.markersonly:
         plotkwargs["linestyle"] = "None"
         plotkwargs["marker"] = "."
-        plotkwargs["markersize"] = 1
+        plotkwargs["markersize"] = 3
         plotkwargs["alpha"] = 0.5
 
         # with no lines, line styles cannot distringuish ions
@@ -1065,7 +1065,8 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         # [['initmasses', ['Ni_56', 'He', 'C', 'Mg']]],
         # ['heating_gamma/gamma_dep'],
         # ["nne", ["_ymin", 1e5], ["_ymax", 1e10]],
-        ["TR", ["_yscale", "linear"], ["_ymin", 1000], ["_ymax", 15000]],
+        ["rho", ["_yscale", "log"]],
+        ["TR", ["_yscale", "linear"]],  # , ["_ymin", 1000], ["_ymax", 15000]
         # ["Te"],
         # ["Te", "TR"],
         [["averageionisation", ["Sr"]]],

From 53ff452b5ddcd66955843229f08de2e1ea948d61 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Mon, 13 Nov 2023 20:34:11 +0100
Subject: [PATCH 143/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index b28ebcee4..903d7521a 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -561,14 +561,6 @@ def get_xlist(
     elif xvariable == "timestep":
         estimators = estimators.with_columns(xvalue=pl.col("timestep"), plotpointid=pl.col("timestep"))
     elif xvariable == "time":
-        timearray = at.get_timestep_times(modelpath)
-        estimators = estimators.lazy().join(
-            pl.DataFrame({"timestep": range(len(timearray)), "time_mid": timearray})
-            .with_columns(pl.col("timestep").cast(pl.Int32))
-            .lazy(),
-            on="timestep",
-            how="left",
-        )
         estimators = estimators.with_columns(xvalue=pl.col("time_mid"), plotpointid=pl.col("timestep"))
     elif xvariable in {"velocity", "beta"}:
         velcolumn = "vel_r_mid"
@@ -1178,6 +1170,18 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
 
         dfmodel = dfmodel.filter(pl.col("vel_r_mid") <= modelmeta["vmax_cmps"])
         estimators = estimators.join(dfmodel, on="modelgridindex")
+        tmids = at.get_timestep_times(modelpath, loc="mid")
+        estimators = estimators.join(
+            pl.DataFrame({"timestep": range(len(tmids)), "time_mid": tmids})
+            .with_columns(pl.col("timestep").cast(pl.Int32))
+            .lazy(),
+            on="timestep",
+            how="left",
+        )
+        estimators = estimators.with_columns(
+            rho_init=pl.col("rho"),
+            rho=pl.col("rho") * (modelmeta["t_model_init_days"] / pl.col("time_mid")) ** 3,
+        )
 
         if args.readonlymgi:
             estimators = estimators.filter(pl.col("modelgridindex").is_in(args.modelgridindex))

From 91766f2c776dd085708b69358562b65c6bb46e80 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Mon, 13 Nov 2023 20:36:47 +0100
Subject: [PATCH 144/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 903d7521a..853c67554 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -396,7 +396,7 @@ def get_iontuple(ionstr):
         else:
             key = f"{seriestype}_{ionstr}"
 
-        print(f"Plotting {seriestype} {ionstr}")
+        print(f"Plotting {seriestype} {ionstr.replace('_', ' ')}")
 
         if seriestype != "populations" or args.ionpoptype == "absolute":
             scalefactor = pl.lit(1)

From 2d5e6dd1a1a586c649d746ff399c6feec981d8bb Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Mon, 13 Nov 2023 20:45:10 +0100
Subject: [PATCH 145/150] Fix test failure

---
 artistools/estimators/plotestimators.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index 853c67554..b4f486d5c 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -1118,6 +1118,14 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
             timestep=tuple(timesteps_included),
         )
     assert estimators is not None
+    tmids = at.get_timestep_times(modelpath, loc="mid")
+    estimators = estimators.join(
+        pl.DataFrame({"timestep": range(len(tmids)), "time_mid": tmids})
+        .with_columns(pl.col("timestep").cast(pl.Int32))
+        .lazy(),
+        on="timestep",
+        how="left",
+    )
 
     for ts in reversed(timesteps_included):
         tswithdata = estimators.select("timestep").unique().collect().to_series()
@@ -1170,14 +1178,6 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
 
         dfmodel = dfmodel.filter(pl.col("vel_r_mid") <= modelmeta["vmax_cmps"])
         estimators = estimators.join(dfmodel, on="modelgridindex")
-        tmids = at.get_timestep_times(modelpath, loc="mid")
-        estimators = estimators.join(
-            pl.DataFrame({"timestep": range(len(tmids)), "time_mid": tmids})
-            .with_columns(pl.col("timestep").cast(pl.Int32))
-            .lazy(),
-            on="timestep",
-            how="left",
-        )
         estimators = estimators.with_columns(
             rho_init=pl.col("rho"),
             rho=pl.col("rho") * (modelmeta["t_model_init_days"] / pl.col("time_mid")) ** 3,

From 3f5e0b6254cfd9e70298c28fa3da983f526c21f5 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Mon, 13 Nov 2023 21:56:49 +0100
Subject: [PATCH 146/150] Update estimators.py

---
 artistools/estimators/estimators.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py
index fc1f87541..619824864 100755
--- a/artistools/estimators/estimators.py
+++ b/artistools/estimators/estimators.py
@@ -35,6 +35,7 @@ def get_variableunits(key: str | None = None) -> str | dict[str, str]:
         "heating": "erg/s/cm3",
         "heating_dep/total_dep": "Ratio",
         "cooling": "erg/s/cm3",
+        "rho": "g/cm3",
         "velocity": "km/s",
         "beta": "v/c",
         "vel_r_max_kmps": "km/s",
@@ -58,6 +59,7 @@ def get_varname_formatted(varname: str) -> str:
     replacements = {
         "nne": r"n$_{\rm e}$",
         "lognne": r"Log n$_{\rm e}$",
+        "rho": r"$\rho$",
         "Te": r"T$_{\rm e}$",
         "TR": r"T$_{\rm R}$",
         "TJ": r"T$_{\rm J}$",

From dd8ab913f4b5d9d5af92ee507c75f1456f9de25d Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Mon, 13 Nov 2023 21:57:57 +0100
Subject: [PATCH 147/150] Update extensions.json

---
 .vscode/extensions.json | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.vscode/extensions.json b/.vscode/extensions.json
index fd21859c2..5f1c0a7eb 100644
--- a/.vscode/extensions.json
+++ b/.vscode/extensions.json
@@ -1,6 +1,7 @@
 {
     "recommendations": [
         "ms-python.mypy-type-checker",
-        "charliermarsh.ruff"
+        "charliermarsh.ruff",
+        "sourcery.sourcery"
     ]
 }
\ No newline at end of file

From ce0cd4e3410e16bd5a993a8489677bc622686766 Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Tue, 14 Nov 2023 23:34:01 +0100
Subject: [PATCH 148/150] Update plotinitialcomposition.py

---
 artistools/inputmodel/plotinitialcomposition.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/artistools/inputmodel/plotinitialcomposition.py b/artistools/inputmodel/plotinitialcomposition.py
index 2151505b3..47fa1614d 100755
--- a/artistools/inputmodel/plotinitialcomposition.py
+++ b/artistools/inputmodel/plotinitialcomposition.py
@@ -59,8 +59,11 @@ def plot_slice_modelcolumn(ax, dfmodelslice, modelmeta, colname, plotaxis1, plot
 
     if args.logcolorscale:
         # logscale for colormap
+        floorval = 1e-16
+        colorscale = [floorval if x < floorval or not math.isfinite(x) else x for x in colorscale]
         with np.errstate(divide="ignore"):
             colorscale = np.log10(colorscale)
+        # np.nan_to_num(colorscale, posinf=-99, neginf=-99)
 
     normalise_between_0_and_1 = False
     if normalise_between_0_and_1:
@@ -144,6 +147,7 @@ def plot_2d_initial_abundances(modelpath, args=None) -> None:
         axeschars: list[AxisType] = ["x", "y", "z"]
         plotaxis1 = next(ax for ax in axeschars if ax != sliceaxis)
         plotaxis2 = next(ax for ax in axeschars if ax not in {sliceaxis, plotaxis1})
+        print(f"Plotting slice through {sliceaxis}=0, plotting {plotaxis1} vs {plotaxis2}")
 
         df2dslice = get_2D_slice_through_3d_model(
             dfmodel=dfmodel, modelmeta=modelmeta, sliceaxis=sliceaxis, plotaxis1=plotaxis1, plotaxis2=plotaxis2
@@ -188,12 +192,14 @@ def plot_2d_initial_abundances(modelpath, args=None) -> None:
     xlabel = r"v$_{" + f"{plotaxis1}" + r"}$ [$c$]"
     ylabel = r"v$_{" + f"{plotaxis2}" + r"}$ [$c$]"
 
-    cbar = fig.colorbar(scaledmap, cax=axcbar, location="top", use_gridspec=True)
+    cbar = fig.colorbar(im, cax=axcbar, location="top", use_gridspec=True)
     axes[0].set_xlabel(xlabel)
     axes[0].set_ylabel(ylabel)
 
     if "cellYe" not in args.plotvars and "tracercount" not in args.plotvars:
-        cbar.set_label(r"log10($\rho$) [g/cm3]" if args.logcolorscale else r"$\rho$ [g/cm3]")
+        cbar.set_label(r"log10($\rho$ [g/cm3])" if args.logcolorscale else r"$\rho$ [g/cm3]")
+    else:
+        cbar.set_label("Ye" if "cellYe" in args.plotvars else "tracercount")
 
     defaultfilename = Path(modelpath) / f"plotcomposition_{','.join(args.plotvars)}.pdf"
     if args.outputfile and Path(args.outputfile).is_dir():

From 2730eeb0c3aab28cffd433d4df8a026cdb5e141d Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Tue, 14 Nov 2023 23:34:15 +0100
Subject: [PATCH 149/150] Update plotestimators.py

---
 artistools/estimators/plotestimators.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py
index b4f486d5c..5ebf2b2de 100755
--- a/artistools/estimators/plotestimators.py
+++ b/artistools/estimators/plotestimators.py
@@ -485,7 +485,7 @@ def get_iontuple(ionstr):
             ymin = max(ymin, ymax / 1e10)
             ax.set_ylim(bottom=ymin)
             # make space for the legend
-            new_ymax = ymax * 10 ** (0.3 * math.log10(ymax / ymin))
+            new_ymax = ymax * 10 ** (0.1 * math.log10(ymax / ymin))
             if ymin > 0 and new_ymax > ymin and np.isfinite(new_ymax):
                 ax.set_ylim(top=new_ymax)
 
@@ -826,7 +826,7 @@ def make_plot(
         outfilename = str(args.outputfile).format(timestep=timestep, timedays=timedays, format=args.format)
 
     if not args.notitle:
-        axes[0].set_title(figure_title, fontsize=8)
+        axes[0].set_title(figure_title, fontsize=12)
 
     print(f"Saving {outfilename} ...")
     fig.savefig(outfilename, dpi=300)
@@ -1057,7 +1057,7 @@ def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None
         # [['initmasses', ['Ni_56', 'He', 'C', 'Mg']]],
         # ['heating_gamma/gamma_dep'],
         # ["nne", ["_ymin", 1e5], ["_ymax", 1e10]],
-        ["rho", ["_yscale", "log"]],
+        ["rho", ["_yscale", "log"], ["_ymin", 1e-16]],
         ["TR", ["_yscale", "linear"]],  # , ["_ymin", 1000], ["_ymax", 15000]
         # ["Te"],
         # ["Te", "TR"],

From 79b638b0b9a19a72c206aeb217db523688ad758e Mon Sep 17 00:00:00 2001
From: Luke Shingles <luke.shingles@gmail.com>
Date: Tue, 14 Nov 2023 23:55:44 +0100
Subject: [PATCH 150/150] Update test_spectra.py

---
 artistools/spectra/test_spectra.py | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/artistools/spectra/test_spectra.py b/artistools/spectra/test_spectra.py
index 97062e8d1..d9c579856 100755
--- a/artistools/spectra/test_spectra.py
+++ b/artistools/spectra/test_spectra.py
@@ -170,16 +170,16 @@ def test_spectra_get_spectrum_polar_angles_frompackets() -> None:
     print(f"expected_results = {results_pkts!r}")
 
     expected_results = {
-        0: (4.3529439120747186e-12, 1.0314066492861092e-11),
-        10: (3.780678513149676e-12, 9.529723964592956e-12),
-        20: (4.4246323512365474e-12, 1.0166369617322156e-11),
-        30: (3.851546304731156e-12, 9.243745813551232e-12),
-        40: (4.067455987977417e-12, 9.994481885245608e-12),
-        50: (4.062094857762276e-12, 9.823422691843564e-12),
-        60: (3.858054725612363e-12, 9.157894155554632e-12),
-        70: (3.997110745634441e-12, 9.534252566271398e-12),
-        80: (4.121413251153407e-12, 9.480857140587749e-12),
-        90: (4.299536899380859e-12, 9.957108958110141e-12),
+        0: (4.353162807671065e-12, 1.0314585154204157e-11),
+        10: (3.780868631353459e-12, 9.530203183864417e-12),
+        20: (4.4248548518147095e-12, 1.016688085146278e-11),
+        30: (3.851739986649016e-12, 9.244210651898158e-12),
+        40: (4.067660527301169e-12, 9.994984475703157e-12),
+        50: (4.062299127491974e-12, 9.823916680282592e-12),
+        60: (3.858248734817849e-12, 9.158354676696867e-12),
+        70: (3.997311747521441e-12, 9.53473201327172e-12),
+        80: (4.121620503814969e-12, 9.481333902503268e-12),
+        90: (4.29975310930973e-12, 9.95760966920298e-12),
     }
 
     for dirbin in results_pkts: