Speed up estimators reading and plotting

.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
 

.vscode/extensions.json

@@ -1,6 +1,7 @@
 {
     "recommendations": [
         "ms-python.mypy-type-checker",
-        "charliermarsh.ruff"
+        "charliermarsh.ruff",
+        "sourcery.sourcery"
     ]
 }

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
@@ -27,6 +26,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
@@ -60,6 +60,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
@@ -84,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
@@ -100,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
 
 

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:

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])
@@ -140,30 +137,30 @@ 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
 
                     tsmgi = (cur_timestep, cur_modelgridindex)
-                    if "populations" not in estimators[tsmgi]:
-                        estimators[tsmgi]["populations"] = {}
 
                     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]["populations"][(atomic_number, ion_stage)] = ionpop
-                                estimators[tsmgi]["populations"].setdefault(atomic_number, 0.0)
-                                estimators[tsmgi]["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]["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"]

artistools/commands.py

@@ -16,11 +16,11 @@
     "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"),
-    "plotmodeldensity": ("inputmodel.plotdensity", "main"),
-    "plotmodeldeposition": ("deposition", "main"),
+    "plotdensity": ("inputmodel.plotdensity", "main"),
+    "plotdeposition": ("deposition", "main"),
     "plotmacroatom": ("macroatom", "main"),
     "plotnltepops": ("nltepops.plotnltepops", "main"),
     "plotnonthermal": ("nonthermal.plotnonthermal", "main"),
@@ -58,36 +58,21 @@ 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",
-        ),
+        "plotartisdensity": ("artistools.inputmodel.plotdensity", "main"),
+        "plotartisdeposition": ("artistools.deposition", "main"),
+        "plotartisestimators": ("artistools.estimators.plotestimators", "main"),
+        "plotartislightcurve": ("artistools.lightcurve.plotlightcurve", "main"),
         "plotartislinefluxes": ("artistools.linefluxes", "main"),
         "plotartismacroatom": ("artistools.macroatom", "main"),
         "plotartisnltepops": ("artistools.nltepops.plotnltepops", "main"),
         "plotartisnonthermal": ("artistools.nonthermal", "main"),
         "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.inputmodel.plotinitialcomposition", "main"),
+        "plotartisviewingangles": ("artistools.viewing_angles_visualization", "main"),
     }
 
 

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():

artistools/estimators/__init__.py

@@ -6,16 +6,14 @@
 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
-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
-from artistools.estimators.estimators import parse_estimfile
+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
 from artistools.estimators.plotestimators import addargs
 from artistools.estimators.plotestimators import main as plot