Fix parsing of abundance tables from arrays

pyxsim/source_models/thermal_sources.py

@@ -2,8 +2,8 @@
 
 import numpy as np
 from more_itertools import chunked
-from soxs.constants import abund_tables, atomic_weights, elem_names, metal_elem
-from soxs.utils import parse_prng, regrid_spectrum
+from soxs.constants import atomic_weights, elem_names, metal_elem
+from soxs.utils import _parse_abund_table, parse_prng, regrid_spectrum
 from unyt.array import unyt_quantity
 from yt.data_objects.static_output import Dataset
 from yt.utilities.exceptions import YTFieldNotFound
@@ -93,7 +93,8 @@ def __init__(
         self.pbar = None
         self.Zconvert = 1.0
         self.mconvert = {}
-        self.atable = abund_tables[abund_table].copy()
+        self.abund_table = abund_table
+        self.atable = _parse_abund_table(abund_table)
         if h_fraction is None:
             h_fraction = compute_H_abund(abund_table)
         self.h_fraction = h_fraction

pyxsim/utils.py

@@ -172,10 +172,26 @@ def merge_files(input_files, output_file, overwrite=False, add_exposure_times=Fa
     f_out.close()
 
 
+def _parse_abund_table(abund_table):
+    if not isinstance(abund_table, str):
+        if len(abund_table) != 30:
+            raise RuntimeError(
+                "User-supplied abundance tables must be 30 elements long!"
+            )
+        atable = np.concatenate([[0.0], np.array(abund_table)])
+    else:
+        if abund_table not in abund_tables:
+            raise KeyError(
+                f"Abundance table {abund_table} not found! Options are: {list(abund_tables.keys())}"
+            )
+        atable = abund_tables[abund_table].copy()
+    return atable
+
+
 def compute_elem_mass_fraction(elem, abund_table="angr"):
     if isinstance(elem, str):
         elem = elem_names.index(elem)
-    atable = abund_tables[abund_table]
+    atable = _parse_abund_table(abund_table)
     mZ = (atomic_weights[3:] * atable[3:]).sum()
     mE = atomic_weights[elem] * atable[elem]
     return mE / mZ
@@ -219,19 +235,17 @@ def _metal_field(field, data):
 
 
 def compute_H_abund(abund_table):
-    if abund_table not in abund_tables:
-        raise KeyError(
-            f"Abundance table {abund_table} not found! Options are: {list(abund_tables.keys())}"
-        )
-    return atomic_weights[1] / (atomic_weights * abund_tables[abund_table]).sum()
+    atable = _parse_abund_table(abund_table)
+    return atomic_weights[1] / (atomic_weights * atable).sum()
 
 
 def compute_zsolar(abund_table):
+    atable = _parse_abund_table(abund_table)
     if abund_table not in abund_tables:
         raise KeyError(
             f"Abundance table {abund_table} not found! Options are: {list(abund_tables.keys())}"
         )
-    elems = atomic_weights * abund_tables[abund_table]
+    elems = atomic_weights * atable
     return elems[3:].sum() / elems.sum()