Move and Rename SDECData to VisualizationData

tardis/visualization/plot_util.py

@@ -1,8 +1,13 @@
 """Utility functions to be used in plotting."""
 
 import re
+
 import numpy as np
 
+from tardis.visualization.tools.simulation_packet_data import (
+    SimulationPacketData,
+)
+
 
 def axis_label_in_latex(label_text, unit, only_text=True):
     """
@@ -79,3 +84,63 @@
     """
     color_tuple_255 = tuple([int(x * 255) for x in color_tuple[:3]])
     return f"rgb{color_tuple_255}"
+
+def create_packet_data_dict_from_simulation(sim):
+    """
+    Create a dictionary containing virtual and real packet data based on simulation state.
+
+    Parameters
+    ----------
+    sim : tardis.simulation.Simulation
+        TARDIS Simulation object produced by running a simulation
+
+    Returns
+    -------
+    dict
+        Dictionary containing 'virtual' and 'real' SimulationPacketData instances
+    """
+    packet_data = {
+        "real": SimulationPacketData.from_simulation(sim, "real")
+    }
+    if sim.transport.transport_state.virt_logging:
+        packet_data["virtual"] = SimulationPacketData.from_simulation(sim, "virtual")
+    else:
+        packet_data["virtual"] = None
+
+    return packet_data
+
+def create_packet_data_dict_from_hdf(hdf_fpath, packets_mode=None):
+    """
+    Create a dictionary containing virtual and real packet data from HDF file.
+
+    Parameters
+    ----------
+    hdf_fpath : str
+        Valid path to the HDF file where simulation is saved
+    packets_mode : {'virtual', 'real', None}
+        Mode of packets to be considered. If None, both modes are returned.
+
+    Returns
+    -------
+    dict
+        Dictionary containing 'virtual' and 'real' SimulationPacketData instances
+    """
+    if packets_mode not in [None, "virtual", "real"]:
+        raise ValueError(
+            "Invalid value passed to packets_mode. Only "
+            "allowed values are 'virtual', 'real' or None"
+        )
+    if packets_mode == "virtual":
+        return {
+            "virtual": SimulationPacketData.from_hdf(hdf_fpath, "virtual"),
+            "real": None
+        }
+    if packets_mode == "real":
+        return {
+            "virtual": None,
+            "real": SimulationPacketData.from_hdf(hdf_fpath, "real")
+        }
+    return {
+        "virtual": SimulationPacketData.from_hdf(hdf_fpath, "virtual"),
+        "real": SimulationPacketData.from_hdf(hdf_fpath, "real")
+    }

tardis/visualization/tools/liv_plot.py

@@ -1,16 +1,17 @@
 import logging
-import matplotlib.pyplot as plt
+
+import astropy.units as u
 import matplotlib.cm as cm
-import plotly.graph_objects as go
+import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
-import astropy.units as u
+import plotly.graph_objects as go
 
+import tardis.visualization.tools.sdec_plot as sdec
 from tardis.util.base import (
     atomic_number2element_symbol,
     int_to_roman,
 )
-import tardis.visualization.tools.sdec_plot as sdec
 from tardis.visualization import plot_util as pu
 
 logger = logging.getLogger(__name__)
@@ -37,7 +38,6 @@
         velocity : astropy.units.Quantity
             Velocity array from the simulation.
         """
-
         self.data = data
         self.time_explosion = time_explosion
         self.velocity = velocity
@@ -57,12 +57,8 @@
         -------
         LIVPlotter
         """
-
         return cls(
-            dict(
-                virtual=sdec.SDECData.from_simulation(sim, "virtual"),
-                real=sdec.SDECData.from_simulation(sim, "real"),
-            ),
+            pu.create_packet_data_dict_from_simulation(sim),
             sim.plasma.time_explosion,
             sim.simulation_state.velocity,
         )
@@ -90,14 +86,11 @@
             velocity = pd.concat(
                 [v_inner, pd.Series([v_outer.iloc[-1]])], ignore_index=True
             ).tolist() * (u.cm / u.s)
-            return cls(
-                dict(
-                    virtual=sdec.SDECData.from_hdf(hdf_fpath, "virtual"),
-                    real=sdec.SDECData.from_hdf(hdf_fpath, "real"),
-                ),
-                time_explosion,
-                velocity,
-            )
+        return cls(
+            pu.create_packet_data_dict_from_hdf(hdf_fpath),
+            time_explosion,
+            velocity,
+        )
 
     def _parse_species_list(self, species_list, packets_mode, nelements=None):
         """