Update spencerfano.py

lukeshingles · Oct 15, 2024 · 269a728 · 269a728
1 parent ccf1ff7
commit 269a728
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/pynonthermal/spencerfano.py b/pynonthermal/spencerfano.py
@@ -41,7 +41,7 @@ class SpencerFanoSolver:
     sourcevec: npt.NDArray[np.float64]
     E_init_ev: float
     sfmatrix: npt.NDArray[np.float64]
-    adata_polars: pd.DataFrame
+    adata_polars: pd.DataFrame | None
 
     def __init__(
         self,
@@ -177,6 +177,8 @@ def add_ion_ltepopexcitation(self, Z, ion_stage, n_ion, temperature=3000, adata_
                 derived_transitions_columns=["epsilon_trans_ev", "lambda_angstroms", "lower_g", "upper_g"],
             )
 
+        assert self.adata_polars is not None
+
         ion = self.adata_polars.filter(pl.col("Z") == Z).filter(pl.col("ion_stage") == ion_stage)
 
         dfpops_thision = ion["levels"].item()
@@ -732,6 +734,7 @@ def plot_yspectrum(
     ) -> None:
         assert self._solved
         fs = 12
+        fig = None
         if axis is None:
             fig, axis = plt.subplots(
                 nrows=1,
@@ -761,6 +764,7 @@ def plot_yspectrum(
         if axis is None:
             if outputfilename is not None:
                 print(f"Saving '{outputfilename}'")
+                assert fig is not None
                 fig.savefig(str(outputfilename))
                 plt.close()
             else:
@@ -769,6 +773,7 @@ def plot_yspectrum(
     def plot_channels(self, outputfilename=None, axis=None, xscalelog=False) -> None:
         assert self._solved
         fs = 12
+        fig = None
         if axis is None:
             fig, axis = plt.subplots(
                 nrows=1,
@@ -903,6 +908,7 @@ def plot_channels(self, outputfilename=None, axis=None, xscalelog=False) -> None
         if axis is None:
             if outputfilename is not None:
                 print(f"Saving '{outputfilename}'")
+                assert fig is not None
                 fig.savefig(str(outputfilename))
                 plt.close()
             else: