Add more functionality for rimscode website creation (#10)

* bump version, add config reading to init
* add table creator for config parser
* Add way to overwrite darkmode of config

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rimsschemedrawer"
-version = "3.0.0.dev2"
+version = "3.0.0.dev3"
 description = "Drawing of RIMS schemes in python and/or with a Python GUI."
 authors = [
     { name = "Reto Trappitsch", email = "[email protected]" }

src/rimsschemedrawer/__init__.py

@@ -1,5 +1,6 @@
 # Init package
 
-from .plotter import Plotter
+from rimsschemedrawer.json_parser import json_reader, ConfigParser
+from rimsschemedrawer.plotter import Plotter
 
-__all__ = ["Plotter"]
+__all__ = ["ConfigParser", "json_reader", "Plotter"]

src/rimsschemedrawer/json_parser.py

@@ -2,7 +2,7 @@
 
 import json
 from pathlib import Path
-from typing import Dict, Tuple
+from typing import Dict, List, Tuple
 
 import numpy as np
 
@@ -60,6 +60,11 @@ def number_of_levels(self) -> int:
         """Get the number of steps in the scheme."""
         return self._num_steps
 
+    @property
+    def rims_scheme(self) -> Dict:
+        """Return the full scheme parsed into a dictionary."""
+        return
+
     @property
     def step_levels(self) -> np.ndarray:
         """Get all levels of the scheme."""
@@ -155,6 +160,139 @@ def sett_title(self) -> str:
         """Get the title for the plot."""
         return self._sett_title
 
+    # METHODS
+
+    def scheme_table(self, prec: int = 3, prec_strength: int = 1) -> Tuple[List, List]:
+        """Create a scheme table for further processing.
+
+        The headers are the following:
+        - Step
+        - λ (nm)
+        - From (cm⁻¹)  - including term symbol formatted
+        - To (cm⁻¹) - including term symbol formatted
+        - Forbidden - filled in as "x" if forbidden (only if there are forbidden steps)
+        - Strength (s⁻¹) - formatted as LaTeX string to given precision
+            (only present if any transitions were given)
+
+        Table: All entries are formatted as strings!
+
+        :param prec: Precision for the wavelength and steps.
+        :param prec_strength: Precision for the transition strength.
+
+        :return: Tuple with headers and the scheme table.
+        """
+
+        def reshuffle_list_low_lying(lst: List) -> List:
+            """Reshuffle a given list when low-lying states are present.
+
+            Low-lying states are stored first, then the step from the ground state, and
+            then the actual steps. This function reshuffles the list to have the step
+            from the ground state first, then the low-lying states, and then the actual
+            steps. If no low-lying states are present, do nothing.
+
+            :param lst: List to reshuffle.
+            """
+            if np.any(self._low_lying):
+                idx_first_step = np.where(~self._low_lying)[0][0]
+                tmp_value = lst.pop(idx_first_step)
+                lst.insert(0, tmp_value)
+                return lst
+
+            return lst
+
+        headers = [
+            "Step",
+            "λ (nm)",
+            "From (cm⁻¹)",
+            "To (cm⁻¹)",
+        ]
+
+        first_no_lowlying = np.where(~self._low_lying)[0][0]
+
+        steps = np.ones(self.number_of_levels, dtype=int)
+        for it in range(first_no_lowlying + 1, len(steps)):
+            steps[it] += steps[it - 1]
+
+        lambdas = [f"{it:.{prec}f}" for it in self._steps_nm]
+        lambdas = reshuffle_list_low_lying(lambdas)
+
+        from_level = ["" for _ in range(self.number_of_levels)]
+        # add low-lying states
+        for it in range(first_no_lowlying):
+            tmp_str = f"{self.step_levels[it]:.{prec}f}"
+            if term := self.step_terms[it]:
+                tmp_str += f" ({term})"
+            from_level[it] = tmp_str
+        # add ground state
+        gs = self.gs_level
+        if gs == 0:
+            tmp_str = "0"
+        else:
+            tmp_str = f"{gs:.{prec}f}"
+        if term := self.gs_term:
+            tmp_str += f" ({term})"
+        from_level[first_no_lowlying] = tmp_str
+        # add steps
+        for it in range(
+            first_no_lowlying + 1, self.number_of_levels
+        ):  # above ground state
+            tmp_str = f"{self.step_levels[it - 1]:.{prec}f}"
+            if term := self.step_terms[it - 1]:
+                tmp_str += f" ({term})"
+            from_level[it] = tmp_str
+        from_level = reshuffle_list_low_lying(from_level)
+
+        to_level = ["" for _ in range(self.number_of_levels)]
+        # add low-lying states and first step
+        for it in range(first_no_lowlying + 1):
+            tmp_str = f"{self.step_levels[first_no_lowlying]:.{prec}f}"
+            if term := self.step_terms[first_no_lowlying]:
+                tmp_str += f" ({term})"
+            to_level[it] = tmp_str
+        # add steps
+        for it in range(first_no_lowlying + 1, self.number_of_levels):
+            tmp_str = f"{self.step_levels[it]:.{prec}f}"
+            if term := self.step_terms[it]:
+                tmp_str += f" ({term})"
+            to_level[it] = tmp_str
+        to_level = reshuffle_list_low_lying(to_level)
+
+        # forbidden transitions
+        if np.any(self.step_forbidden):
+            headers.append("Forbidden")
+            forbidden = ["x" if it else "" for it in self.step_forbidden]
+            forbidden = reshuffle_list_low_lying(forbidden)
+
+        # transition strength
+        if np.any(self.transition_strengths):
+            transition_strengths = []
+            headers.append("Strength (s⁻¹)")
+            for val in self.transition_strengths:
+                if val != 0:
+                    transition_strengths.append(ut.my_exp_formatter(val, prec_strength))
+                else:
+                    transition_strengths.append("")
+            transition_strengths = reshuffle_list_low_lying(transition_strengths)
+
+        # create table
+        table = []
+        for idx in range(self.number_of_levels):
+            row = [
+                str(steps[idx]),
+                lambdas[idx],
+                from_level[idx],
+                to_level[idx],
+            ]
+            if np.any(self.step_forbidden):
+                row.append(forbidden[idx])
+            if np.any(self.transition_strengths):
+                row.append(transition_strengths[idx])
+            table.append(row)
+
+        return headers, table
+
+    # PRIVATE METHODS
+
     def _parse_data_scheme(self):
         """Parse the data of the scheme and save it to class variables."""
         # variable that defines if input is in nm (True). Otherwise in cm^-1 (False)

src/rimsschemedrawer/plotter.py

@@ -18,6 +18,7 @@ def __init__(self, data: dict, **kwargs):
                 This number can be higher than the number of available steps!
             - fig_ax: Tuple of matplotlib figure and axes to plot on. Defaults to
                 creating new ones.
+            - darkmode: Overwrite the darkmode settings from the config file.
         """
         self.config_parser = ConfigParser(data)
 
@@ -30,13 +31,15 @@ def __init__(self, data: dict, **kwargs):
         matplotlib.rc("xtick", labelsize=fsz_axes, direction="in")
         matplotlib.rc("ytick", labelsize=fsz_axes, direction="in")
 
+        darkmode = kwargs.get("darkmode", self.config_parser.sett_darkmode)
+
         # figure stuff
-        if self.config_parser.sett_darkmode:
+        if darkmode:
             plt.style.use("dark_background")
         self._figure, self._axes = kwargs.get("fig_ax", plt.subplots(1, 1))
 
         # Colors for arrows
-        if self.config_parser.sett_darkmode:
+        if darkmode:
             self.colmain = "#ffffff"
             self.colir = "#d27878"
             self.coluv = "#8c96df"

tests/data/raised_ground_cm.json

@@ -5,7 +5,7 @@
         "ip_level": "63427.70",
         "ip_term": "",
         "step_forbidden0": false,
-        "step_forbidden1": false,
+        "step_forbidden1": true,
         "step_forbidden2": false,
         "step_forbidden3": false,
         "step_forbidden4": false,
@@ -32,9 +32,9 @@
         "step_term4": "",
         "step_term5": "",
         "step_term6": "",
-        "trans_strength0": "",
+        "trans_strength0": "1e6",
         "trans_strength1": "",
-        "trans_strength2": "",
+        "trans_strength2": "3.015e5",
         "trans_strength3": "",
         "trans_strength4": "",
         "trans_strength5": "",

tests/test_scheme_table.py

@@ -0,0 +1,81 @@
+# Tests for ConfigParser scheme table creation.
+
+import rimsschemedrawer.json_parser as jp
+from rimsschemedrawer.utils import term_to_string as tts
+
+
+def test_ti_new_json(data_path):
+    """Return the correct table for the new Ti scheme."""
+    json_file = data_path.joinpath("ti_new.json")
+    data = jp.json_reader(json_file)
+
+    parser = jp.ConfigParser(data)
+
+    header_exp = [
+        "Step",
+        "λ (nm)",
+        "From (cm⁻¹)",
+        "To (cm⁻¹)",
+    ]
+
+    table_exp = [
+        ["1", "465.777", f"0 ({tts('3F2')})", f"21469.500 ({tts('3G3')})"],
+        ["1", "469.498", f"170.150 ({tts('3F3')})", f"21469.500 ({tts('3G3')})"],
+        ["1", "474.324", f"386.880 ({tts('3F4')})", f"21469.500 ({tts('3G3')})"],
+        ["2", "416.158", f"21469.500 ({tts('3G3')})", f"45498.850 ({tts('3G4')})"],
+        ["3", "881.399", f"45498.850 ({tts('3G4')})", "56844.450"],
+    ]
+
+    header, table = parser.scheme_table()
+
+    assert header == header_exp
+    assert table == table_exp
+
+
+def test_raised_ground(data_path):
+    """Return the correct table for a scheme with a raised ground level."""
+    json_file = data_path.joinpath("raised_ground_cm.json")
+    data = jp.json_reader(json_file)
+
+    parser = jp.ConfigParser(data)
+
+    header_exp = [
+        "Step",
+        "λ (nm)",
+        "From (cm⁻¹)",
+        "To (cm⁻¹)",
+        "Forbidden",
+        "Strength (s⁻¹)",
+    ]
+
+    table_exp = [
+        [
+            "1",
+            "400.262",
+            f"1000.000 ({tts('5D0')})",
+            f"25983.609 ({tts('5F1')})",
+            "",
+            "$1.0 \\times 10^{6}$",
+        ],
+        [
+            "2",
+            "407.469",
+            f"25983.609 ({tts('5F1')})",
+            f"50525.354 ({tts('5D2')})",
+            "x",
+            "",
+        ],
+        [
+            "3",
+            "771.908",
+            f"50525.354 ({tts('5D2')})",
+            f"63480.266 ({tts('AI')})",
+            "",
+            "$3.0 \\times 10^{5}$",
+        ],
+    ]
+
+    header, table = parser.scheme_table()
+
+    assert header == header_exp
+    assert table == table_exp