Fixed to the montecarlo_numba file to compy w/ issue tardis-sn#680

atharva-2001 · Dec 2, 2020 · e28936b · e28936b
1 parent c671075
commit e28936b
Show file tree

Hide file tree

Showing 8 changed files with 84 additions and 101 deletions.
diff --git a/tardis/montecarlo/base.py b/tardis/montecarlo/base.py
@@ -140,13 +140,8 @@ def _initialize_estimator_arrays(self, tau_sobolev_shape):
 
         Parameters
         ----------
-<<<<<<< HEAD
-
-        tau_sobolev_shape: tuple
+        tau_sobolev_shape : tuple
             tuple for the tau_sobolev_shape
-=======
-        model : Radial1DModel
->>>>>>> 56c506760fd38eb20d9ab479bf71e38c7ee4c43a
         """
 
         # Estimators
@@ -462,7 +457,6 @@ def calculate_radiationfield_properties(self):
         -------
         t_rad : astropy.units.Quantity (float)
         w : numpy.ndarray (float)
-
         """
 
         t_rad = (

diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py
@@ -110,14 +110,14 @@ def montecarlo_main_loop(
 
     Parameters
     ----------
-    packet_collection: PacketCollection
-    numba_model: NumbaModel
-    estimators: NumbaEstimators
-    spectrum_frequency: astropy.units.Quantity
+    packet_collection : PacketCollection
+    numba_model : NumbaModel
+    estimators : NumbaEstimators
+    spectrum_frequency : astropy.units.Quantity
         frequency bins
-    number_of_vpackets: int
+    number_of_vpackets : int
         VPackets released per interaction
-    packet_seeds: numpy.array
+    packet_seeds : numpy.array
     """
     output_nus = np.empty_like(packet_collection.packets_output_nu)
     last_interaction_types = (

diff --git a/tardis/montecarlo/montecarlo_numba/interaction.py b/tardis/montecarlo/montecarlo_numba/interaction.py
@@ -30,9 +30,8 @@ def thomson_scatter(r_packet, time_explosion):
     Parameters
     ----------
     r_packet : RPacket
-    time_explosion: float
+    time_explosion : float
         time since explosion in seconds
-
     """
     old_doppler_factor = get_doppler_factor(
         r_packet.r, r_packet.mu, time_explosion
@@ -56,10 +55,13 @@ def thomson_scatter(r_packet, time_explosion):
 def line_scatter(r_packet, time_explosion, line_interaction_type, numba_plasma):
     """
     Line scatter function that handles the scattering itself, including new angle drawn, and calculating nu out using macro atom
-    r_packet: RPacket
-    time_explosion: float
-    line_interaction_type: enum
-    numba_plasma: NumbaPlasma
+    
+    Parameters
+    ----------
+    r_packet : RPacket
+    time_explosion : float
+    line_interaction_type : enum
+    numba_plasma : NumbaPlasma
     """
 
     old_doppler_factor = get_doppler_factor(
@@ -89,12 +91,11 @@ def line_emission(r_packet, emission_line_id, time_explosion, numba_plasma):
     Sets the frequency of the RPacket properly given the emission channel
 
     Parameters
-    -----------
-
-    r_packet: RPacket
-    emission_line_id: int
-    time_explosion: float
-    numba_plasma: NumbaPlasma
+    ----------
+    r_packet : RPacket
+    emission_line_id : int
+    time_explosion : float
+    numba_plasma : NumbaPlasma
     """
     r_packet.last_line_interaction_out_id = emission_line_id
 

diff --git a/tardis/montecarlo/montecarlo_numba/macro_atom.py b/tardis/montecarlo/montecarlo_numba/macro_atom.py
@@ -21,15 +21,13 @@ class MacroAtomTransitionType(IntEnum):
 @njit(**njit_dict)
 def macro_atom(r_packet, numba_plasma):
     """
-
     Parameters
     ----------
-    r_packet: tardis.montecarlo.montecarlo_numba.r_packet.RPacket
-    numba_plasma: tardis.montecarlo.numba_interface.numba_plasma
+    r_packet : tardis.montecarlo.montecarlo_numba.r_packet.RPacket
+    numba_plasma : tardis.montecarlo.numba_interface.numba_plasma
 
     Returns
     -------
-
     """
     activation_level_id = numba_plasma.line2macro_level_upper[
         r_packet.next_line_id

diff --git a/tardis/montecarlo/montecarlo_numba/numba_interface.py b/tardis/montecarlo/montecarlo_numba/numba_interface.py
@@ -29,9 +29,9 @@ def __init__(self, r_inner, r_outer, time_explosion):
 
         Parameters
         ----------
-        r_inner: numpy.ndarray
-        r_outer: numpy.ndarray
-        time_explosion: float
+        r_inner : numpy.ndarray
+        r_outer : numpy.ndarray
+        time_explosion : float
         """
         self.r_inner = r_inner
         self.r_outer = r_outer
@@ -67,18 +67,18 @@ def __init__(
     ):
         """
         Plasma for the Numba code
+        
         Parameters
         ----------
-
-        electron_density: numpy.array
-        line_list_nu: numpy.array
-        tau_sobolev: numpy.array
-        transition_probabilities: numpy.array
-        line2macro_level_upper: numpy.array
-        macro_block_references: numpy.array
-        transition_type: numpy.array
-        destination_level_id: numpy.array
-        transition_line_id: numpy.array
+        electron_density : numpy.array
+        line_list_nu : numpy.array
+        tau_sobolev : numpy.array
+        transition_probabilities : numpy.array
+        line2macro_level_upper : numpy.array
+        macro_block_references : numpy.array
+        transition_type : numpy.array
+        destination_level_id : numpy.array
+        transition_line_id : numpy.array
         """
 
         self.electron_density = electron_density
@@ -102,10 +102,9 @@ def numba_plasma_initialize(plasma, line_interaction_type):
     Initialize the NumbaPlasma object and copy over the data over from TARDIS Plasma
 
     Parameters
-    -----------
-
-    plasma: tardis.plasma.BasePlasma
-    line_interaction_type: enum
+    ----------
+    plasma : tardis.plasma.BasePlasma
+    line_interaction_type : enum
     """
     electron_densities = plasma.electron_densities.values
     line_list_nu = plasma.atomic_data.lines.nu.values

diff --git a/tardis/montecarlo/montecarlo_numba/r_packet.py b/tardis/montecarlo/montecarlo_numba/r_packet.py
@@ -83,16 +83,15 @@ def calculate_distance_boundary(r, mu, r_inner, r_outer):
     """
     Calculate distance to shell boundary in cm.
 
-    Parameters:
-    ------------
-
-    r: float
+    Parameters
+    ----------
+    r : float
        radial coordinate of the RPacket
-    mu: float
+    mu : float
        cosine of the direction of movement
-    r_inner: float
+    r_inner : float
        inner radius of current shell
-    r_outer: float
+    r_outer : float
        outer radius of current shell
     """
 
@@ -129,21 +128,21 @@ def calculate_distance_line(
 ):
     """
     Calculate distance until RPacket is in resonance with the next line
+    
     Parameters
     ----------
-    r_packet: RPacket
-    comov_nu: float
+    r_packet : RPacket
+    comov_nu : float
         comoving frequency at the CURRENT position of the RPacket
-    nu_line: float
+    nu_line : float
         line to check the distance to
-    time_explosion: float
+    time_explosion : float
         time since explosion in seconds
-    is_last_line: bool
+    is_last_line : bool
         return MISS_DISTANCE if at the end of the line list
 
     Returns
     -------
-
     """
 
     nu = r_packet.nu
@@ -204,8 +203,10 @@ def calculate_distance_electron(electron_density, tau_event):
     """
     Calculate distance to Thomson Scattering
 
-    electron_density: float
-    tau_event: float
+    Parameters
+    ----------
+    electron_density : float
+    tau_event : float
     """
     # add full_relativity here
     return tau_event / (electron_density * numba_config.SIGMA_THOMSON)
@@ -216,11 +217,10 @@ def calculate_tau_electron(electron_density, distance):
     """
     Calculate tau for Thomson scattering
 
-    Parameters:
-    ------------
-
-    electron_density: float
-    distance: float
+    Parameters
+    ----------
+    electron_density : float
+    distance : float
     """
     return electron_density * numba_config.SIGMA_THOMSON * distance
 
@@ -251,12 +251,11 @@ def get_inverse_doppler_factor(r, mu, time_explosion):
     """
     Calculate doppler factor for frame transformation
 
-    Parameters:
-    ------------
-
-    r: float
-    mu: float
-    time_explosion: float
+    Parameters
+    ----------
+    r : float
+    mu : float
+    time_explosion : float
     """
     inv_c = 1 / C_SPEED_OF_LIGHT
     inv_t = 1 / time_explosion
@@ -291,12 +290,11 @@ def update_line_estimators(
 
     Parameters
     ----------
-    estimators: Estimators
-    r_packet: RPacket
-    cur_line_id: int
-    distance_trace: float
-    time_explosion: float
-
+    estimators : Estimators
+    r_packet : RPacket
+    cur_line_id : int
+    distance_trace : float
+    time_explosion : float
     """
 
     """ Actual calculation - simplified below
@@ -341,16 +339,15 @@ def trace_packet(r_packet, numba_model, numba_plasma, estimators):
     """
     Traces the RPacket through the ejecta and stops when an interaction happens (heart of the calculation)
 
-    r_packet: RPacket
     Parameters
     ----------
-    numba_model: tardis.montecarlo.montecarlo_numba.numba_interface.NumbaModel
-    numba_plasma: tardis.montecarlo.montecarlo_numba.numba_interface.NumbaPlasma
-    estimators: tardis.montecarlo.montecarlo_numba.numba_interface.Estimators
+    r_packet : RPacket
+    numba_model : tardis.montecarlo.montecarlo_numba.numba_interface.NumbaModel
+    numba_plasma : tardis.montecarlo.montecarlo_numba.numba_interface.NumbaPlasma
+    estimators : tardis.montecarlo.montecarlo_numba.numba_interface.Estimators
 
     Returns
     -------
-
     """
 
     r_inner = numba_model.r_inner[r_packet.current_shell_id]
@@ -499,12 +496,11 @@ def move_r_packet(r_packet, distance, time_explosion, numba_estimator):
 
     Parameters
     ----------
-
-    r_packet: tardis.montecarlo.montecarlo_numba.r_packet.RPacket
+    r_packet : tardis.montecarlo.montecarlo_numba.r_packet.RPacket
         r_packet objects
-    time_explosion: float
+    time_explosion : float
         time since explosion in s
-    numba_estimator: tardis.montecarlo.montecarlo_numba.numba_interface.NumbaEstimator
+    numba_estimator : tardis.montecarlo.montecarlo_numba.numba_interface.NumbaEstimator
         Estimators object
     distance : float
         distance in cm
@@ -577,10 +573,10 @@ def move_packet_across_shell_boundary(packet, delta_shell, no_of_shells):
     distance : float
         distance to move to shell boundary
 
-    delta_shell: int
+    delta_shell : int
         is +1 if moving outward or -1 if moving inward
 
-    no_of_shells: int
+    no_of_shells : int
         number of shells in TARDIS simulation
     """
     next_shell_id = packet.current_shell_id + delta_shell

diff --git a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py
@@ -36,14 +36,13 @@ def single_packet_loop(
     r_packet, numba_model, numba_plasma, estimators, vpacket_collection
 ):
     """
-
     Parameters
     ----------
-    r_packet: tardis.montecarlo.montecarlo_numba.r_packet.RPacket
-    numba_model: tardis.montecarlo.montecarlo_numba.numba_interface.NumbaModel
-    numba_plasma: tardis.montecarlo.montecarlo_numba.numba_interface.NumbaPlasma
-    estimators: tardis.montecarlo.montecarlo_numba.numba_interface.Estimators
-    vpacket_collection: tardis.montecarlo.montecarlo_numba.numba_interface.VPacketCollection
+    r_packet : tardis.montecarlo.montecarlo_numba.r_packet.RPacket
+    numba_model : tardis.montecarlo.montecarlo_numba.numba_interface.NumbaModel
+    numba_plasma : tardis.montecarlo.montecarlo_numba.numba_interface.NumbaPlasma
+    estimators : tardis.montecarlo.montecarlo_numba.numba_interface.Estimators
+    vpacket_collection : tardis.montecarlo.montecarlo_numba.numba_interface.VPacketCollection
 
     Returns
     -------

diff --git a/tardis/montecarlo/packet_source.py b/tardis/montecarlo/packet_source.py
@@ -38,15 +38,9 @@ def create_zero_limb_darkening_packet_mus(no_of_packets, rng):
     @staticmethod
     def create_uniform_packet_energies(no_of_packets, rng):
         """
-<<<<<<< HEAD
         Uniformly distribute energy in arbitrary units where the ensemble of
         packets has energy of 1.
 
-=======
-        Uniformly distribute energy in arbitrary units where the ensemble of 
-        packets has energy of 1. 
-        
->>>>>>> 56c506760fd38eb20d9ab479bf71e38c7ee4c43a
         Parameters
         ----------
         no_of_packets : int
@@ -79,13 +73,15 @@ def create_blackbody_packet_nus(T, no_of_packets, rng, l_samples=1000):
         .. math::
             x = -\\ln{(\\xi_1\\xi_2\\xi_3\\xi_4)}/l_{\\rm min}\\; .
         where :math:`x=h\\nu/kT`
+        
         Parameters
         ----------
         T : float
             temperature
         no_of_packets : int
         l_samples : int
             number of l_samples needed in the algorithm
+            
         Returns
         -------
             : numpy.ndarray