diff --git a/pyproject.toml b/pyproject.toml index 64ef7f2190f..dd4aed49e24 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,8 +22,6 @@ exclude = ''' | buck-out | build | dist - | model # temporary - to remove later - | montecarlo # temporary )/ | setup.py | docs/conf.py diff --git a/tardis/__init__.py b/tardis/__init__.py index bad783a0ccb..63597e684aa 100644 --- a/tardis/__init__.py +++ b/tardis/__init__.py @@ -17,8 +17,8 @@ from astropy import physical_constants, astronomical_constants -physical_constants.set('codata2014') -astronomical_constants.set('iau2012') +physical_constants.set("codata2014") +astronomical_constants.set("iau2012") # ---------------------------------------------------------------------------- diff --git a/tardis/analysis/opacities.py b/tardis/analysis/opacities.py index 56ab03cece5..a0ae8b0e196 100644 --- a/tardis/analysis/opacities.py +++ b/tardis/analysis/opacities.py @@ -342,7 +342,7 @@ def _calc_thomson_scattering_opacity(self): sigma_T = const.sigma_T except AttributeError: logger.warning("using astropy < 1.1.1: setting sigma_T manually") - sigma_T = 6.65245873e-29 * units.m ** 2 + sigma_T = 6.65245873e-29 * units.m**2 edens = self.mdl.plasma.electron_densities.values @@ -350,7 +350,7 @@ def _calc_thomson_scattering_opacity(self): edens.to("1/cm^3") except AttributeError: logger.info("setting electron density units by hand (cm^-3)") - edens = edens / units.cm ** 3 + edens = edens / units.cm**3 kappa_thom = edens * sigma_T @@ -376,9 +376,7 @@ def _calc_planck_mean_opacity(self): bb_nu = Blackbody(T) tmp = ( - bb_nu(self.nu_bins[:-1]) - * delta_nu - * self.kappa_tot[:, 0] + bb_nu(self.nu_bins[:-1]) * delta_nu * self.kappa_tot[:, 0] ).sum() tmp /= (bb_nu(self.nu_bins[:-1], T) * delta_nu).sum() diff --git a/tardis/conftest.py b/tardis/conftest.py index 52505bddd69..64232ff8783 100644 --- a/tardis/conftest.py +++ b/tardis/conftest.py @@ -171,6 +171,7 @@ def config_montecarlo_1e5_verysimple(): config = Configuration.from_yaml(path) return config + @pytest.fixture(scope="session") def simulation_verysimple(config_verysimple, atomic_dataset): atomic_data = deepcopy(atomic_dataset) diff --git a/tardis/gui/widgets.py b/tardis/gui/widgets.py index 5608a42065d..cb5b2dcc958 100644 --- a/tardis/gui/widgets.py +++ b/tardis/gui/widgets.py @@ -162,8 +162,8 @@ def shell_picker(self, shell, mouseevent): """Enable picking shells in the shell plot.""" if mouseevent.xdata is None: return False, dict() - mouse_r2 = mouseevent.xdata ** 2 + mouseevent.ydata ** 2 - if shell.r_inner ** 2 < mouse_r2 < shell.r_outer ** 2: + mouse_r2 = mouseevent.xdata**2 + mouseevent.ydata**2 + if shell.r_inner**2 < mouse_r2 < shell.r_outer**2: return True, dict() return False, dict() diff --git a/tardis/io/logger/tests/test_logging.py b/tardis/io/logger/tests/test_logging.py index 05291950ddb..1e512f277d6 100644 --- a/tardis/io/logger/tests/test_logging.py +++ b/tardis/io/logger/tests/test_logging.py @@ -7,7 +7,10 @@ from tardis import run_tardis import pytest -pytestmark = pytest.mark.skip(reason='logging testing slow and disabled for now') +pytestmark = pytest.mark.skip( + reason="logging testing slow and disabled for now" +) + def test_logging_simulation(atomic_data_fname, caplog): """ diff --git a/tardis/io/tests/test_model_reader.py b/tardis/io/tests/test_model_reader.py index 60a110b1a91..6cd9e76fe16 100644 --- a/tardis/io/tests/test_model_reader.py +++ b/tardis/io/tests/test_model_reader.py @@ -55,8 +55,8 @@ def test_simple_read_artis_density(artis_density_fname): assert np.isclose(0.00114661 * u.day, time_of_model, atol=1e-7 * u.day) assert np.isclose( mean_density[23], - 0.2250048 * u.g / u.cm ** 3, - atol=1.0e-6 * u.g / u.cm ** 3, + 0.2250048 * u.g / u.cm**3, + atol=1.0e-6 * u.g / u.cm**3, ) assert len(mean_density) == 69 assert len(velocity) == len(mean_density) + 1 @@ -115,11 +115,11 @@ def test_simple_read_cmfgen_density(cmfgen_fname): assert np.isclose(0.976 * u.day, time_of_model, atol=1e-7 * u.day) assert np.isclose( mean_density[4], - 4.2539537e-09 * u.g / u.cm ** 3, - atol=1.0e-6 * u.g / u.cm ** 3, + 4.2539537e-09 * u.g / u.cm**3, + atol=1.0e-6 * u.g / u.cm**3, ) assert np.isclose( - electron_densities[5], 2.6e14 * u.cm ** -3, atol=1.0e-6 * u.cm ** -3 + electron_densities[5], 2.6e14 * u.cm**-3, atol=1.0e-6 * u.cm**-3 ) assert len(mean_density) == 9 assert len(velocity) == len(mean_density) + 1 diff --git a/tardis/model/tests/test_base.py b/tardis/model/tests/test_base.py index 52cda7a02c8..307578b22c1 100644 --- a/tardis/model/tests/test_base.py +++ b/tardis/model/tests/test_base.py @@ -173,31 +173,31 @@ def test_abundances(self): def test_densities(self): assert_almost_equal( self.model.density[0].to(u.Unit("g/cm3")).value, - 9.7656229e-11 / 13.0 ** 3, + 9.7656229e-11 / 13.0**3, ) assert_almost_equal( self.model.density[1].to(u.Unit("g/cm3")).value, - 4.8170911e-11 / 13.0 ** 3, + 4.8170911e-11 / 13.0**3, ) assert_almost_equal( self.model.density[2].to(u.Unit("g/cm3")).value, - 2.5600000e-11 / 13.0 ** 3, + 2.5600000e-11 / 13.0**3, ) assert_almost_equal( self.model.density[3].to(u.Unit("g/cm3")).value, - 1.4450533e-11 / 13.0 ** 3, + 1.4450533e-11 / 13.0**3, ) assert_almost_equal( self.model.density[4].to(u.Unit("g/cm3")).value, - 8.5733893e-11 / 13.0 ** 3, + 8.5733893e-11 / 13.0**3, ) assert_almost_equal( self.model.density[5].to(u.Unit("g/cm3")).value, - 5.3037103e-11 / 13.0 ** 3, + 5.3037103e-11 / 13.0**3, ) assert_almost_equal( self.model.density[6].to(u.Unit("g/cm3")).value, - 3.3999447e-11 / 13.0 ** 3, + 3.3999447e-11 / 13.0**3, ) @@ -313,6 +313,7 @@ def test_model_decay(simple_isotope_abundance): class TestModelState: """Test the ModelState class.""" + def setup(self): """Initialize config and model.""" filename = "tardis_configv1_verysimple.yml" @@ -387,6 +388,7 @@ def test_density(self): def to_hdf_buffer(hdf_file_path, simulation_verysimple): simulation_verysimple.model.to_hdf(hdf_file_path, overwrite=True) + model_scalar_attrs = ["t_inner"] diff --git a/tardis/montecarlo/base.py b/tardis/montecarlo/base.py index eff2f8f55a6..ce3664ead89 100644 --- a/tardis/montecarlo/base.py +++ b/tardis/montecarlo/base.py @@ -29,7 +29,7 @@ logger = logging.getLogger(__name__) -MAX_SEED_VAL = 2 ** 32 - 1 +MAX_SEED_VAL = 2**32 - 1 # MAX_SEED_VAL must be multiple orders of magnitude larger than no_of_packets; # otherwise, each packet would not have its own seed. Here, we set the max @@ -73,14 +73,14 @@ class MontecarloRunner(HDFWriterMixin): hdf_name = "runner" w_estimator_constant = ( - (const.c ** 2 / (2 * const.h)) - * (15 / np.pi ** 4) + (const.c**2 / (2 * const.h)) + * (15 / np.pi**4) * (const.h / const.k_B) ** 4 / (4 * np.pi) ).cgs.value t_rad_estimator_constant = ( - (np.pi ** 4 / (15 * 24 * zeta(5, 1))) * (const.h / const.k_B) + (np.pi**4 / (15 * 24 * zeta(5, 1))) * (const.h / const.k_B) ).cgs.value def __init__( @@ -126,7 +126,7 @@ def __init__( self.seed = seed self._integrator = None self._spectrum_integrated = None - self.use_gpu=use_gpu + self.use_gpu = use_gpu self.virt_logging = virtual_packet_logging self.virt_packet_last_interaction_type = np.ones(2) * -1 @@ -269,10 +269,11 @@ def spectrum_virtual(self): @property def spectrum_integrated(self): if self._spectrum_integrated is None: - #This was changed from unpacking to specific attributes as compute - #is not used in calculate_spectrum + # This was changed from unpacking to specific attributes as compute + # is not used in calculate_spectrum self._spectrum_integrated = self.integrator.calculate_spectrum( - self.spectrum_frequency[:-1], points=self.integrator_settings.points, + self.spectrum_frequency[:-1], + points=self.integrator_settings.points, interpolate_shells=self.integrator_settings.interpolate_shells, ) return self._spectrum_integrated @@ -560,7 +561,7 @@ def calculate_radiationfield_properties(self): w = self.j_estimator / ( 4 * const.sigma_sb.cgs.value - * t_rad ** 4 + * t_rad**4 * self.time_of_simulation.value * self.volume.value ) @@ -584,7 +585,7 @@ def calculate_luminosity_inner(self, model): * np.pi * const.sigma_sb.cgs * model.r_inner[0] ** 2 - * model.t_inner ** 4 + * model.t_inner**4 ).to("erg/s") def calculate_time_of_simulation(self, model): @@ -641,8 +642,8 @@ def from_config( config.spectrum.start.to("Hz", u.spectral()), num=config.spectrum.num + 1, ) - running_mode=config.spectrum.integrated.compute.upper() - + running_mode = config.spectrum.integrated.compute.upper() + if running_mode == "GPU": if cuda.is_available(): use_gpu = True @@ -663,7 +664,7 @@ def from_config( """An invalid option for compute was passed. The three valid values are 'GPU', 'CPU', and 'Automatic'.""" ) - + mc_config_module.disable_line_scattering = ( config.plasma.disable_line_scattering ) diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py index ac8a34af192..be5b8f06e1e 100644 --- a/tardis/montecarlo/montecarlo_numba/base.py +++ b/tardis/montecarlo/montecarlo_numba/base.py @@ -26,6 +26,7 @@ from numba.typed import List from tardis.util.base import update_iterations_pbar, update_packet_pbar + def montecarlo_radial1d( model, plasma, @@ -100,22 +101,28 @@ def montecarlo_radial1d( runner.last_line_interaction_out_id = last_line_interaction_out_id if montecarlo_configuration.VPACKET_LOGGING and number_of_vpackets > 0: - runner.virt_packet_nus = np.concatenate( - virt_packet_nus).ravel() + runner.virt_packet_nus = np.concatenate(virt_packet_nus).ravel() runner.virt_packet_energies = np.concatenate( - virt_packet_energies).ravel() + virt_packet_energies + ).ravel() runner.virt_packet_initial_mus = np.concatenate( - virt_packet_initial_mus).ravel() + virt_packet_initial_mus + ).ravel() runner.virt_packet_initial_rs = np.concatenate( - virt_packet_initial_rs).ravel() + virt_packet_initial_rs + ).ravel() runner.virt_packet_last_interaction_in_nu = np.concatenate( - virt_packet_last_interaction_in_nu).ravel() + virt_packet_last_interaction_in_nu + ).ravel() runner.virt_packet_last_interaction_type = np.concatenate( - virt_packet_last_interaction_type).ravel() + virt_packet_last_interaction_type + ).ravel() runner.virt_packet_last_line_interaction_in_id = np.concatenate( - virt_packet_last_line_interaction_in_id).ravel() + virt_packet_last_line_interaction_in_id + ).ravel() runner.virt_packet_last_line_interaction_out_id = np.concatenate( - virt_packet_last_line_interaction_out_id).ravel() + virt_packet_last_line_interaction_out_id + ).ravel() update_iterations_pbar(1) # Condition for Checking if RPacket Tracking is enabled if montecarlo_configuration.RPACKET_TRACKING: @@ -145,7 +152,7 @@ def montecarlo_main_loop( ---------- packet_collection : PacketCollection numba_model : NumbaModel - numba_plasma : NumbaPlasma + numba_plasma : NumbaPlasma estimators : NumbaEstimators spectrum_frequency : astropy.units.Quantity frequency binspas @@ -188,7 +195,7 @@ def montecarlo_main_loop( ) ) rpacket_trackers.append(RPacketTracker()) - + # Arrays for vpacket logging virt_packet_nus = [] virt_packet_energies = [] @@ -233,7 +240,7 @@ def montecarlo_main_loop( numba_plasma, estimators, vpacket_collection, - rpacket_tracker + rpacket_tracker, ) output_nus[i] = r_packet.nu diff --git a/tardis/montecarlo/montecarlo_numba/calculate_distances.py b/tardis/montecarlo/montecarlo_numba/calculate_distances.py index 8f896b4b7a6..44f35a506e9 100644 --- a/tardis/montecarlo/montecarlo_numba/calculate_distances.py +++ b/tardis/montecarlo/montecarlo_numba/calculate_distances.py @@ -15,7 +15,10 @@ ) from tardis.montecarlo.montecarlo_numba.utils import MonteCarloException -from tardis.montecarlo.montecarlo_numba.r_packet import print_r_packet_properties +from tardis.montecarlo.montecarlo_numba.r_packet import ( + print_r_packet_properties, +) + @njit(**njit_dict_no_parallel) def calculate_distance_boundary(r, mu, r_inner, r_outer): @@ -96,9 +99,7 @@ def calculate_distance_line( if nu_diff >= 0: distance = (nu_diff / nu) * C_SPEED_OF_LIGHT * time_explosion else: - raise MonteCarloException( - "nu difference is less than 0.0" - ) + raise MonteCarloException("nu difference is less than 0.0") if nc.ENABLE_FULL_RELATIVITY: return calculate_distance_line_full_relativity( diff --git a/tardis/montecarlo/montecarlo_numba/continuum/r_packet_transport_continuum.py b/tardis/montecarlo/montecarlo_numba/continuum/r_packet_transport_continuum.py index dd275509944..6a30d68872f 100644 --- a/tardis/montecarlo/montecarlo_numba/continuum/r_packet_transport_continuum.py +++ b/tardis/montecarlo/montecarlo_numba/continuum/r_packet_transport_continuum.py @@ -3,24 +3,21 @@ from tardis.montecarlo import montecarlo_configuration from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel -from tardis.montecarlo.montecarlo_numba.calculate_distances import \ - calculate_distance_boundary, \ - calculate_distance_line -from tardis.montecarlo.montecarlo_numba.estimators import \ - update_line_estimators -from tardis.montecarlo.montecarlo_numba.frame_transformations import \ - get_doppler_factor +from tardis.montecarlo.montecarlo_numba.calculate_distances import ( + calculate_distance_boundary, + calculate_distance_line, +) +from tardis.montecarlo.montecarlo_numba.estimators import update_line_estimators +from tardis.montecarlo.montecarlo_numba.frame_transformations import ( + get_doppler_factor, +) from tardis.montecarlo.montecarlo_numba.r_packet import InteractionType + @njit(**njit_dict_no_parallel) def trace_packet_continuum( - r_packet, - numba_model, - numba_plasma, - estimators, - chi_continuum, - escat_prob + r_packet, numba_model, numba_plasma, estimators, chi_continuum, escat_prob ): """ Traces the RPacket through the ejecta and stops when an interaction happens (heart of the calculation) @@ -42,11 +39,7 @@ def trace_packet_continuum( ( distance_boundary, delta_shell, - ) = calculate_distance_boundary(r_packet.r, - r_packet.mu, - r_inner, - r_outer - ) + ) = calculate_distance_boundary(r_packet.r, r_packet.mu, r_inner, r_outer) # defining start for line interaction start_line_id = r_packet.next_line_id @@ -168,4 +161,4 @@ def trace_packet_continuum( distance = distance_boundary interaction_type = InteractionType.BOUNDARY - return distance, interaction_type, delta_shell \ No newline at end of file + return distance, interaction_type, delta_shell diff --git a/tardis/montecarlo/montecarlo_numba/estimators.py b/tardis/montecarlo/montecarlo_numba/estimators.py index 73d819fbbbc..630853f839f 100644 --- a/tardis/montecarlo/montecarlo_numba/estimators.py +++ b/tardis/montecarlo/montecarlo_numba/estimators.py @@ -13,6 +13,7 @@ calc_packet_energy_full_relativity, ) + @njit(**njit_dict_no_parallel) def set_estimators(r_packet, distance, numba_estimator, comov_nu, comov_energy): """ diff --git a/tardis/montecarlo/montecarlo_numba/frame_transformations.py b/tardis/montecarlo/montecarlo_numba/frame_transformations.py index 122a7402d95..e6245b278a2 100644 --- a/tardis/montecarlo/montecarlo_numba/frame_transformations.py +++ b/tardis/montecarlo/montecarlo_numba/frame_transformations.py @@ -7,9 +7,8 @@ ) from tardis.montecarlo.montecarlo_numba import numba_config as nc -from tardis.montecarlo.montecarlo_numba.numba_config import ( - C_SPEED_OF_LIGHT -) +from tardis.montecarlo.montecarlo_numba.numba_config import C_SPEED_OF_LIGHT + @njit(**njit_dict_no_parallel) def get_doppler_factor(r, mu, time_explosion): diff --git a/tardis/montecarlo/montecarlo_numba/interaction.py b/tardis/montecarlo/montecarlo_numba/interaction.py index 0b96f7ab508..fa958a35e32 100644 --- a/tardis/montecarlo/montecarlo_numba/interaction.py +++ b/tardis/montecarlo/montecarlo_numba/interaction.py @@ -14,11 +14,13 @@ angle_aberration_CMF_to_LF, ) from tardis.montecarlo.montecarlo_numba.r_packet import ( - InteractionType, PacketStatus + InteractionType, + PacketStatus, ) from tardis.montecarlo.montecarlo_numba.utils import get_random_mu from tardis.montecarlo.montecarlo_numba.macro_atom import ( - macro_atom, MacroAtomTransitionType + macro_atom, + MacroAtomTransitionType, ) from tardis import constants as const @@ -26,11 +28,10 @@ H = const.h.cgs.value - @njit(**njit_dict_no_parallel) -def determine_bf_macro_activation_idx(numba_plasma, - nu, chi_bf_contributions, active_continua - ): +def determine_bf_macro_activation_idx( + numba_plasma, nu, chi_bf_contributions, active_continua +): """ Determine the macro atom activation level after bound-free absorption. @@ -52,7 +53,6 @@ def determine_bf_macro_activation_idx(numba_plasma, # Perform a MC experiment to determine the continuum for absorption index = np.searchsorted(chi_bf_contributions, np.random.random()) continuum_id = active_continua[index] - # Perform a MC experiment to determine whether thermal or # ionization energy is created @@ -61,15 +61,18 @@ def determine_bf_macro_activation_idx(numba_plasma, if ( np.random.random() < fraction_ionization ): # Create ionization energy (i-packet) - destination_level_idx = numba_plasma.photo_ion_activation_idx[continuum_id] + destination_level_idx = numba_plasma.photo_ion_activation_idx[ + continuum_id + ] else: # Create thermal energy (k-packet) destination_level_idx = numba_plasma.k_packet_idx return destination_level_idx + @njit(**njit_dict_no_parallel) -def determine_continuum_macro_activation_idx(numba_plasma, - nu, chi_bf, chi_ff, chi_bf_contributions, active_continua - ): +def determine_continuum_macro_activation_idx( + numba_plasma, nu, chi_bf, chi_ff, chi_bf_contributions, active_continua +): """ Determine the macro atom activation level after a continuum absorption. @@ -117,6 +120,7 @@ def sample_nu_free_free(numba_plasma, shell): zrand = np.random.random() return -K_B * T / H * np.log(zrand) + @njit(**njit_dict_no_parallel) def sample_nu_free_bound(numba_plasma, shell, continuum_id): """ @@ -134,9 +138,10 @@ def sample_nu_free_bound(numba_plasma, shell, continuum_id): zrand = np.random.random() idx = np.searchsorted(em, zrand, side="right") - return phot_nus_block[idx] - (em[idx] - zrand) / ( - em[idx] - em[idx - 1] - ) * (phot_nus_block[idx] - phot_nus_block[idx - 1]) + return phot_nus_block[idx] - (em[idx] - zrand) / (em[idx] - em[idx - 1]) * ( + phot_nus_block[idx] - phot_nus_block[idx - 1] + ) + @njit(**njit_dict_no_parallel) def scatter(r_packet, time_explosion): @@ -152,12 +157,20 @@ def scatter(r_packet, time_explosion): ) r_packet.energy = comov_energy * inverse_new_doppler_factor - + return comov_nu, inverse_new_doppler_factor + @njit(**njit_dict_no_parallel) -def continuum_event(r_packet, time_explosion, numba_plasma, - chi_bf_tot, chi_ff, chi_bf_contributions, current_continua): +def continuum_event( + r_packet, + time_explosion, + numba_plasma, + chi_bf_tot, + chi_ff, + chi_bf_contributions, + current_continua, +): """ continuum event handler - activate the macroatom and run the handler @@ -169,33 +182,38 @@ def continuum_event(r_packet, time_explosion, numba_plasma, continuum : tardis.montecarlo.montecarlo_numba.numba_interface.Continuum """ old_doppler_factor = get_doppler_factor( - r_packet.r, - r_packet.mu, - time_explosion - ) + r_packet.r, r_packet.mu, time_explosion + ) r_packet.mu = get_random_mu() inverse_doppler_factor = get_inverse_doppler_factor( - r_packet.r, - r_packet.mu, - time_explosion - ) + r_packet.r, r_packet.mu, time_explosion + ) comov_energy = r_packet.energy * old_doppler_factor - comov_nu = r_packet.nu * old_doppler_factor # make sure frequency should be updated + comov_nu = ( + r_packet.nu * old_doppler_factor + ) # make sure frequency should be updated r_packet.energy = comov_energy * inverse_doppler_factor r_packet.nu = comov_nu * inverse_doppler_factor destination_level_idx = determine_continuum_macro_activation_idx( - numba_plasma, comov_nu, chi_bf_tot, chi_ff, - chi_bf_contributions, current_continua) + numba_plasma, + comov_nu, + chi_bf_tot, + chi_ff, + chi_bf_contributions, + current_continua, + ) + + macro_atom_event( + destination_level_idx, r_packet, time_explosion, numba_plasma + ) - macro_atom_event(destination_level_idx, - r_packet, time_explosion, - numba_plasma) @njit(**njit_dict_no_parallel) -def macro_atom_event(destination_level_idx, - r_packet, time_explosion, numba_plasma): +def macro_atom_event( + destination_level_idx, r_packet, time_explosion, numba_plasma +): """ Macroatom event handler - run the macroatom and handle the result @@ -206,44 +224,38 @@ def macro_atom_event(destination_level_idx, time_explosion : float numba_plasma : tardis.montecarlo.montecarlo_numba.numba_interface.NumbaPlasma """ - + transition_id, transition_type = macro_atom( - destination_level_idx, - r_packet.current_shell_id, - numba_plasma - ) - - if (montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED and - transition_type == MacroAtomTransitionType.FF_EMISSION): - free_free_emission( - r_packet, - time_explosion, - numba_plasma - ) - - elif (montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED and - transition_type == MacroAtomTransitionType.BF_EMISSION): + destination_level_idx, r_packet.current_shell_id, numba_plasma + ) + + if ( + montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED + and transition_type == MacroAtomTransitionType.FF_EMISSION + ): + free_free_emission(r_packet, time_explosion, numba_plasma) + + elif ( + montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED + and transition_type == MacroAtomTransitionType.BF_EMISSION + ): bound_free_emission( - r_packet, - time_explosion, - numba_plasma, - transition_id - ) - elif (montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED and - transition_type == MacroAtomTransitionType.BF_COOLING): + r_packet, time_explosion, numba_plasma, transition_id + ) + elif ( + montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED + and transition_type == MacroAtomTransitionType.BF_COOLING + ): bf_cooling(r_packet, time_explosion, numba_plasma) - - elif (montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED and - transition_type == MacroAtomTransitionType.ADIABATIC_COOLING): + + elif ( + montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED + and transition_type == MacroAtomTransitionType.ADIABATIC_COOLING + ): adiabatic_cooling(r_packet) elif transition_type == MacroAtomTransitionType.BB_EMISSION: - line_emission( - r_packet, - transition_id, - time_explosion, - numba_plasma - ) + line_emission(r_packet, transition_id, time_explosion, numba_plasma) else: raise Exception("No Interaction Found!") @@ -259,22 +271,19 @@ def bf_cooling(r_packet, time_explosion, numba_plasma): time_explosion : float numba_plasma : tardis.montecarlo.montecarlo_numba.numba_interface.NumbaPlasma """ - - fb_cooling_prob = numba_plasma.p_fb_deactivation[:, - r_packet.current_shell_id] + + fb_cooling_prob = numba_plasma.p_fb_deactivation[ + :, r_packet.current_shell_id + ] p = fb_cooling_prob[0] i = 0 zrand = np.random.random() - while p <= zrand: # Can't search-sorted this because it's not cumulative + while p <= zrand: # Can't search-sorted this because it's not cumulative i += 1 p += fb_cooling_prob[i] continuum_idx = i - bound_free_emission( - r_packet, - time_explosion, - numba_plasma, - continuum_idx - ) + bound_free_emission(r_packet, time_explosion, numba_plasma, continuum_idx) + @njit(**njit_dict_no_parallel) def adiabatic_cooling(r_packet): @@ -288,6 +297,7 @@ def adiabatic_cooling(r_packet): r_packet.status = PacketStatus.ADIABATIC_COOLING + @njit(**njit_dict_no_parallel) def get_current_line_id(nu, line_list): """ @@ -319,23 +329,21 @@ def free_free_emission(r_packet, time_explosion, numba_plasma): time_explosion : float numba_plasma : tardis.montecarlo.montecarlo_numba.numba_interface.NumbaPlasma """ - + inverse_doppler_factor = get_inverse_doppler_factor( r_packet.r, r_packet.mu, time_explosion ) comov_nu = sample_nu_free_free(numba_plasma, r_packet.current_shell_id) r_packet.nu = comov_nu * inverse_doppler_factor - current_line_id = get_current_line_id( - comov_nu, - numba_plasma.line_list_nu - ) - r_packet.next_line_id = current_line_id - + current_line_id = get_current_line_id(comov_nu, numba_plasma.line_list_nu) + r_packet.next_line_id = current_line_id + if montecarlo_configuration.full_relativity: r_packet.mu = angle_aberration_CMF_to_LF( r_packet, time_explosion, r_packet.mu ) + @njit(**njit_dict_no_parallel) def bound_free_emission(r_packet, time_explosion, numba_plasma, continuum_id): """ @@ -348,19 +356,16 @@ def bound_free_emission(r_packet, time_explosion, numba_plasma, continuum_id): numba_plasma : tardis.montecarlo.montecarlo_numba.numba_interface.NumbaPlasma continuum_id : int """ - + inverse_doppler_factor = get_inverse_doppler_factor( r_packet.r, r_packet.mu, time_explosion ) - comov_nu = sample_nu_free_bound(numba_plasma, - r_packet.current_shell_id, - continuum_id) + comov_nu = sample_nu_free_bound( + numba_plasma, r_packet.current_shell_id, continuum_id + ) r_packet.nu = comov_nu * inverse_doppler_factor - current_line_id = get_current_line_id( - comov_nu, - numba_plasma.line_list_nu - ) + current_line_id = get_current_line_id(comov_nu, numba_plasma.line_list_nu) r_packet.next_line_id = current_line_id if montecarlo_configuration.full_relativity: @@ -384,7 +389,7 @@ def thomson_scatter(r_packet, time_explosion): time_explosion : float time since explosion in seconds """ - + old_doppler_factor = get_doppler_factor( r_packet.r, r_packet.mu, time_explosion ) @@ -407,8 +412,7 @@ def thomson_scatter(r_packet, time_explosion): @njit(**njit_dict_no_parallel) -def line_scatter(r_packet, time_explosion, - line_interaction_type, numba_plasma): +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 @@ -437,17 +441,18 @@ def line_scatter(r_packet, time_explosion, r_packet, r_packet.next_line_id, time_explosion, numba_plasma ) else: # includes both macro atom and downbranch - encoded in the transition probabilities - comov_nu = r_packet.nu * old_doppler_factor # Is this necessary? - r_packet.nu = comov_nu * inverse_new_doppler_factor + comov_nu = r_packet.nu * old_doppler_factor # Is this necessary? + r_packet.nu = comov_nu * inverse_new_doppler_factor activation_level_id = numba_plasma.line2macro_level_upper[ r_packet.next_line_id ] - macro_atom_event(activation_level_id, r_packet, - time_explosion, numba_plasma) + macro_atom_event( + activation_level_id, r_packet, time_explosion, numba_plasma + ) + @njit(**njit_dict_no_parallel) -def line_emission(r_packet, emission_line_id, time_explosion, - numba_plasma): +def line_emission(r_packet, emission_line_id, time_explosion, numba_plasma): """ Sets the frequency of the RPacket properly given the emission channel @@ -458,7 +463,7 @@ def line_emission(r_packet, emission_line_id, time_explosion, time_explosion : float numba_plasma : tardis.montecarlo.montecarlo_numba.numba_interface.NumbaPlasma """ - + r_packet.last_line_interaction_out_id = emission_line_id if emission_line_id != r_packet.next_line_id: diff --git a/tardis/montecarlo/montecarlo_numba/macro_atom.py b/tardis/montecarlo/montecarlo_numba/macro_atom.py index b04a41ef4a4..19916c60b1f 100644 --- a/tardis/montecarlo/montecarlo_numba/macro_atom.py +++ b/tardis/montecarlo/montecarlo_numba/macro_atom.py @@ -67,4 +67,7 @@ def macro_atom(activation_level_id, current_shell_id, numba_plasma): ) # current_transition_type = MacroAtomTransitionType(current_transition_type) - return numba_plasma.transition_line_id[transition_id], current_transition_type + return ( + numba_plasma.transition_line_id[transition_id], + current_transition_type, + ) diff --git a/tardis/montecarlo/montecarlo_numba/numba_interface.py b/tardis/montecarlo/montecarlo_numba/numba_interface.py index 8d8a4df109e..4971db407a6 100644 --- a/tardis/montecarlo/montecarlo_numba/numba_interface.py +++ b/tardis/montecarlo/montecarlo_numba/numba_interface.py @@ -57,9 +57,9 @@ def __init__(self, r_inner, r_outer, time_explosion): ("x_sect", float64[:]), ("phot_nus", float64[:]), ("ff_opacity_factor", float64[:]), - ("emissivities", float64[:,:]), + ("emissivities", float64[:, :]), ("photo_ion_activation_idx", int64[:]), - ("k_packet_idx", int64) + ("k_packet_idx", int64), ] @@ -88,7 +88,7 @@ def __init__( ff_opacity_factor, emissivities, photo_ion_activation_idx, - k_packet_idx + k_packet_idx, ): """ Plasma for the Numba code @@ -139,6 +139,7 @@ def __init__( self.photo_ion_activation_idx = photo_ion_activation_idx self.k_packet_idx = k_packet_idx + def numba_plasma_initialize(plasma, line_interaction_type): """ Initialize the NumbaPlasma object and copy over the data over from TARDIS Plasma @@ -182,8 +183,8 @@ def numba_plasma_initialize(plasma, line_interaction_type): macro_block_references = plasma.macro_block_references else: macro_block_references = plasma.atomic_data.macro_atom_references[ - "block_references" - ].values + "block_references" + ].values transition_type = plasma.macro_atom_data["transition_type"].values # Destination level is not needed and/or generated for downbranch @@ -196,7 +197,8 @@ def numba_plasma_initialize(plasma, line_interaction_type): plasma.level2continuum_idx.index ].values p_fb_deactivation = np.ascontiguousarray( - plasma.p_fb_deactivation.values.copy(), dtype=np.float64) + plasma.p_fb_deactivation.values.copy(), dtype=np.float64 + ) phot_nus = plasma.photo_ion_cross_sections.nu.loc[ plasma.level2continuum_idx.index @@ -207,8 +209,12 @@ def numba_plasma_initialize(plasma, line_interaction_type): .values.cumsum(), [1, 0], ) - photo_ion_nu_threshold_mins = phot_nus.groupby(level=[0, 1, 2], sort=False).first().values - photo_ion_nu_threshold_maxs = phot_nus.groupby(level=[0, 1, 2], sort=False).last().values + photo_ion_nu_threshold_mins = ( + phot_nus.groupby(level=[0, 1, 2], sort=False).first().values + ) + photo_ion_nu_threshold_maxs = ( + phot_nus.groupby(level=[0, 1, 2], sort=False).last().values + ) chi_bf = plasma.chi_bf.loc[plasma.level2continuum_idx.index].values x_sect = plasma.photo_ion_cross_sections.x_sect.loc[ @@ -216,8 +222,10 @@ def numba_plasma_initialize(plasma, line_interaction_type): ].values phot_nus = phot_nus.values - ff_opacity_factor = plasma.ff_cooling_factor/np.sqrt(t_electrons) - emissivities = plasma.fb_emission_cdf.loc[plasma.level2continuum_idx.index].values + ff_opacity_factor = plasma.ff_cooling_factor / np.sqrt(t_electrons) + emissivities = plasma.fb_emission_cdf.loc[ + plasma.level2continuum_idx.index + ].values photo_ion_activation_idx = plasma.photo_ion_idx.loc[ plasma.level2continuum_idx.index, "destination_level_idx" ].values @@ -228,11 +236,11 @@ def numba_plasma_initialize(plasma, line_interaction_type): photo_ion_nu_threshold_mins = np.zeros(0, dtype=np.float64) photo_ion_nu_threshold_maxs = np.zeros(0, dtype=np.float64) photo_ion_block_references = np.zeros(0, dtype=np.int64) - chi_bf = np.zeros((0,0), dtype=np.float64) + chi_bf = np.zeros((0, 0), dtype=np.float64) x_sect = np.zeros(0, dtype=np.float64) phot_nus = np.zeros(0, dtype=np.float64) ff_opacity_factor = np.zeros(0, dtype=np.float64) - emissivities = np.zeros((0,0), dtype=np.float64) + emissivities = np.zeros((0, 0), dtype=np.float64) photo_ion_activation_idx = np.zeros(0, dtype=np.int64) k_packet_idx = np.int64(-1) @@ -258,7 +266,7 @@ def numba_plasma_initialize(plasma, line_interaction_type): ff_opacity_factor, emissivities, photo_ion_activation_idx, - k_packet_idx + k_packet_idx, ) @@ -406,6 +414,7 @@ def set_properties( self.last_interaction_out_id[self.idx] = last_interaction_out_id self.idx += 1 + rpacket_tracker_spec = [ ("length", int64), ("seed", int64), @@ -419,6 +428,7 @@ def set_properties( ("interact_id", int64), ] + @jitclass(rpacket_tracker_spec) class RPacketTracker(object): """ @@ -520,7 +530,6 @@ def finalize_array(self): ] - @jitclass(base_estimators_spec + continuum_estimators_spec) class Estimators(object): def __init__( diff --git a/tardis/montecarlo/montecarlo_numba/opacities.py b/tardis/montecarlo/montecarlo_numba/opacities.py index 72af275f4a5..66c4bc5f119 100644 --- a/tardis/montecarlo/montecarlo_numba/opacities.py +++ b/tardis/montecarlo/montecarlo_numba/opacities.py @@ -18,7 +18,7 @@ C = const.c.cgs.value FF_OPAC_CONST = ( - (2 * np.pi / (3 * M_E * K_B)) ** 0.5 * 4 * E ** 6 / (3 * M_E * H * C) + (2 * np.pi / (3 * M_E * K_B)) ** 0.5 * 4 * E**6 / (3 * M_E * H * C) ) # See Eq. 6.1.8 in http://personal.psu.edu/rbc3/A534/lec6.pdf @@ -42,6 +42,7 @@ def chi_electron_calculator(numba_plasma, nu, shell): """ return numba_plasma.electron_density[shell] * SIGMA_THOMSON + @njit(**njit_dict_no_parallel) def calculate_tau_electron(electron_density, distance): """ @@ -59,6 +60,7 @@ def calculate_tau_electron(electron_density, distance): """ return electron_density * SIGMA_THOMSON * distance + @njit(**njit_dict_no_parallel) def get_current_bound_free_continua(numba_plasma, nu): """ @@ -77,11 +79,10 @@ def get_current_bound_free_continua(numba_plasma, nu): """ nu_mins = numba_plasma.photo_ion_nu_threshold_mins nu_maxs = numba_plasma.photo_ion_nu_threshold_maxs - current_continua = np.where( - np.logical_and(nu >= nu_mins, nu <= nu_maxs) - )[0] + current_continua = np.where(np.logical_and(nu >= nu_mins, nu <= nu_maxs))[0] return current_continua + @njit(**njit_dict_no_parallel) def chi_bf_interpolator(numba_plasma, nu, shell): """ @@ -111,7 +112,7 @@ def chi_bf_interpolator(numba_plasma, nu, shell): Photoionization cross-sections of all bound-free continua for which absorption is possible for frequency `nu`. """ - + current_continua = get_current_bound_free_continua(numba_plasma, nu) chi_bfs = np.zeros(len(current_continua)) x_sect_bfs = np.zeros(len(current_continua)) @@ -120,16 +121,22 @@ def chi_bf_interpolator(numba_plasma, nu, shell): end = numba_plasma.photo_ion_block_references[continuum_id + 1] phot_nus_continuum = numba_plasma.phot_nus[start:end] nu_idx = np.searchsorted(phot_nus_continuum, nu) - interval = phot_nus_continuum[nu_idx] - phot_nus_continuum[nu_idx-1] - high_weight = (nu - phot_nus_continuum[nu_idx-1]) - low_weight = (phot_nus_continuum[nu_idx] - nu) + interval = phot_nus_continuum[nu_idx] - phot_nus_continuum[nu_idx - 1] + high_weight = nu - phot_nus_continuum[nu_idx - 1] + low_weight = phot_nus_continuum[nu_idx] - nu chi_bfs_continuum = numba_plasma.chi_bf[start:end, shell] - chi_bfs[i] = (chi_bfs_continuum[nu_idx]*high_weight + chi_bfs_continuum[nu_idx-1]*low_weight)/interval + chi_bfs[i] = ( + chi_bfs_continuum[nu_idx] * high_weight + + chi_bfs_continuum[nu_idx - 1] * low_weight + ) / interval x_sect_bfs_continuum = numba_plasma.x_sect[start:end] - x_sect_bfs[i] = (x_sect_bfs_continuum[nu_idx]*high_weight + x_sect_bfs_continuum[nu_idx-1]*low_weight)/interval + x_sect_bfs[i] = ( + x_sect_bfs_continuum[nu_idx] * high_weight + + x_sect_bfs_continuum[nu_idx - 1] * low_weight + ) / interval chi_bf_contributions = chi_bfs.cumsum() - + # If we are outside the range of frequencies # for which we have photo-ionization cross sections # we will have no local continuua and therefore @@ -148,6 +155,7 @@ def chi_bf_interpolator(numba_plasma, nu, shell): x_sect_bfs, ) + @njit(**njit_dict_no_parallel) def chi_ff_calculator(numba_plasma, nu, shell): """ @@ -167,11 +175,12 @@ def chi_ff_calculator(numba_plasma, nu, shell): chi_ff = ( FF_OPAC_CONST * numba_plasma.ff_opacity_factor[shell] - / nu ** 3 + / nu**3 * (1 - np.exp(-H * nu / (K_B * numba_plasma.t_electrons[shell]))) ) return chi_ff + @njit(**njit_dict_no_parallel) def chi_continuum_calculator(numba_plasma, nu, shell): """ @@ -185,7 +194,7 @@ def chi_continuum_calculator(numba_plasma, nu, shell): Returns ------- - + chi_bf_tot : float Total bound-free opacity at frequency `nu`. chi_bf_contributions : numpy.ndarray, dtype float @@ -212,4 +221,4 @@ def chi_continuum_calculator(numba_plasma, nu, shell): current_continua, x_sect_bfs, chi_ff, - ) \ No newline at end of file + ) diff --git a/tardis/montecarlo/montecarlo_numba/r_packet.py b/tardis/montecarlo/montecarlo_numba/r_packet.py index 0a05b1cfe23..f151a140298 100644 --- a/tardis/montecarlo/montecarlo_numba/r_packet.py +++ b/tardis/montecarlo/montecarlo_numba/r_packet.py @@ -13,6 +13,7 @@ from tardis.montecarlo.montecarlo_numba import numba_config as nc from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel + class InteractionType(IntEnum): BOUNDARY = 1 LINE = 2 @@ -26,6 +27,7 @@ class PacketStatus(IntEnum): REABSORBED = 2 ADIABATIC_COOLING = 4 + rpacket_spec = [ ("r", float64), ("mu", float64), @@ -42,6 +44,7 @@ class PacketStatus(IntEnum): ("last_line_interaction_out_id", int64), ] + @jitclass(rpacket_spec) class RPacket(object): def __init__(self, r, mu, nu, energy, seed, index=0): @@ -72,8 +75,6 @@ def initialize_line_id(self, numba_plasma, numba_model): self.next_line_id = next_line_id - - @njit(**njit_dict_no_parallel) def print_r_packet_properties(r_packet): """ @@ -88,5 +89,9 @@ def print_r_packet_properties(r_packet): print("R-Packet information:") with objmode: for r_packet_attribute_name, _ in rpacket_spec: - print(r_packet_attribute_name, "=", str(getattr(r_packet, r_packet_attribute_name))) - print("-" * 80) \ No newline at end of file + print( + r_packet_attribute_name, + "=", + str(getattr(r_packet, r_packet_attribute_name)), + ) + print("-" * 80) diff --git a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py index 6f750775a59..f9302a0fcf3 100644 --- a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py +++ b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py @@ -4,13 +4,17 @@ PacketStatus, ) from tardis.montecarlo.montecarlo_numba.r_packet_transport import ( - move_r_packet, move_packet_across_shell_boundary) -from tardis.montecarlo.montecarlo_numba.continuum.r_packet_transport_continuum import trace_packet_continuum + move_r_packet, + move_packet_across_shell_boundary, +) +from tardis.montecarlo.montecarlo_numba.continuum.r_packet_transport_continuum import ( + trace_packet_continuum, +) from tardis.montecarlo.montecarlo_numba.frame_transformations import ( get_inverse_doppler_factor, - get_doppler_factor + get_doppler_factor, ) from tardis.montecarlo.montecarlo_numba.interaction import ( InteractionType, @@ -26,22 +30,25 @@ from tardis import constants as const from tardis.montecarlo.montecarlo_numba.opacities import ( - chi_continuum_calculator, chi_electron_calculator + chi_continuum_calculator, + chi_electron_calculator, ) from tardis.montecarlo.montecarlo_numba.estimators import ( - update_bound_free_estimators + update_bound_free_estimators, ) C_SPEED_OF_LIGHT = const.c.to("cm/s").value @njit -def single_packet_loop(r_packet, +def single_packet_loop( + r_packet, numba_model, numba_plasma, estimators, vpacket_collection, - rpacket_tracker): + rpacket_tracker, +): """ Parameters ---------- @@ -92,15 +99,17 @@ def single_packet_loop(r_packet, x_sect_bfs, chi_ff, ) = chi_continuum_calculator( - numba_plasma, - comov_nu, - r_packet.current_shell_id + numba_plasma, comov_nu, r_packet.current_shell_id ) chi_continuum = chi_e + chi_bf_tot + chi_ff - escat_prob = chi_e / chi_continuum # probability of e-scatter + escat_prob = chi_e / chi_continuum # probability of e-scatter distance, interaction_type, delta_shell = trace_packet_continuum( - r_packet, numba_model, numba_plasma, - estimators, chi_continuum, escat_prob + r_packet, + numba_model, + numba_plasma, + estimators, + chi_continuum, + escat_prob, ) update_bound_free_estimators( comov_nu, @@ -117,10 +126,14 @@ def single_packet_loop(r_packet, escat_prob = 1.0 chi_continuum = chi_e distance, interaction_type, delta_shell = trace_packet_continuum( - r_packet, numba_model, numba_plasma, - estimators, chi_continuum, escat_prob + r_packet, + numba_model, + numba_plasma, + estimators, + chi_continuum, + escat_prob, ) - + # If continuum processes: update continuum estimators if interaction_type == InteractionType.BOUNDARY: @@ -143,8 +156,7 @@ def single_packet_loop(r_packet, numba_plasma, ) trace_vpacket_volley( - r_packet, vpacket_collection, - numba_model, numba_plasma + r_packet, vpacket_collection, numba_model, numba_plasma ) elif interaction_type == InteractionType.ESCATTERING: @@ -158,15 +170,23 @@ def single_packet_loop(r_packet, trace_vpacket_volley( r_packet, vpacket_collection, numba_model, numba_plasma ) - elif (montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED and - interaction_type == InteractionType.CONTINUUM_PROCESS): + elif ( + montecarlo_configuration.CONTINUUM_PROCESSES_ENABLED + and interaction_type == InteractionType.CONTINUUM_PROCESS + ): r_packet.last_interaction_type = InteractionType.CONTINUUM_PROCESS move_r_packet( r_packet, distance, numba_model.time_explosion, estimators ) - continuum_event(r_packet, numba_model.time_explosion, - numba_plasma, chi_bf_tot, chi_ff, chi_bf_contributions, - current_continua) + continuum_event( + r_packet, + numba_model.time_explosion, + numba_plasma, + chi_bf_tot, + chi_ff, + chi_bf_contributions, + current_continua, + ) trace_vpacket_volley( r_packet, vpacket_collection, numba_model, numba_plasma diff --git a/tardis/montecarlo/montecarlo_numba/tests/conftest.py b/tardis/montecarlo/montecarlo_numba/tests/conftest.py index e769aec34ed..6b9be84d2ba 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/conftest.py +++ b/tardis/montecarlo/montecarlo_numba/tests/conftest.py @@ -59,7 +59,7 @@ def verysimple_numba_model(nb_simulation_verysimple): @pytest.fixture(scope="package") def verysimple_estimators(nb_simulation_verysimple): runner = nb_simulation_verysimple.runner - + return Estimators( runner.j_estimator, runner.nu_bar_estimator, @@ -69,7 +69,7 @@ def verysimple_estimators(nb_simulation_verysimple): runner.stim_recomb_estimator, runner.bf_heating_estimator, runner.stim_recomb_cooling_estimator, - runner.photo_ion_estimator_statistics + runner.photo_ion_estimator_statistics, ) @@ -127,6 +127,7 @@ def packet(verysimple_packet_collection): index=0, ) + @pytest.fixture(scope="function") def static_packet(): return RPacket( diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_base.py b/tardis/montecarlo/montecarlo_numba/tests/test_base.py index fde614fc50d..a8b261cbdf2 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_base.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_base.py @@ -14,6 +14,7 @@ def test_montecarlo_radial1d(): assert False + def test_montecarlo_main_loop( config_montecarlo_1e5_verysimple, atomic_dataset, @@ -21,7 +22,7 @@ def test_montecarlo_main_loop( tmpdir, set_seed_fixture, random_call_fixture, - request + request, ): montecarlo_configuration.LEGACY_MODE_ENABLED = True @@ -30,16 +31,20 @@ def test_montecarlo_main_loop( config_montecarlo_1e5_verysimple.montecarlo.last_no_of_packets = 1e5 config_montecarlo_1e5_verysimple.montecarlo.no_of_virtual_packets = 0 config_montecarlo_1e5_verysimple.montecarlo.iterations = 1 - config_montecarlo_1e5_verysimple.plasma.line_interaction_type = 'macroatom' + config_montecarlo_1e5_verysimple.plasma.line_interaction_type = "macroatom" del config_montecarlo_1e5_verysimple["config_dirname"] - sim = Simulation.from_config(config_montecarlo_1e5_verysimple, atom_data=atomic_data) + sim = Simulation.from_config( + config_montecarlo_1e5_verysimple, atom_data=atomic_data + ) sim.run() - compare_fname = os.path.join(tardis_ref_path, "montecarlo_1e5_compare_data.h5") + compare_fname = os.path.join( + tardis_ref_path, "montecarlo_1e5_compare_data.h5" + ) if request.config.getoption("--generate-reference"): sim.to_hdf(compare_fname, overwrite=True) - + # Load compare data from refdata expected_nu = pd.read_hdf( compare_fname, key="/simulation/runner/output_nu" @@ -54,7 +59,6 @@ def test_montecarlo_main_loop( compare_fname, key="/simulation/runner/j_estimator" ).values - actual_energy = sim.runner.output_energy actual_nu = sim.runner.output_nu actual_nu_bar_estimator = sim.runner.nu_bar_estimator diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py index 9ecc4d51272..bcdaa8f1736 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py @@ -32,7 +32,7 @@ def test_line_scatter( line_interaction_type, packet, verysimple_numba_model, - verysimple_numba_plasma + verysimple_numba_plasma, ): init_mu = packet.mu init_nu = packet.nu diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py b/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py index bda1b8f0657..752b2dba81d 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py @@ -25,7 +25,7 @@ def test_macro_atom( result, transition_type = macro_atom.macro_atom( activation_level_id, static_packet.current_shell_id, - verysimple_numba_plasma + verysimple_numba_plasma, ) assert result == expected assert transition_type == -1 # line transition diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_numba_formal_integral.py b/tardis/montecarlo/montecarlo_numba/tests/test_numba_formal_integral.py index a20657a3d7c..b019c787620 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_numba_formal_integral.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_numba_formal_integral.py @@ -9,6 +9,7 @@ import tardis.montecarlo.montecarlo_numba.formal_integral as formal_integral from tardis.montecarlo.montecarlo_numba.numba_interface import NumbaModel + @pytest.mark.parametrize( ["nu", "T"], [ @@ -41,6 +42,7 @@ def test_trapezoid_integration(N): def calculate_z(r, p): return np.sqrt(r * r - p * p) + TESTDATA = [ { "r": np.linspace(1, 2, 3, dtype=np.float64), @@ -48,28 +50,24 @@ def calculate_z(r, p): { "r": np.linspace(0, 1, 3), }, - #{"r": np.linspace(1, 2, 10, dtype=np.float64)}, + # {"r": np.linspace(1, 2, 10, dtype=np.float64)}, ] @pytest.fixture(scope="function", params=TESTDATA) def formal_integral_model(request): r = request.param["r"] - model = NumbaModel( - r[:-1], - r[1:], - 1/c.c.cgs.value) + model = NumbaModel(r[:-1], r[1:], 1 / c.c.cgs.value) return model -@pytest.mark.parametrize( - 'p', [0.0, 0.5, 1.0] -) + +@pytest.mark.parametrize("p", [0.0, 0.5, 1.0]) def test_calculate_z(formal_integral_model, p): - + func = formal_integral.calculate_z inv_t = 1.0 / formal_integral_model.time_explosion size = len(formal_integral_model.r_outer) - r_outer = formal_integral_model.r_outer + r_outer = formal_integral_model.r_outer for r in r_outer: actual = func(r, p, inv_t) diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py b/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py index 75b287e0faa..207aa1a60a8 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py @@ -65,7 +65,7 @@ def test_VPacketCollection_set_properties(verysimple_3vpacket_collection): nus = [3.0e15, 0.0, 1e15, 1e5] energies = [0.4, 0.1, 0.6, 1e10] - initial_mus = [.1, 0, 1, .9] + initial_mus = [0.1, 0, 1, 0.9] initial_rs = [3e42, 4.5e45, 0, 9.0e40] last_interaction_in_nus = np.array( [3.0e15, 0.0, 1e15, 1e5], dtype=np.float64 diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py index 8600f897663..c227a87feff 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py @@ -50,7 +50,7 @@ def estimators(): stim_recomb_estimator=np.empty((0, 0), dtype=np.float64), bf_heating_estimator=np.empty((0, 0), dtype=np.float64), stim_recomb_cooling_estimator=np.empty((0, 0), dtype=np.float64), - photo_ion_estimator_statistics=np.empty((0, 0), dtype=np.int64) + photo_ion_estimator_statistics=np.empty((0, 0), dtype=np.int64), ) @@ -254,6 +254,7 @@ def test_update_line_estimators( assert_allclose(estimators.j_blue_estimator, expected_j_blue) assert_allclose(estimators.Edotlu_estimator, expected_Edotlu) + # TODO set RNG consistently # TODO: update this test to use the correct trace_packet @pytest.mark.xfail(reason="Need to fix estimator differences across runs") @@ -324,7 +325,9 @@ def test_move_r_packet( packet.r, packet.mu, model.time_explosion ) - r_packet_transport.move_r_packet(packet, distance, model.time_explosion, estimators) + r_packet_transport.move_r_packet( + packet, distance, model.time_explosion, estimators + ) assert_almost_equal(packet.mu, expected_params["mu"]) assert_almost_equal(packet.r, expected_params["r"]) @@ -402,5 +405,3 @@ def test_move_packet_across_shell_boundary_increment( packet, delta_shell, no_of_shells ) assert packet.current_shell_id == current_shell_id + delta_shell - - diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_single_packet_loop.py b/tardis/montecarlo/montecarlo_numba/tests/test_single_packet_loop.py index 2cf7f70f4ee..58f4e03d304 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_single_packet_loop.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_single_packet_loop.py @@ -7,7 +7,8 @@ single_packet_loop, ) -@pytest.mark.xfail(reason='Need to fix estimator differences across runs') + +@pytest.mark.xfail(reason="Need to fix estimator differences across runs") # TODO set RNG consistently def test_verysimple_single_packet_loop( verysimple_numba_model, diff --git a/tardis/montecarlo/montecarlo_numba/vpacket.py b/tardis/montecarlo/montecarlo_numba/vpacket.py index c6890709e6e..f7d11724cbe 100644 --- a/tardis/montecarlo/montecarlo_numba/vpacket.py +++ b/tardis/montecarlo/montecarlo_numba/vpacket.py @@ -13,8 +13,9 @@ from tardis.montecarlo.montecarlo_numba.r_packet import ( PacketStatus, ) -from tardis.montecarlo.montecarlo_numba.r_packet_transport import \ - move_packet_across_shell_boundary +from tardis.montecarlo.montecarlo_numba.r_packet_transport import ( + move_packet_across_shell_boundary, +) from tardis.montecarlo.montecarlo_numba.calculate_distances import ( calculate_distance_boundary, @@ -83,7 +84,7 @@ def trace_vpacket_within_shell(v_packet, numba_model, numba_plasma): v_packet.current_shell_id ] chi_e = cur_electron_density * SIGMA_THOMSON - + # Calculating doppler factor doppler_factor = get_doppler_factor( v_packet.r, v_packet.mu, numba_model.time_explosion @@ -94,8 +95,7 @@ def trace_vpacket_within_shell(v_packet, numba_model, numba_plasma): tau_continuum = chi_continuum * distance_boundary tau_trace_combined = tau_continuum - - + cur_line_id = start_line_id for cur_line_id in range(start_line_id, len(numba_plasma.line_list_nu)): @@ -259,8 +259,8 @@ def trace_vpacket_volley( v_packet_energy = r_packet.energy * weight * doppler_factor_ratio # TODO: Make sure we have a new continuum object for each vpacket - #comov_nu = v_packet_nu * v_packet_doppler_factor - #continuum.calculate(comov_nu, r_packet.current_shell_id) + # comov_nu = v_packet_nu * v_packet_doppler_factor + # continuum.calculate(comov_nu, r_packet.current_shell_id) v_packet = VPacket( r_packet.r, diff --git a/tardis/montecarlo/packet_source.py b/tardis/montecarlo/packet_source.py index d801f3e39bb..78bcc275b8c 100644 --- a/tardis/montecarlo/packet_source.py +++ b/tardis/montecarlo/packet_source.py @@ -88,7 +88,7 @@ def create_blackbody_packet_nus(T, no_of_packets, rng, l_samples=1000): """ l_samples = l_samples l_array = np.cumsum(np.arange(1, l_samples, dtype=np.float64) ** -4) - l_coef = np.pi ** 4 / 90.0 + l_coef = np.pi**4 / 90.0 # For testing purposes if montecarlo_configuration.LEGACY_MODE_ENABLED: diff --git a/tardis/montecarlo/spectrum.py b/tardis/montecarlo/spectrum.py index 57639a81112..2479206bfe3 100644 --- a/tardis/montecarlo/spectrum.py +++ b/tardis/montecarlo/spectrum.py @@ -37,8 +37,10 @@ def __init__(self, _frequency, luminosity): self._frequency = _frequency.to("Hz", u.spectral()) self.luminosity = luminosity.to("erg / s") - l_nu_unit = u.def_unit('erg\ s^{-1}\ Hz^{-1}', u.Unit('erg/(s Hz)')) - l_lambda_unit = u.def_unit('erg\ s^{-1}\ \\AA^{-1}', u.Unit('erg/(s AA)')) + l_nu_unit = u.def_unit("erg\ s^{-1}\ Hz^{-1}", u.Unit("erg/(s Hz)")) + l_lambda_unit = u.def_unit( + "erg\ s^{-1}\ \\AA^{-1}", u.Unit("erg/(s AA)") + ) self.frequency = self._frequency[:-1] self.delta_frequency = self._frequency[1] - self._frequency[0] @@ -119,13 +121,17 @@ def plot(self, ax=None, mode="wavelength", **kwargs): `matplotlib documentation `_ for a list of all possible arguments. - """ + """ if ax is None: from matplotlib.pyplot import gca ax = gca() if mode == "wavelength": - ax.plot(self.wavelength.value, self.luminosity_density_lambda.value, **kwargs) + ax.plot( + self.wavelength.value, + self.luminosity_density_lambda.value, + **kwargs, + ) ax.set_xlabel( f"Wavelength [{self.wavelength.unit.to_string('latex_inline')}]" ) @@ -133,7 +139,9 @@ def plot(self, ax=None, mode="wavelength", **kwargs): f"$L_\\lambda$ [{self.luminosity_density_lambda.unit.to_string('latex_inline')}]" ) elif mode == "frequency": - ax.plot(self.frequency.value, self.luminosity_density_nu.value, **kwargs) + ax.plot( + self.frequency.value, self.luminosity_density_nu.value, **kwargs + ) ax.set_xlabel( f"Frequency [{self.frequency.unit.to_string('latex_inline')}]" ) diff --git a/tardis/montecarlo/tests/conftest.py b/tardis/montecarlo/tests/conftest.py index bad1613fd68..26a498689f2 100644 --- a/tardis/montecarlo/tests/conftest.py +++ b/tardis/montecarlo/tests/conftest.py @@ -10,4 +10,3 @@ def runner(): config = config_reader.Configuration.from_yaml(config_fname) runner = MontecarloRunner.from_config(config) return runner - diff --git a/tardis/montecarlo/tests/test_montecarlo.py b/tardis/montecarlo/tests/test_montecarlo.py index 195344880d6..a24b8c837fb 100644 --- a/tardis/montecarlo/tests/test_montecarlo.py +++ b/tardis/montecarlo/tests/test_montecarlo.py @@ -77,6 +77,7 @@ from tardis import __path__ as path + @pytest.fixture(scope="module") def continuum_compare_data_fname(): fname = "continuum_compare_data.hdf" @@ -108,6 +109,7 @@ def ff_emissivity(t_electron): return ff_emissivity + @pytest.fixture(scope="module") def ion_edges(): return [ @@ -128,6 +130,7 @@ def ion_edges(): }, ] + """ Important Tests: ---------------- @@ -152,7 +155,7 @@ def ion_edges(): 0.0, 0, 3, - {"result": 0}, # This one might need to check for a bounds error + {"result": 0}, # This one might need to check for a bounds error ), ], ) @@ -216,7 +219,7 @@ def test_line_search(nu, nu_insert, number_of_lines, expected_params): 8.0, 0, 3, - {"result": 0}, # this one too + {"result": 0}, # this one too ), ( [2.0, 4.0, 6.0, 7.0], @@ -230,9 +233,7 @@ def test_line_search(nu, nu_insert, number_of_lines, expected_params): def test_binary_search(x, x_insert, imin, imax, expected_params): obtained_result = 0 - obtained_result = formal_integral.binary_search( - x, x_insert, imin, imax - ) + obtained_result = formal_integral.binary_search(x, x_insert, imin, imax) assert obtained_result == expected_params["result"] @@ -313,9 +314,7 @@ def test_unphysical_inverse_doppler_factor(mu, r, inv_t_exp): (0, 1, 1 / 2.6e7, 1.0), ], ) -def test_get_doppler_factor_full_relativity( - mu, r, inv_t_exp, expected -): +def test_get_doppler_factor_full_relativity(mu, r, inv_t_exp, expected): # Set the params from test cases here # TODO: add relativity tests mc.full_relativity = True @@ -338,9 +337,7 @@ def test_get_doppler_factor_full_relativity( (0, 1, 1 / 2.6e7, 1.0), ], ) -def test_get_inverse_doppler_factor_full_relativity( - mu, r, inv_t_exp, expected -): +def test_get_inverse_doppler_factor_full_relativity(mu, r, inv_t_exp, expected): # Set the params from test cases here # TODO: add relativity tests mc.full_relativity = True @@ -401,9 +398,7 @@ def test_both_angle_aberrations(mu, r, time_explosion): energy = 0.9 packet = r_packet.RPacket(r, mu, nu, energy) packet.r = r - obtained_mu = angle_aberration_LF_to_CMF( - packet, time_explosion, mu - ) + obtained_mu = angle_aberration_LF_to_CMF(packet, time_explosion, mu) inverse_obtained_mu = angle_aberration_CMF_to_LF( packet, time_explosion, obtained_mu ) @@ -422,9 +417,7 @@ def test_both_angle_aberrations_inverse(mu, r, time_explosion): energy = 0.9 packet = r_packet.RPacket(r, mu, nu, energy) packet.r = r - obtained_mu = angle_aberration_CMF_to_LF( - packet, time_explosion, mu - ) + obtained_mu = angle_aberration_CMF_to_LF(packet, time_explosion, mu) inverse_obtained_mu = angle_aberration_LF_to_CMF( packet, time_explosion, obtained_mu ) @@ -493,7 +486,7 @@ def test_move_packet_across_shell_boundary_increment( [(0, 1, 0, 0), (0, 1, 1, 0), (0, 1, 0, 1)], ) def test_packet_energy_limit_one(distance_trace, time_explosion, mu, r): - initial_energy = 0.9 + initial_energy = 0.9 nu = 0.4 packet = r_packet.RPacket(r, mu, nu, initial_energy) new_energy = r_packet.calc_packet_energy( @@ -510,9 +503,7 @@ def test_packet_energy_limit_one(distance_trace, time_explosion, mu, r): ({"mu": -0.3, "r": 7.5e14}, {"d_boundary": 709376919351035.9}), ], ) -def test_compute_distance2boundary( - packet_params, expected_params -): +def test_compute_distance2boundary(packet_params, expected_params): mu = packet_params["mu"] r = packet_params["r"] r_inner = np.array([6.912e14, 8.64e14], dtype=np.float64) @@ -561,9 +552,7 @@ def test_compute_distance2line(packet_params, expected_params): time_explosion = 5.2e7 - doppler_factor = get_doppler_factor( - packet.r, packet.mu, time_explosion - ) + doppler_factor = get_doppler_factor(packet.r, packet.mu, time_explosion) comov_nu = packet.nu * doppler_factor d_line = 0 @@ -620,9 +609,7 @@ def test_compute_distance2line(packet_params, expected_params): ), ], ) -def test_move_packet( - packet_params, expected_params, full_relativity -): +def test_move_packet(packet_params, expected_params, full_relativity): distance = 1e13 r, mu, nu, energy = 7.5e14, 0.3, 0.4, 0.9 time_explosion = 5.2e7 @@ -634,9 +621,7 @@ def test_move_packet( # model.full_relativity = full_relativity mc.full_relativity = full_relativity - doppler_factor = get_doppler_factor( - packet.r, packet.mu, time_explosion - ) + doppler_factor = get_doppler_factor(packet.r, packet.mu, time_explosion) numba_estimator = Estimators( packet_params["j"], packet_params["nu_bar"], 0, 0 ) @@ -769,9 +754,7 @@ def test_angle_transformation_invariance(mu, r, inv_t_exp): mc.full_relativity = True mu1 = angle_aberration_CMF_to_LF(packet, 1 / inv_t_exp, mu) - mu_obtained = angle_aberration_LF_to_CMF( - packet, 1 / inv_t_exp, mu1 - ) + mu_obtained = angle_aberration_LF_to_CMF(packet, 1 / inv_t_exp, mu1) mc.full_relativity = False assert_almost_equal(mu_obtained, mu) @@ -860,4 +843,4 @@ def test_rpacket_tracking(index, seed, r, nu, mu, energy): assert test_rpacket.nu == tracked_rpacket_properties.nu assert test_rpacket.mu == tracked_rpacket_properties.mu assert test_rpacket.energy == tracked_rpacket_properties.energy - assert tracked_rpacket_properties.interact_id == 1 + assert tracked_rpacket_properties.interact_id == 1 diff --git a/tardis/montecarlo/tests/test_spectrum.py b/tardis/montecarlo/tests/test_spectrum.py index 2d94ff6de4a..418317f822c 100644 --- a/tardis/montecarlo/tests/test_spectrum.py +++ b/tardis/montecarlo/tests/test_spectrum.py @@ -113,7 +113,7 @@ def test_luminosity_to_flux(): def test_f_nu_to_f_lambda(spectrum): expected = ( spectrum.luminosity_density_nu - * spectrum.frequency ** 2 + * spectrum.frequency**2 / c.c.to("angstrom/s") ).to("erg / (s angstrom)") print(expected) @@ -122,7 +122,7 @@ def test_f_nu_to_f_lambda(spectrum): ) expected = ( spectrum.luminosity_density_nu - * spectrum.frequency ** 2 + * spectrum.frequency**2 / c.c.to("angstrom/s") ).to("erg / (s angstrom)") np.testing.assert_allclose( diff --git a/tardis/plasma/base.py b/tardis/plasma/base.py index e6612a3e565..4a060b631c8 100644 --- a/tardis/plasma/base.py +++ b/tardis/plasma/base.py @@ -410,7 +410,7 @@ def write_to_tex(self, fname_graph, scale=0.5, args=None, latex_label=True): r"\begin{tikzpicture}[>=latex',line join=bevel,]", r"\begin{tikzpicture}" r"[>=latex',line join=bevel," - fr"scale={scale}]", + rf"scale={scale}]", ), end="", ) diff --git a/tardis/plasma/properties/atomic.py b/tardis/plasma/properties/atomic.py index 118294b3105..99dc1068231 100644 --- a/tardis/plasma/properties/atomic.py +++ b/tardis/plasma/properties/atomic.py @@ -177,9 +177,9 @@ class PhotoIonizationData(ProcessingPlasmaProperty): ) def calculate(self, atomic_data, continuum_interaction_species): - #photoionization_data = atomic_data.photoionization_data.set_index( + # photoionization_data = atomic_data.photoionization_data.set_index( # ["atomic_number", "ion_number", "level_number"] - #) + # ) photoionization_data = atomic_data.photoionization_data mask_selected_species = photoionization_data.index.droplevel( "level_number" @@ -754,7 +754,7 @@ def calculate_yg_van_regemorter( nu_lines = lines_filtered.nu.values yg = f_lu * (I_H / (H * nu_lines)) ** 2 - coll_const = A0 ** 2 * np.pi * np.sqrt(8 * K_B / (np.pi * M_E)) + coll_const = A0**2 * np.pi * np.sqrt(8 * K_B / (np.pi * M_E)) yg = 14.5 * coll_const * t_electrons * yg[:, np.newaxis] u0 = nu_lines[np.newaxis].T / t_electrons * (H / K_B) diff --git a/tardis/plasma/properties/continuum_processes.py b/tardis/plasma/properties/continuum_processes.py index b7262426be3..25efa9c5d57 100644 --- a/tardis/plasma/properties/continuum_processes.py +++ b/tardis/plasma/properties/continuum_processes.py @@ -57,15 +57,15 @@ A0 = const.a0.cgs.value M_E = const.m_e.cgs.value E = const.e.esu.value -BETA_COLL = (H ** 4 / (8 * K_B * M_E ** 3 * np.pi ** 3)) ** 0.5 +BETA_COLL = (H**4 / (8 * K_B * M_E**3 * np.pi**3)) ** 0.5 F_K = ( 16 / (3.0 * np.sqrt(3)) - * np.sqrt((2 * np.pi) ** 3 * K_B / (H ** 2 * M_E ** 3)) - * (E ** 2 / C) ** 3 + * np.sqrt((2 * np.pi) ** 3 * K_B / (H**2 * M_E**3)) + * (E**2 / C) ** 3 ) # See Eq. 19 in Sutherland, R. S. 1998, MNRAS, 300, 321 FF_OPAC_CONST = ( - (2 * np.pi / (3 * M_E * K_B)) ** 0.5 * 4 * E ** 6 / (3 * M_E * H * C) + (2 * np.pi / (3 * M_E * K_B)) ** 0.5 * 4 * E**6 / (3 * M_E * H * C) ) # See Eq. 6.1.8 in http://personal.psu.edu/rbc3/A534/lec6.pdf logger = logging.getLogger(__name__) @@ -316,7 +316,7 @@ def calculate( x_sect = photo_ion_cross_sections["x_sect"].values nu = photo_ion_cross_sections["nu"].values - alpha_sp = 8 * np.pi * x_sect * nu ** 2 / C ** 2 + alpha_sp = 8 * np.pi * x_sect * nu**2 / C**2 alpha_sp = alpha_sp[:, np.newaxis] alpha_sp = alpha_sp * boltzmann_factor_photo_ion alpha_sp = integrate_array_by_blocks( @@ -351,7 +351,7 @@ def calculate( x_sect = photo_ion_cross_sections["x_sect"].values nu = photo_ion_cross_sections["nu"].values factor = (1 - nu_i / photo_ion_cross_sections["nu"]).values - alpha_sp = (8 * np.pi * x_sect * factor * nu ** 3 / C ** 2) * H + alpha_sp = (8 * np.pi * x_sect * factor * nu**3 / C**2) * H alpha_sp = alpha_sp[:, np.newaxis] alpha_sp = alpha_sp * boltzmann_factor_photo_ion alpha_sp = integrate_array_by_blocks( @@ -389,7 +389,7 @@ def calculate( nu = photo_ion_cross_sections["nu"].values # alpha_sp_E will be missing a lot of prefactors since we are only # interested in relative values here - alpha_sp_E = nu ** 3 * x_sect + alpha_sp_E = nu**3 * x_sect alpha_sp_E = alpha_sp_E[:, np.newaxis] alpha_sp_E = alpha_sp_E * boltzmann_factor_photo_ion alpha_sp_E = cumulative_integrate_array_by_blocks( @@ -940,7 +940,7 @@ def _calculate_ff_cooling_factor(ion_number_density, electron_densities): ion_charge = ion_number_density.index.get_level_values(1).values factor = ( electron_densities - * ion_number_density.multiply(ion_charge ** 2, axis=0).sum() + * ion_number_density.multiply(ion_charge**2, axis=0).sum() ) return factor @@ -965,7 +965,7 @@ def chi_ff(nu, shell): chi_ff = ( FF_OPAC_CONST * ff_opacity_factor[shell] - / nu ** 3 + / nu**3 * (1 - np.exp(-H * nu / (K_B * t_electrons[shell]))) ) return chi_ff @@ -1193,7 +1193,7 @@ def chi_bf_interpolator(nu, shell): Photoionization cross-sections of all bound-free continua for which absorption is possible for frequency `nu`. """ - + current_continua = get_current_bound_free_continua(nu) chi_bfs = np.zeros(len(current_continua)) x_sect_bfs = np.zeros(len(current_continua)) @@ -1202,15 +1202,23 @@ def chi_bf_interpolator(nu, shell): end = photo_ion_block_references[continuum_id + 1] phot_nus_continuum = phot_nus[start:end] nu_idx = np.searchsorted(phot_nus_continuum, nu) - interval = phot_nus_continuum[nu_idx] - phot_nus_continuum[nu_idx-1] - high_weight = (nu - phot_nus_continuum[nu_idx-1]) - low_weight = (phot_nus_continuum[nu_idx] - nu) + interval = ( + phot_nus_continuum[nu_idx] - phot_nus_continuum[nu_idx - 1] + ) + high_weight = nu - phot_nus_continuum[nu_idx - 1] + low_weight = phot_nus_continuum[nu_idx] - nu chi_bfs_continuum = chi_bf[start:end, shell] - chi_bfs[i] = (chi_bfs_continuum[nu_idx]*high_weight + chi_bfs_continuum[nu_idx-1]*low_weight)/interval + chi_bfs[i] = ( + chi_bfs_continuum[nu_idx] * high_weight + + chi_bfs_continuum[nu_idx - 1] * low_weight + ) / interval x_sect_bfs_continuum = x_sect[start:end] - x_sect_bfs[i] = (x_sect_bfs_continuum[nu_idx]*high_weight + x_sect_bfs_continuum[nu_idx-1]*low_weight)/interval + x_sect_bfs[i] = ( + x_sect_bfs_continuum[nu_idx] * high_weight + + x_sect_bfs_continuum[nu_idx - 1] * low_weight + ) / interval chi_bf_contributions = chi_bfs.cumsum() - + # If we are outside the range of frequencies # for which we have photo-ionization cross sections # we will have no local continuua and therefore diff --git a/tardis/plasma/properties/general.py b/tardis/plasma/properties/general.py index 8560529883d..040b860b53a 100644 --- a/tardis/plasma/properties/general.py +++ b/tardis/plasma/properties/general.py @@ -57,7 +57,7 @@ class GElectron(ProcessingPlasmaProperty): def calculate(self, beta_rad): return ( (2 * np.pi * const.m_e.cgs.value / beta_rad) - / (const.h.cgs.value ** 2) + / (const.h.cgs.value**2) ) ** 1.5 @@ -149,7 +149,7 @@ class LuminosityInner(ProcessingPlasmaProperty): @staticmethod def calculate(r_inner, t_inner): return ( - 4 * np.pi * const.sigma_sb.cgs * r_inner[0] ** 2 * t_inner ** 4 + 4 * np.pi * const.sigma_sb.cgs * r_inner[0] ** 2 * t_inner**4 ).to("erg/s") diff --git a/tardis/plasma/properties/nlte.py b/tardis/plasma/properties/nlte.py index 47b2a7e82c9..9989561c493 100644 --- a/tardis/plasma/properties/nlte.py +++ b/tardis/plasma/properties/nlte.py @@ -119,7 +119,7 @@ def calculate_helium_one( level_boltzmann_factor.loc[2, 0] * (1.0 / (2 * g.loc[2, 1, 0])) * (1 / g_electron) - * (1 / (w ** 2.0)) + * (1 / (w**2.0)) * np.exp(ionization_data.loc[2, 1] * beta_rad) ) diff --git a/tardis/plasma/properties/partition_function.py b/tardis/plasma/properties/partition_function.py index 4e00332d7cd..0fafaeee976 100644 --- a/tardis/plasma/properties/partition_function.py +++ b/tardis/plasma/properties/partition_function.py @@ -208,7 +208,7 @@ def _main_nlte_calculation( dtype=np.float64, ) r_ul_matrix_reshaped = r_ul_matrix.reshape( - (number_of_levels ** 2, len(t_electrons)) + (number_of_levels**2, len(t_electrons)) ) r_ul_matrix_reshaped[r_ul_index] = ( A_uls[np.newaxis].T + B_uls[np.newaxis].T * j_blues_filtered @@ -216,7 +216,7 @@ def _main_nlte_calculation( r_ul_matrix_reshaped[r_ul_index] *= beta_sobolevs_filtered r_lu_matrix = np.zeros_like(r_ul_matrix) r_lu_matrix_reshaped = r_lu_matrix.reshape( - (number_of_levels ** 2, len(t_electrons)) + (number_of_levels**2, len(t_electrons)) ) r_lu_matrix_reshaped[r_lu_index] = ( B_lus[np.newaxis].T * j_blues_filtered * beta_sobolevs_filtered diff --git a/tardis/plasma/properties/radiative_properties.py b/tardis/plasma/properties/radiative_properties.py index ef40885c136..f3b5671de46 100644 --- a/tardis/plasma/properties/radiative_properties.py +++ b/tardis/plasma/properties/radiative_properties.py @@ -135,10 +135,10 @@ def __init__(self, plasma_parent): super(TauSobolev, self).__init__(plasma_parent) self.sobolev_coefficient = ( ( - ((np.pi * const.e.gauss ** 2) / (const.m_e.cgs * const.c.cgs)) + ((np.pi * const.e.gauss**2) / (const.m_e.cgs * const.c.cgs)) * u.cm * u.s - / u.cm ** 3 + / u.cm**3 ) .to(1) .value diff --git a/tardis/plasma/properties/util/macro_atom.py b/tardis/plasma/properties/util/macro_atom.py index 68b141fa51a..b89971776ff 100644 --- a/tardis/plasma/properties/util/macro_atom.py +++ b/tardis/plasma/properties/util/macro_atom.py @@ -6,7 +6,7 @@ h_cgs = const.h.cgs.value c = const.c.to("cm/s").value kb = const.k_B.cgs.value -inv_c2 = 1 / (c ** 2) +inv_c2 = 1 / (c**2) @njit(**njit_dict) diff --git a/tardis/scripts/cmfgen2tardis.py b/tardis/scripts/cmfgen2tardis.py index 9c4c2c77bf5..7f01fba3940 100644 --- a/tardis/scripts/cmfgen2tardis.py +++ b/tardis/scripts/cmfgen2tardis.py @@ -79,7 +79,7 @@ def convert_format(file_path): abundances_df.loc[element_symbol] = abundances density_df = pd.DataFrame.from_records( - [velocity, temperature * 10 ** 4, density, electron_density] + [velocity, temperature * 10**4, density, electron_density] ).transpose() density_df.columns = [ "velocity", diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index fd7b6df73c0..b0cac4eb8dd 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -193,7 +193,8 @@ def __init__( self._cb_next_id = 0 mc_config_module.CONTINUUM_PROCESSES_ENABLED = ( - not self.plasma.continuum_interaction_species.empty) + not self.plasma.continuum_interaction_species.empty + ) def estimate_t_inner( self, input_t_inner, luminosity_requested, t_inner_update_exponent=-0.5 @@ -206,7 +207,7 @@ def estimate_t_inner( (emitted_luminosity / luminosity_requested).to(1).value ) - return input_t_inner * luminosity_ratios ** t_inner_update_exponent + return input_t_inner * luminosity_ratios**t_inner_update_exponent @staticmethod def damped_converge(value, estimated_value, damping_factor): diff --git a/tardis/tests/test_util.py b/tardis/tests/test_util.py index 9fa47d19c40..e4814b321d5 100644 --- a/tardis/tests/test_util.py +++ b/tardis/tests/test_util.py @@ -119,10 +119,10 @@ def test_calculate_luminosity(string_io, distance, result): @pytest.mark.parametrize( ["nu", "t", "i"], [ - (10 ** 6, 1000, 3.072357852080765e-22), - (10 ** 6, 300, 9.21707305730458e-23), - (10 ** 8, 1000, 6.1562660718558254e-24), - (10 ** 8, 300, 1.846869480674048e-24), + (10**6, 1000, 3.072357852080765e-22), + (10**6, 300, 9.21707305730458e-23), + (10**8, 1000, 6.1562660718558254e-24), + (10**8, 300, 1.846869480674048e-24), ], ) def test_intensity_black_body(nu, t, i): diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py index 8d4b6e44731..4dffec02d91 100644 --- a/tardis/visualization/tools/sdec_plot.py +++ b/tardis/visualization/tools/sdec_plot.py @@ -1066,13 +1066,13 @@ def _calculate_photosphere_luminosity(self, packets_mode): """ bb_lam = BlackBody( self.data[packets_mode].t_inner, - scale=1.0 * u.erg / (u.cm ** 2 * u.AA * u.s * u.sr) + scale=1.0 * u.erg / (u.cm**2 * u.AA * u.s * u.sr), ) L_lambda_ph = ( bb_lam(self.plot_wavelength) * 4 - * np.pi ** 2 + * np.pi**2 * self.data[packets_mode].r_inner[0] ** 2 * u.sr ).to("erg / (AA s)") diff --git a/tardis/visualization/tools/tests/test_sdec_plot.py b/tardis/visualization/tools/tests/test_sdec_plot.py index 1a4fd79f4ec..5e62044d7c0 100644 --- a/tardis/visualization/tools/tests/test_sdec_plot.py +++ b/tardis/visualization/tools/tests/test_sdec_plot.py @@ -143,7 +143,7 @@ def observed_spectrum(self): observed_spectrum_wavelength, observed_spectrum_flux = test_data.T observed_spectrum_wavelength = observed_spectrum_wavelength * u.AA observed_spectrum_flux = ( - observed_spectrum_flux * u.erg / (u.s * u.cm ** 2 * u.AA) + observed_spectrum_flux * u.erg / (u.s * u.cm**2 * u.AA) ) return observed_spectrum_wavelength, observed_spectrum_flux