From 3164f609f7d38e754684cb879790d0be9d4f1d32 Mon Sep 17 00:00:00 2001
From: Andrew Fullard <andrewgfullard@gmail.com>
Date: Wed, 12 Jun 2024 16:58:12 -0400
Subject: [PATCH] Fixes packet source re-init for custom sources

Also adjusts notebook to deal with quantity-based packet frequencies
---
 docs/io/optional/how_to_custom_source.ipynb   |  3 +--
 .../parse_packet_source_configuration.py      | 21 ++-----------------
 tardis/model/base.py                          |  4 +++-
 tardis/simulation/base.py                     |  6 ++----
 4 files changed, 8 insertions(+), 26 deletions(-)

diff --git a/docs/io/optional/how_to_custom_source.ipynb b/docs/io/optional/how_to_custom_source.ipynb
index 99bd5da7915..046faf00780 100644
--- a/docs/io/optional/how_to_custom_source.ipynb
+++ b/docs/io/optional/how_to_custom_source.ipynb
@@ -131,7 +131,6 @@
     "        truncation_frequency = (\n",
     "            u.Quantity(self.truncation_wavelength, u.Angstrom)\n",
     "            .to(u.Hz, equivalencies=u.spectral())\n",
-    "            .value\n",
     "        )\n",
     "\n",
     "        # Draw nus from blackbody distribution and reject based on truncation_frequency.\n",
@@ -220,7 +219,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.5"
+   "version": "3.12.3"
   }
  },
  "nbformat": 4,
diff --git a/tardis/io/model/parse_packet_source_configuration.py b/tardis/io/model/parse_packet_source_configuration.py
index 23394e73949..90aab02e556 100644
--- a/tardis/io/model/parse_packet_source_configuration.py
+++ b/tardis/io/model/parse_packet_source_configuration.py
@@ -6,9 +6,7 @@
 )
 
 
-def initialize_packet_source(
-    config, geometry, packet_source, legacy_mode_enabled
-):
+def initialize_packet_source(packet_source, config, geometry):
     """
     Initialize the packet source based on config and geometry
 
@@ -31,18 +29,6 @@ def initialize_packet_source(
     ValueError
         If both t_inner and luminosity_requested are None.
     """
-    if config.montecarlo.enable_full_relativity:
-        packet_source = BlackBodySimpleSourceRelativistic(
-            base_seed=config.montecarlo.seed,
-            time_explosion=config.supernova.time_explosion,
-            legacy_mode_enabled=legacy_mode_enabled,
-        )
-    else:
-        packet_source = BlackBodySimpleSource(
-            base_seed=config.montecarlo.seed,
-            legacy_mode_enabled=legacy_mode_enabled,
-        )
-
     luminosity_requested = config.supernova.luminosity_requested
     if config.plasma.initial_t_inner > 0.0 * u.K:
         packet_source.radius = geometry.r_inner_active[0]
@@ -57,7 +43,6 @@ def initialize_packet_source(
         raise ValueError(
             "Both t_inner and luminosity_requested cannot be None."
         )
-
     return packet_source
 
 
@@ -89,6 +74,4 @@ def parse_packet_source_from_config(config, geometry, legacy_mode_enabled):
             legacy_mode_enabled=legacy_mode_enabled,
         )
 
-    return initialize_packet_source(
-        config, geometry, packet_source, legacy_mode_enabled
-    )
+    return initialize_packet_source(packet_source, config, geometry)
diff --git a/tardis/model/base.py b/tardis/model/base.py
index 95fc0a47411..776e2f920fd 100644
--- a/tardis/model/base.py
+++ b/tardis/model/base.py
@@ -291,7 +291,9 @@ def from_config(cls, config, atom_data, legacy_mode_enabled=False):
 
         geometry = parse_geometry_from_config(config, time_explosion)
 
-        composition, electron_densities = parse_composition_from_config(atom_data, config, time_explosion, geometry)
+        composition, electron_densities = parse_composition_from_config(
+            atom_data, config, time_explosion, geometry
+        )
 
         packet_source = parse_packet_source_from_config(
             config, geometry, legacy_mode_enabled
diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py
index eba875c7884..7efa377a314 100644
--- a/tardis/simulation/base.py
+++ b/tardis/simulation/base.py
@@ -679,12 +679,10 @@ def from_config(
                     atom_data=atom_data,
                     legacy_mode_enabled=legacy_mode_enabled,
                 )
+            # Override with custom packet source from function argument if present
             if packet_source is not None:
                 simulation_state.packet_source = initialize_packet_source(
-                    config,
-                    simulation_state.geometry,
-                    packet_source,
-                    legacy_mode_enabled,
+                    packet_source, config, simulation_state.geometry
                 )
         if "plasma" in kwargs:
             plasma = kwargs["plasma"]