diff --git a/docs/development/continuous_integration.rst b/docs/development/continuous_integration.rst index 09d6e62c3f4..de4adf34f33 100644 --- a/docs/development/continuous_integration.rst +++ b/docs/development/continuous_integration.rst @@ -172,6 +172,9 @@ to paths in Azure: See the Azure documentation to learn more about `checking out multiple repositories`_. +Azure provides a list of agent hosts that can run the pipeline on a virtual machine. In our pipelines, we +use the vm_Images: Ubuntu 16.04 and macOs-10.14. + Jobs ---- diff --git a/docs/development/debug_numba.rst b/docs/development/debug_numba.rst new file mode 100644 index 00000000000..2f5e1c210ef --- /dev/null +++ b/docs/development/debug_numba.rst @@ -0,0 +1,14 @@ +************************** +Debugging numba_montecarlo +************************** +To facilitate more in-depth debugging when interfacing with the `montecarlo_numba` +module, we provide a set of debugging configurations. PyCharm debugging +configurations, in addition to related scripts and .yml files, are contained in +`tardis.scripts.debug`. Currently, these include the ability to run TARDIS +in asingle-packet mode, with the packet seed identified at debug time. +`tardis_example_single.yml` is the configuration filethat is used to set up the +single-packet TARDIS run; `run_numba_single.py` is thePython script that runs +this .yml file; `run_numba_single.xml` is the PyCharmdebug configuration file +that can be used in conjunction with the above files. + +Note that this method is EXPERIMENTAL. diff --git a/docs/development/index.rst b/docs/development/index.rst index 1c253696d18..61bec8a18fd 100644 --- a/docs/development/index.rst +++ b/docs/development/index.rst @@ -11,8 +11,8 @@ to the Astropy team for designing it. .. toctree:: :maxdepth: 2 - running_tests issues + debug_numba diff --git a/docs/graphics/numba_comparison_real_quickstart.png b/docs/graphics/numba_comparison_real_quickstart.png new file mode 100644 index 00000000000..e1ced13fc3a Binary files /dev/null and b/docs/graphics/numba_comparison_real_quickstart.png differ diff --git a/docs/graphics/numba_comparison_virtual_quickstart.png b/docs/graphics/numba_comparison_virtual_quickstart.png new file mode 100644 index 00000000000..1c5cc3ae3df Binary files /dev/null and b/docs/graphics/numba_comparison_virtual_quickstart.png differ diff --git a/docs/index.rst b/docs/index.rst index d4d95b98488..62b92c4e393 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -69,6 +69,14 @@ Mission Statement Physics overview physics/plasma/index +.. toctree:: + :maxdepth: 2 + :caption: Testing TARDIS + :hidden: + + testing/automated_tests + testing/numba + .. toctree:: :maxdepth: 2 diff --git a/docs/running/custom_source.ipynb b/docs/running/custom_source.ipynb index 6a2e756d4ad..1516cd0734b 100644 --- a/docs/running/custom_source.ipynb +++ b/docs/running/custom_source.ipynb @@ -52,9 +52,10 @@ " \n", " def __init__(self, seed, truncation_wavelength):\n", " super().__init__(seed)\n", + " self.rng = np.random.default_rng(seed=seed)\n", " self.truncation_wavelength = truncation_wavelength\n", " \n", - " def create_packets(self, T, no_of_packets,\n", + " def create_packets(self, T, no_of_packets, rng,\n", " drawing_sample_size=None):\n", " \"\"\"\n", " Packet source that generates a truncated Blackbody source.\n", @@ -70,11 +71,10 @@ " Only wavelengths higher than the truncation wavelength\n", " will be sampled.\n", " \"\"\"\n", - " \n", - " \n", + "\n", " # Use mus and energies from normal blackbody source.\n", - " mus = self.create_zero_limb_darkening_packet_mus(no_of_packets)\n", - " energies = self.create_uniform_packet_energies(no_of_packets)\n", + " mus = self.create_zero_limb_darkening_packet_mus(no_of_packets, self.rng)\n", + " energies = self.create_uniform_packet_energies(no_of_packets, self.rng)\n", "\n", " # If not specified, draw 2 times as many packets and reject any beyond no_of_packets.\n", " if drawing_sample_size is None:\n", @@ -86,7 +86,7 @@ " \n", " # Draw nus from blackbody distribution and reject based on truncation_frequency.\n", " # If more nus.shape[0] > no_of_packets use only the first no_of_packets.\n", - " nus = self.create_blackbody_packet_nus(T, drawing_sample_size)\n", + " nus = self.create_blackbody_packet_nus(T, drawing_sample_size, self.rng)\n", " nus = nus[nus no_of_packets.\n", " while nus.shape[0] < no_of_packets:\n", " additional_nus = self.create_blackbody_packet_nus(\n", - " T, drawing_sample_size\n", + " T, drawing_sample_size, self.rng\n", " )\n", " mask = additional_nus < truncation_frequency\n", " additional_nus = additional_nus[mask][:no_of_packets]\n", @@ -170,4 +170,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} +} \ No newline at end of file diff --git a/tardis/gui/datahandler.py b/tardis/gui/datahandler.py index 4e9798a2761..0cb92c2fd4f 100644 --- a/tardis/gui/datahandler.py +++ b/tardis/gui/datahandler.py @@ -39,15 +39,15 @@ class Node(object): Attributes ---------- - parent: None/Node + parent: None/Node The parent of the node. children: list of Node The children of the node. - data: list of string + data: list of string The data stored on the node. Can be a key or a value. - siblings: dictionary - A dictionary of nodes that are siblings of this node. The - keys are the values of the nodes themselves. This is + siblings: dictionary + A dictionary of nodes that are siblings of this node. The + keys are the values of the nodes themselves. This is used to keep track of which value the user has selected if the parent of this node happens to be a key that can take values from a list. @@ -61,18 +61,18 @@ def __init__(self, data, parent=None): ---------- data: list of string The data that is intended to be stored on the node. - parent: Node - Another node which is the parent of this node. The + parent: Node + Another node which is the parent of this node. The root node has parent set to None. Note ---- - A leaf node is a node that is the only child of its parent. + A leaf node is a node that is the only child of its parent. For this tree this will always be the case. This is because the tree stores the key at every node except the leaf node where it stores the value for the key. So if in the dictionary - the value of a key is another dictionary, then it will - be a node with no leafs. If the key has a value that is a + the value of a key is another dictionary, then it will + be a node with no leafs. If the key has a value that is a value or a list then it will have one child that is a leaf. The leaf can have no children. For example- @@ -86,12 +86,12 @@ def __init__(self, data, parent=None): val7, val5 and val6 and is currently set to val5. - In the tree shown above all quoted values are keys in the - dictionary and are non-leaf nodes in the tree. All values - of the form valx are leaf nodes and are not dictionaries + In the tree shown above all quoted values are keys in the + dictionary and are non-leaf nodes in the tree. All values + of the form valx are leaf nodes and are not dictionaries themselves. If the keys have non-dictionary values then they have a leaf attached. And no leaf can have a child. - + """ self.parent = parent self.children = [] @@ -105,9 +105,9 @@ def append_child(self, child): child.parent = self def get_child(self, i): - """Get the ith child of this node. + """Get the ith child of this node. - No error is raised if the cild requested doesn't exist. A + No error is raised if the cild requested doesn't exist. A None is returned in such cases. """ @@ -141,7 +141,7 @@ def get_parent(self): return self.parent def get_index_of_self(self): - """Returns the number at which it comes in the list of its + """Returns the number at which it comes in the list of its parent's children. For root the index 0 is returned. """ @@ -170,10 +170,10 @@ class TreeModel(QtCore.QAbstractItemModel): ---------- root: Node Root node of the tree. - disabledNodes: list of Node + disabledNodes: list of Node List of leaf nodes that are not editable currently. - typenodes: list of Node - List of nodes that correspond to keys that set container + typenodes: list of Node + List of nodes that correspond to keys that set container types. Look at tardis configuration template. These are the nodes that have values that can be set from a list. @@ -184,9 +184,9 @@ def __init__(self, dictionary, parent=None): Parameters ---------- - dictionary: dictionary + dictionary: dictionary The dictionary that needs to be converted to the tree. - parent: None + parent: None Used to instantiate the QAbstractItemModel """ @@ -267,7 +267,7 @@ def headerData(self, section, orientation, role): return None def index(self, row, column, parent=QtCore.QModelIndex()): - """Create a model index for the given row and column. For a + """Create a model index for the given row and column. For a tree model, the row is the set of nodes with the same parents and the column indexes the data in the node. @@ -315,8 +315,8 @@ def parent(self, index): return self.createIndex(parentItem.get_index_of_self(), 0, parentItem) def rowCount(self, parent=QtCore.QModelIndex()): - """The number of rows for a given node. - + """The number of rows for a given node. + (The number of rows is just the number of children for a node.) """ @@ -325,7 +325,7 @@ def rowCount(self, parent=QtCore.QModelIndex()): return parentItem.num_hildren() def setData(self, index, value, role=QtCore.Qt.EditRole): - """Set the value as the data at the location pointed by the + """Set the value as the data at the location pointed by the index. """ @@ -360,7 +360,7 @@ def dict_to_tree(self, dictionary, root): ---------- dictionary: dictionary The dictionary that is to be converted to the tree. - root: Node + root: Node The root node of the tree. """ @@ -411,7 +411,7 @@ def tree_from_node(self, dictionary, root): self.typenodes.append(child) def dict_from_node(self, node): - """Take a node and convert the whole subtree rooted at it into a + """Take a node and convert the whole subtree rooted at it into a dictionary. """ @@ -429,7 +429,7 @@ def dict_from_node(self, node): class TreeDelegate(QtWidgets.QStyledItemDelegate): - """Create a custom delegate to modify the columnview that displays the + """Create a custom delegate to modify the columnview that displays the TreeModel. """ @@ -457,7 +457,7 @@ def createEditor(self, parent, option, index): def setModelData(self, editor, model, index): """Called when new data id set in the model. This is where the - siblings of type nodes are enabled or disabled according to the + siblings of type nodes are enabled or disabled according to the new choice made. """ diff --git a/tardis/gui/interface.py b/tardis/gui/interface.py index dce94639fd0..040499f21ab 100644 --- a/tardis/gui/interface.py +++ b/tardis/gui/interface.py @@ -28,15 +28,15 @@ def show(model): """Take an instance of tardis model and display it. - If IPython functions were successfully imported then QApplication instance - is created using get_app_qt4. This will only create the app instance if - it doesn't already exists. Otherwise it is explicitly started. + If IPython functions were successfully imported then QApplication instance + is created using get_app_qt4. This will only create the app instance if + it doesn't already exists. Otherwise it is explicitly started. - Then the mainwindow is created, the model is attached to it and its + Then the mainwindow is created, the model is attached to it and its show method is called. Finally the eventloop is started using IPython functions (which start them - consistently) if they were imported. Otherwise it is started explicitly. + consistently) if they were imported. Otherwise it is started explicitly. """ if importFailed: diff --git a/tardis/gui/widgets.py b/tardis/gui/widgets.py index 0b26f580d08..cb4bbd3cca0 100644 --- a/tardis/gui/widgets.py +++ b/tardis/gui/widgets.py @@ -667,7 +667,7 @@ def change_spectrum_to_spec_virtual_flux_angstrom(self): def change_spectrum_to_spec_flux_angstrom(self): """Change spectrum data back from virtual spectrum. (See the - method above).""" + method above).""" if self.model.runner.spectrum.luminosity_density_lambda is None: luminosity_density_lambda = np.zeros_like( self.model.runner.spectrum.wavelength @@ -946,7 +946,10 @@ def __init__(self, parent, wavelength_start, wavelength_end, tablecreator): self.setGeometry(180 + len(self.parent.line_info) * 20, 150, 250, 400) self.setWindowTitle( "Line Interaction: %.2f - %.2f (A) " - % (wavelength_start, wavelength_end,) + % ( + wavelength_start, + wavelength_end, + ) ) self.layout = QtWidgets.QVBoxLayout() packet_nu_line_interaction = analysis.LastLineInteraction.from_model( @@ -1155,8 +1158,10 @@ def __init__( self.line_interaction_analysis = line_interaction_analysis self.atom_data = atom_data self.lines_data = lines_data.reset_index().set_index("line_id") - line_interaction_species_group = line_interaction_analysis.last_line_in.groupby( - ["atomic_number", "ion_number"] + line_interaction_species_group = ( + line_interaction_analysis.last_line_in.groupby( + ["atomic_number", "ion_number"] + ) ) self.species_selected = sorted( line_interaction_species_group.groups.keys() diff --git a/tardis/io/config_reader.py b/tardis/io/config_reader.py index ad3ddb63d95..e6a6a699aa3 100644 --- a/tardis/io/config_reader.py +++ b/tardis/io/config_reader.py @@ -239,9 +239,10 @@ class ConfigWriterMixin(HDFWriterMixin): """ Overrides HDFWriterMixin to obtain HDF properties from configuration keys """ + def get_properties(self): data = yaml.dump(self) - data = pd.DataFrame(index=[0], data={'config': data}) + data = pd.DataFrame(index=[0], data={"config": data}) return data @@ -249,6 +250,7 @@ class Configuration(ConfigurationNameSpace, ConfigWriterMixin): """ Tardis configuration class """ + hdf_name = "simulation" @classmethod @@ -358,21 +360,21 @@ def quantity_representer(dumper, data): def cns_representer(dumper, data): """ - Represents Configuration as dict + Represents Configuration as dict - Parameters - ---------- + Parameters + ---------- - dumper : - YAML dumper object + dumper : + YAML dumper object - data : - ConfigurationNameSpace object + data : + ConfigurationNameSpace object - Returns - ------- + Returns + ------- - yaml dumper representation of Configuration as dict + yaml dumper representation of Configuration as dict """ return dumper.represent_dict(dict(data)) @@ -380,4 +382,4 @@ def cns_representer(dumper, data): yaml.add_representer(u.Quantity, quantity_representer) yaml.add_representer(ConfigurationNameSpace, cns_representer) -yaml.add_representer(Configuration, cns_representer) \ No newline at end of file +yaml.add_representer(Configuration, cns_representer) diff --git a/tardis/io/config_validator.py b/tardis/io/config_validator.py index 8b23aefc8c2..4324575e56f 100644 --- a/tardis/io/config_validator.py +++ b/tardis/io/config_validator.py @@ -40,11 +40,17 @@ def set_defaults(validator, properties, instance, schema): instance.setdefault(property, subschema["default"]) for error in validate_properties( - validator, properties, instance, schema, + validator, + properties, + instance, + schema, ): yield error - return validators.extend(validator_class, {"properties": set_defaults},) + return validators.extend( + validator_class, + {"properties": set_defaults}, + ) DefaultDraft4Validator = extend_with_default(Draft4Validator) diff --git a/tardis/io/decay.py b/tardis/io/decay.py index 59f85b62d75..4a118016de3 100644 --- a/tardis/io/decay.py +++ b/tardis/io/decay.py @@ -101,7 +101,7 @@ def decay(self, t): def as_atoms(self): """ - Merge Isotope dataframe according to atomic number + Merge Isotope dataframe according to atomic number Returns: : merged isotope abundances @@ -111,11 +111,11 @@ def as_atoms(self): def merge(self, other, normalize=True): """ - Merge Isotope dataframe with abundance passed as parameter + Merge Isotope dataframe with abundance passed as parameter Parameters ---------- - other: pd.DataFrame + other: pd.DataFrame normalize : bool If true, resultant dataframe will be normalized diff --git a/tardis/io/schemas/montecarlo.yml b/tardis/io/schemas/montecarlo.yml index b0638270e6c..1e5075c1d2b 100644 --- a/tardis/io/schemas/montecarlo.yml +++ b/tardis/io/schemas/montecarlo.yml @@ -33,21 +33,17 @@ properties: multipleOf: 1.0 default: 0 description: Setting the number of virtual packets for the last iteration. - virtual_spectrum_range: + virtual_spectrum_spawn_range: type: object default: {} properties: start: type: quantity - default: 50 angstrom - stop: + default: 0 angstrom + end: type: quantity - default: 250000 angstrom - num: - type: number - multipleOf: 1.0 - default: 1000000 - description: Limits of virtual packet spectrum (giving maximum and minimum packet + default: inf angstrom + description: Limits of virtual packet spawn spectrum (giving maximum and minimum packet frequency) enable_reflective_inner_boundary: type: boolean @@ -68,6 +64,19 @@ properties: default: false description: Enables a more complete treatment of relativitic effects. This includes angle aberration as well as use of the fully general Doppler formula. + debug_packets: + type: boolean + default: false + description: Decide whether to go into debugging mode. [EXPERIMENTAL FEATURE DO NOT RELY ON IT] + logger_buffer: + type: number + default: 1 + description: Provides option to not log every line. + single_packet_seed: + type: + - number + default: -1 + description: If debug_packets is true, this is the seed for the only packet. required: - no_of_packets diff --git a/tardis/io/schemas/plasma.yml b/tardis/io/schemas/plasma.yml index 56a45371cd6..59405000b74 100644 --- a/tardis/io/schemas/plasma.yml +++ b/tardis/io/schemas/plasma.yml @@ -17,6 +17,11 @@ properties: default: false description: disable electron scattering process in montecarlo part - non-physical only for tests + disable_line_scattering: + type: boolean + default: false + description: disable line scattering process in montecarlo part - non-physical + only for tests ionization: type: string enum: diff --git a/tardis/io/schemas/spectrum.yml b/tardis/io/schemas/spectrum.yml index 8439e75c041..96c9e1534ee 100644 --- a/tardis/io/schemas/spectrum.yml +++ b/tardis/io/schemas/spectrum.yml @@ -50,3 +50,7 @@ properties: default: False description: If True bias v-packet emission based on the electron scattering optical depth + virtual_packet_logging: + type: boolean + default: False + description: If True, enable virtual packet logging output diff --git a/tardis/io/tests/test_config_reader.py b/tardis/io/tests/test_config_reader.py index 4cdc83aeb75..2345329cd69 100644 --- a/tardis/io/tests/test_config_reader.py +++ b/tardis/io/tests/test_config_reader.py @@ -62,9 +62,9 @@ def test_from_config_dict(tardis_config_verysimple): def test_config_hdf(hdf_file_path, tardis_config_verysimple): expected = Configuration.from_config_dict( - tardis_config_verysimple, validate=True, config_dirname="test" - ) + tardis_config_verysimple, validate=True, config_dirname="test" + ) expected.to_hdf(hdf_file_path) actual = pd.read_hdf(hdf_file_path, key="/simulation/config") - expected = expected.get_properties()['config'] + expected = expected.get_properties()["config"] assert actual[0] == expected[0] diff --git a/tardis/io/util.py b/tardis/io/util.py index b75253d7830..7430b0d7850 100644 --- a/tardis/io/util.py +++ b/tardis/io/util.py @@ -284,7 +284,7 @@ def get_properties(self): @property def full_hdf_properties(self): # If tardis was compiled --with-vpacket-logging, add vpacket properties - if hasattr(self, "virt_logging") and self.virt_logging == 1: + if hasattr(self, "virt_logging") and self.virt_logging: self.hdf_properties.extend(self.vpacket_hdf_properties) return self.optional_hdf_properties + self.hdf_properties diff --git a/tardis/montecarlo/base.py b/tardis/montecarlo/base.py index 7a2f5a89039..dced2fc7512 100644 --- a/tardis/montecarlo/base.py +++ b/tardis/montecarlo/base.py @@ -10,49 +10,92 @@ from tardis.util.base import quantity_linspace from tardis.io.util import HDFWriterMixin -from tardis.montecarlo import montecarlo, packet_source as source +from tardis.montecarlo import packet_source as source from tardis.montecarlo.formal_integral import FormalIntegrator +from tardis.montecarlo import montecarlo_configuration as mc_config_module + + +from tardis.montecarlo.montecarlo_numba import montecarlo_radial1d +from tardis.montecarlo.montecarlo_numba import montecarlo_logger as mc_logger +from tardis.montecarlo.montecarlo_numba.numba_interface import ( + configuration_initialize, +) +from tardis.montecarlo.montecarlo_numba import numba_config import numpy as np logger = logging.getLogger(__name__) +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 +# seed val to the maximum allowed by numpy. +# TODO: refactor this into more parts class MontecarloRunner(HDFWriterMixin): """ This class is designed as an interface between the Python part and the montecarlo C-part """ - hdf_properties = ['output_nu', 'output_energy', 'nu_bar_estimator', - 'j_estimator', 'montecarlo_virtual_luminosity', - 'last_interaction_in_nu', - 'last_interaction_type', - 'last_line_interaction_in_id', - 'last_line_interaction_out_id', - 'last_line_interaction_shell_id', - 'packet_luminosity', 'spectrum', - 'spectrum_virtual', 'spectrum_reabsorbed'] - - vpacket_hdf_properties = ["virt_packet_last_interaction_in_nu", - "virt_packet_last_interaction_type", - "virt_packet_last_line_interaction_in_id", - "virt_packet_last_line_interaction_out_id", - "virt_packet_nus", "virt_packet_energies"] - - hdf_name = 'runner' - w_estimator_constant = ((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)).cgs.value - - def __init__(self, seed, spectrum_frequency, virtual_spectrum_range, - sigma_thomson, enable_reflective_inner_boundary, - enable_full_relativity, inner_boundary_albedo, - line_interaction_type, integrator_settings, - v_packet_settings, spectrum_method, - packet_source=None): + + hdf_properties = [ + "output_nu", + "output_energy", + "nu_bar_estimator", + "j_estimator", + "montecarlo_virtual_luminosity", + "last_interaction_in_nu", + "last_interaction_type", + "last_line_interaction_in_id", + "last_line_interaction_out_id", + "last_line_interaction_shell_id", + "packet_luminosity", + "spectrum", + "spectrum_virtual", + "spectrum_reabsorbed", + ] + + vpacket_hdf_properties = [ + "virt_packet_last_interaction_in_nu", + "virt_packet_last_interaction_type", + "virt_packet_last_line_interaction_in_id", + "virt_packet_last_line_interaction_out_id", + "virt_packet_nus", + "virt_packet_energies", + ] + + hdf_name = "runner" + w_estimator_constant = ( + (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) + ).cgs.value + + def __init__( + self, + seed, + spectrum_frequency, + virtual_spectrum_spawn_range, + disable_electron_scattering, + enable_reflective_inner_boundary, + enable_full_relativity, + inner_boundary_albedo, + line_interaction_type, + integrator_settings, + v_packet_settings, + spectrum_method, + virtual_packet_logging, + packet_source=None, + debug_packets=False, + logger_buffer=1, + single_packet_seed=None, + ): self.seed = seed if packet_source is None: @@ -60,37 +103,54 @@ def __init__(self, seed, spectrum_frequency, virtual_spectrum_range, else: self.packet_source = packet_source # inject different packets + self.disable_electron_scattering = disable_electron_scattering self.spectrum_frequency = spectrum_frequency - self.virtual_spectrum_range = virtual_spectrum_range - self.sigma_thomson = sigma_thomson + self.virtual_spectrum_spawn_range = virtual_spectrum_spawn_range self.enable_reflective_inner_boundary = enable_reflective_inner_boundary self.inner_boundary_albedo = inner_boundary_albedo self.enable_full_relativity = enable_full_relativity + numba_config.ENABLE_FULL_RELATIVITY = enable_full_relativity self.line_interaction_type = line_interaction_type + self.single_packet_seed = single_packet_seed self.integrator_settings = integrator_settings self.v_packet_settings = v_packet_settings self.spectrum_method = spectrum_method + self.seed = seed self._integrator = None self._spectrum_integrated = None - if self.spectrum_method == 'integrated': - self.optional_hdf_properties.append('spectrum_integrated') + self.virt_logging = virtual_packet_logging + self.virt_packet_last_interaction_type = np.ones(2) * -1 + self.virt_packet_last_interaction_in_nu = np.ones(2) * -1.0 + self.virt_packet_last_line_interaction_in_id = np.ones(2) * -1 + self.virt_packet_last_line_interaction_out_id = np.ones(2) * -1 + self.virt_packet_nus = np.ones(2) * -1.0 + self.virt_packet_energies = np.ones(2) * -1.0 - def _initialize_estimator_arrays(self, no_of_shells, tau_sobolev_shape): + # set up logger based on config + mc_logger.DEBUG_MODE = debug_packets + mc_logger.BUFFER = logger_buffer + + if self.spectrum_method == "integrated": + self.optional_hdf_properties.append("spectrum_integrated") + + def _initialize_estimator_arrays(self, tau_sobolev_shape): """ Initialize the output arrays of the montecarlo simulation. Parameters ---------- - model: ~Radial1DModel + tau_sobolev_shape: tuple + tuple for the tau_sobolev_shape """ # Estimators - self.j_estimator = np.zeros(no_of_shells, dtype=np.float64) - self.nu_bar_estimator = np.zeros(no_of_shells, dtype=np.float64) + self.j_estimator = np.zeros(tau_sobolev_shape[1], dtype=np.float64) + self.nu_bar_estimator = np.zeros(tau_sobolev_shape[1], dtype=np.float64) self.j_blue_estimator = np.zeros(tau_sobolev_shape) self.Edotlu_estimator = np.zeros(tau_sobolev_shape) + # TODO: this is the wrong attribute naming style. def _initialize_geometry_arrays(self, model): """ @@ -101,15 +161,23 @@ def _initialize_geometry_arrays(self, model): model : model.Radial1DModel """ - self.r_inner_cgs = model.r_inner.to('cm').value - self.r_outer_cgs = model.r_outer.to('cm').value - self.v_inner_cgs = model.v_inner.to('cm/s').value - - def _initialize_packets(self, T, no_of_packets): + self.r_inner_cgs = model.r_inner.to("cm").value + self.r_outer_cgs = model.r_outer.to("cm").value + self.v_inner_cgs = model.v_inner.to("cm/s").value + + def _initialize_packets(self, T, no_of_packets, iteration): + # the iteration is added each time to preserve randomness + # across different simulations with the same temperature, + # for example. We seed the random module instead of the numpy module + # because we call random.sample, which references a different internal + # state than in the numpy.random module. + seed = self.seed + iteration + rng = np.random.default_rng(seed=seed) + seeds = rng.choice(MAX_SEED_VAL, no_of_packets, replace=True) nus, mus, energies = self.packet_source.create_packets( - T, - no_of_packets - ) + T, no_of_packets, rng + ) + mc_config_module.packet_seeds = seeds self.input_nu = nus self.input_mu = mus self.input_energy = energies @@ -118,71 +186,83 @@ def _initialize_packets(self, T, no_of_packets): self._output_energy = np.ones(no_of_packets, dtype=np.float64) * -99.0 self.last_line_interaction_in_id = -1 * np.ones( - no_of_packets, dtype=np.int64) + no_of_packets, dtype=np.int64 + ) self.last_line_interaction_out_id = -1 * np.ones( - no_of_packets, dtype=np.int64) + no_of_packets, dtype=np.int64 + ) self.last_line_interaction_shell_id = -1 * np.ones( - no_of_packets, dtype=np.int64) - self.last_interaction_type = -1 * np.ones( - no_of_packets, dtype=np.int64) + no_of_packets, dtype=np.int64 + ) + self.last_interaction_type = -1 * np.ones(no_of_packets, dtype=np.int64) self.last_interaction_in_nu = np.zeros(no_of_packets, dtype=np.float64) self._montecarlo_virtual_luminosity = u.Quantity( - np.zeros_like(self.spectrum_frequency.value), - 'erg / s' - ) + np.zeros_like(self.spectrum_frequency.value), "erg / s" + ) @property def spectrum(self): return TARDISSpectrum( - self.spectrum_frequency, - self.montecarlo_emitted_luminosity) + self.spectrum_frequency, self.montecarlo_emitted_luminosity + ) @property def spectrum_reabsorbed(self): return TARDISSpectrum( - self.spectrum_frequency, - self.montecarlo_reabsorbed_luminosity) + self.spectrum_frequency, self.montecarlo_reabsorbed_luminosity + ) @property def spectrum_virtual(self): if np.all(self.montecarlo_virtual_luminosity == 0): warnings.warn( - "MontecarloRunner.spectrum_virtual" - "is zero. Please run the montecarlo simulation with" - "no_of_virtual_packets > 0", UserWarning) + "MontecarloRunner.spectrum_virtual" + "is zero. Please run the montecarlo simulation with" + "no_of_virtual_packets > 0", + UserWarning, + ) return TARDISSpectrum( - self.spectrum_frequency, - self.montecarlo_virtual_luminosity) + self.spectrum_frequency, self.montecarlo_virtual_luminosity + ) @property def spectrum_integrated(self): if self._spectrum_integrated is None: self._spectrum_integrated = self.integrator.calculate_spectrum( - self.spectrum_frequency[:-1], **self.integrator_settings) + self.spectrum_frequency[:-1], **self.integrator_settings + ) return self._spectrum_integrated @property def integrator(self): if self._integrator is None: warnings.warn( - "MontecarloRunner.integrator: " - "The FormalIntegrator is not yet available." - "Please run the montecarlo simulation at least once.", - UserWarning) + "MontecarloRunner.integrator: " + "The FormalIntegrator is not yet available." + "Please run the montecarlo simulation at least once.", + UserWarning, + ) if self.enable_full_relativity: raise NotImplementedError( - "The FormalIntegrator is not yet implemented for the full " - "relativity mode. " - "Please run with config option enable_full_relativity: " - "False." + "The FormalIntegrator is not yet implemented for the full " + "relativity mode. " + "Please run with config option enable_full_relativity: " + "False." ) return self._integrator - def run(self, model, plasma, no_of_packets, - no_of_virtual_packets=0, nthreads=1, - last_run=False): + def run( + self, + model, + plasma, + no_of_packets, + no_of_virtual_packets=0, + nthreads=1, + last_run=False, + iteration=0, + ): """ Run the montecarlo calculation @@ -199,47 +279,41 @@ def run(self, model, plasma, no_of_packets, ------- None """ - self._integrator = FormalIntegrator( - model, - plasma, - self) + self._integrator = FormalIntegrator(model, plasma, self) self.time_of_simulation = self.calculate_time_of_simulation(model) self.volume = model.volume - self._initialize_estimator_arrays(self.volume.shape[0], - plasma.tau_sobolevs.shape) + + # Initializing estimator array + self._initialize_estimator_arrays(plasma.tau_sobolevs.shape) + self._initialize_geometry_arrays(model) - self._initialize_packets(model.t_inner.value, - no_of_packets) - - montecarlo.montecarlo_radial1d( - model, plasma, self, - virtual_packet_flag=no_of_virtual_packets, - nthreads=nthreads, - last_run=last_run) - # Workaround so that j_blue_estimator is in the right ordering - # They are written as an array of dimension (no_of_shells, no_of_lines) - # but python expects (no_of_lines, no_of_shells) - self.j_blue_estimator = np.ascontiguousarray( - self.j_blue_estimator.flatten().reshape( - self.j_blue_estimator.shape, order='F') - ) - self.Edotlu_estimator = np.ascontiguousarray( - self.Edotlu_estimator.flatten().reshape( - self.Edotlu_estimator.shape, order='F') - ) + self._initialize_packets(model.t_inner.value, no_of_packets, iteration) + + configuration_initialize(self, no_of_virtual_packets) + montecarlo_radial1d(model, plasma, self) + # montecarlo.montecarlo_radial1d( + # model, plasma, self, + # virtual_packet_flag=no_of_virtual_packets, + # nthreads=nthreads, + # last_run=last_run) def legacy_return(self): - return (self.output_nu, self.output_energy, - self.j_estimator, self.nu_bar_estimator, - self.last_line_interaction_in_id, - self.last_line_interaction_out_id, - self.last_interaction_type, - self.last_line_interaction_shell_id) + return ( + self.output_nu, + self.output_energy, + self.j_estimator, + self.nu_bar_estimator, + self.last_line_interaction_in_id, + self.last_line_interaction_out_id, + self.last_interaction_type, + self.last_line_interaction_shell_id, + ) def get_line_interaction_id(self, line_interaction_type): - return ['scatter', 'downbranch', 'macroatom'].index( - line_interaction_type) + return ["scatter", "downbranch", "macroatom"].index( + line_interaction_type + ) @property def output_nu(self): @@ -255,9 +329,11 @@ def virtual_packet_nu(self): return u.Quantity(self.virt_packet_nus, u.Hz) except AttributeError: warnings.warn( - "MontecarloRunner.virtual_packet_nu:" - "compile with --with-vpacket-logging" - "to access this property", UserWarning) + "MontecarloRunner.virtual_packet_nu:" + "compile with --with-vpacket-logging" + "to access this property", + UserWarning, + ) return None @property @@ -266,9 +342,11 @@ def virtual_packet_energy(self): return u.Quantity(self.virt_packet_energies, u.erg) except AttributeError: warnings.warn( - "MontecarloRunner.virtual_packet_energy:" - "compile with --with-vpacket-logging" - "to access this property", UserWarning) + "MontecarloRunner.virtual_packet_energy:" + "compile with --with-vpacket-logging" + "to access this property", + UserWarning, + ) return None @property @@ -277,9 +355,11 @@ def virtual_packet_luminosity(self): return self.virtual_packet_energy / self.time_of_simulation except TypeError: warnings.warn( - "MontecarloRunner.virtual_packet_luminosity:" - "compile with --with-vpacket-logging" - "to access this property", UserWarning) + "MontecarloRunner.virtual_packet_luminosity:" + "compile with --with-vpacket-logging" + "to access this property", + UserWarning, + ) return None @property @@ -309,51 +389,57 @@ def reabsorbed_packet_luminosity(self): @property def montecarlo_reabsorbed_luminosity(self): return u.Quantity( - np.histogram( - self.reabsorbed_packet_nu, - weights=self.reabsorbed_packet_luminosity, - bins=self.spectrum_frequency.value)[0], - 'erg / s' - ) + np.histogram( + self.reabsorbed_packet_nu, + weights=self.reabsorbed_packet_luminosity, + bins=self.spectrum_frequency.value, + )[0], + "erg / s", + ) @property def montecarlo_emitted_luminosity(self): return u.Quantity( - np.histogram( - self.emitted_packet_nu, - weights=self.emitted_packet_luminosity, - bins=self.spectrum_frequency.value)[0], - 'erg / s' - ) + np.histogram( + self.emitted_packet_nu, + weights=self.emitted_packet_luminosity, + bins=self.spectrum_frequency.value, + )[0], + "erg / s", + ) @property def montecarlo_virtual_luminosity(self): return ( - self._montecarlo_virtual_luminosity[:-1] / - self.time_of_simulation.value) + self._montecarlo_virtual_luminosity[:-1] + / self.time_of_simulation.value + ) - def calculate_emitted_luminosity(self, luminosity_nu_start, - luminosity_nu_end): + def calculate_emitted_luminosity( + self, luminosity_nu_start, luminosity_nu_end + ): luminosity_wavelength_filter = ( - (self.emitted_packet_nu > luminosity_nu_start) & - (self.emitted_packet_nu < luminosity_nu_end)) + self.emitted_packet_nu > luminosity_nu_start + ) & (self.emitted_packet_nu < luminosity_nu_end) emitted_luminosity = self.emitted_packet_luminosity[ - luminosity_wavelength_filter].sum() + luminosity_wavelength_filter + ].sum() return emitted_luminosity def calculate_reabsorbed_luminosity( - self, luminosity_nu_start, - luminosity_nu_end): + self, luminosity_nu_start, luminosity_nu_end + ): luminosity_wavelength_filter = ( - (self.reabsorbed_packet_nu > luminosity_nu_start) & - (self.reabsorbed_packet_nu < luminosity_nu_end)) + self.reabsorbed_packet_nu > luminosity_nu_start + ) & (self.reabsorbed_packet_nu < luminosity_nu_end) reabsorbed_luminosity = self.reabsorbed_packet_luminosity[ - luminosity_wavelength_filter].sum() + luminosity_wavelength_filter + ].sum() return reabsorbed_luminosity @@ -381,22 +467,31 @@ def calculate_radiationfield_properties(self): """ t_rad = ( - self.t_rad_estimator_constant * - self.nu_bar_estimator / - self.j_estimator) + self.t_rad_estimator_constant + * self.nu_bar_estimator + / self.j_estimator + ) w = self.j_estimator / ( - 4 * const.sigma_sb.cgs.value * t_rad ** 4 * - self.time_of_simulation.value * - self.volume.value) + 4 + * const.sigma_sb.cgs.value + * t_rad ** 4 + * self.time_of_simulation.value + * self.volume.value + ) return t_rad * u.K, w def calculate_luminosity_inner(self, model): - return (4 * np.pi * const.sigma_sb.cgs * - model.r_inner[0] ** 2 * model.t_inner ** 4).to('erg/s') + return ( + 4 + * np.pi + * const.sigma_sb.cgs + * model.r_inner[0] ** 2 + * model.t_inner ** 4 + ).to("erg/s") def calculate_time_of_simulation(self, model): - return (1.0 * u.erg / self.calculate_luminosity_inner(model)) + return 1.0 * u.erg / self.calculate_luminosity_inner(model) def calculate_f_nu(self, frequency): pass @@ -419,26 +514,42 @@ def from_config(cls, config, packet_source=None): """ if config.plasma.disable_electron_scattering: - logger.warn('Disabling electron scattering - this is not physical') - sigma_thomson = 1e-200 * (u.cm ** 2) + logger.warn( + "Disabling electron scattering - this is not physical." + "Likely bug in formal integral - " + "will not give same results." + ) + numba_config.SIGMA_THOMSON = 1e-200 + # mc_config_module.disable_electron_scattering = True else: logger.debug("Electron scattering switched on") - sigma_thomson = const.sigma_T.cgs + numba_config.SIGMA_THOMSON = const.sigma_T.to("cm^2").value + # mc_config_module.disable_electron_scattering = False spectrum_frequency = quantity_linspace( - config.spectrum.stop.to('Hz', u.spectral()), - config.spectrum.start.to('Hz', u.spectral()), - num=config.spectrum.num + 1) - - return cls(seed=config.montecarlo.seed, - spectrum_frequency=spectrum_frequency, - virtual_spectrum_range=config.montecarlo.virtual_spectrum_range, - sigma_thomson=sigma_thomson, - enable_reflective_inner_boundary=config.montecarlo.enable_reflective_inner_boundary, - inner_boundary_albedo=config.montecarlo.inner_boundary_albedo, - enable_full_relativity=config.montecarlo.enable_full_relativity, - line_interaction_type=config.plasma.line_interaction_type, - integrator_settings=config.spectrum.integrated, - v_packet_settings=config.spectrum.virtual, - spectrum_method=config.spectrum.method, - packet_source=packet_source) + config.spectrum.stop.to("Hz", u.spectral()), + config.spectrum.start.to("Hz", u.spectral()), + num=config.spectrum.num + 1, + ) + mc_config_module.disable_line_scattering = ( + config.plasma.disable_line_scattering + ) + + return cls( + seed=config.montecarlo.seed, + spectrum_frequency=spectrum_frequency, + virtual_spectrum_spawn_range=config.montecarlo.virtual_spectrum_spawn_range, + enable_reflective_inner_boundary=config.montecarlo.enable_reflective_inner_boundary, + inner_boundary_albedo=config.montecarlo.inner_boundary_albedo, + enable_full_relativity=config.montecarlo.enable_full_relativity, + line_interaction_type=config.plasma.line_interaction_type, + integrator_settings=config.spectrum.integrated, + v_packet_settings=config.spectrum.virtual, + spectrum_method=config.spectrum.method, + disable_electron_scattering=config.plasma.disable_electron_scattering, + packet_source=packet_source, + debug_packets=config.montecarlo.debug_packets, + logger_buffer=config.montecarlo.logger_buffer, + single_packet_seed=config.montecarlo.single_packet_seed, + virtual_packet_logging=config.spectrum.virtual.virtual_packet_logging, + ) diff --git a/tardis/montecarlo/formal_integral.py b/tardis/montecarlo/formal_integral.py index ad9463153e9..7b0fe3f499f 100644 --- a/tardis/montecarlo/formal_integral.py +++ b/tardis/montecarlo/formal_integral.py @@ -5,31 +5,59 @@ from scipy.interpolate import interp1d from astropy import units as u from tardis import constants as const +from numba import njit +import pdb + + +from tardis.montecarlo.montecarlo_numba import njit_dict +from tardis.montecarlo.montecarlo_numba.numba_interface import ( + numba_plasma_initialize, +) from tardis.montecarlo.montecarlo import formal_integral from tardis.montecarlo.spectrum import TARDISSpectrum +C_INV = 3.33564e-11 +M_PI = np.arccos(-1) +KB_CGS = 1.3806488e-16 +H_CGS = 6.62606957e-27 +SIGMA_THOMSON = 6.652486e-25 + class IntegrationError(Exception): pass -class FormalIntegrator(object): +# integrator_spec = [ +# ('model', float64), +# ('plasma', float64), +# ('runner', float64), +# ('points', int64) +# ] +# +# @jitclass(integrator_spec) +class FormalIntegrator(object): def __init__(self, model, plasma, runner, points=1000): self.model = model - self.plasma = plasma + if plasma: + self.plasma = numba_plasma_initialize( + plasma, runner.line_interaction_type + ) + self.atomic_data = plasma.atomic_data + self.original_plasma = plasma self.runner = runner self.points = points def check(self, raises=True): - ''' + """ A method that determines if the formal integral can be performed with the current configuration settings The function returns False if the configuration conflicts with the required settings. If raises evaluates to True, then a IntegrationError is raised instead - ''' + """ + def raise_or_return(message): if raises: raise IntegrationError(message) @@ -40,51 +68,47 @@ def raise_or_return(message): for obj in (self.model, self.plasma, self.runner): if obj is None: return raise_or_return( - 'The integrator is missing either model, plasma or ' - 'runner. Please make sure these are provided to the ' - 'FormalIntegrator.' - ) + "The integrator is missing either model, plasma or " + "runner. Please make sure these are provided to the " + "FormalIntegrator." + ) - if not self.runner.line_interaction_type in ['downbranch', 'macroatom']: + if not self.runner.line_interaction_type in ["downbranch", "macroatom"]: return raise_or_return( - 'The FormalIntegrator currently only works for ' - 'line_interaction_type == "downbranch"' - 'and line_interaction_type == "macroatom"' - ) + "The FormalIntegrator currently only works for " + 'line_interaction_type == "downbranch"' + 'and line_interaction_type == "macroatom"' + ) return True - def calculate_spectrum(self, frequency, points=None, - interpolate_shells=-1, raises=True): + def calculate_spectrum( + self, frequency, points=None, interpolate_shells=-1, raises=True + ): # Very crude implementation # The c extension needs bin centers (or something similar) # while TARDISSpectrum needs bin edges self.check(raises) N = points or self.points self.interpolate_shells = interpolate_shells - frequency = frequency.to('Hz', u.spectral()) + frequency = frequency.to("Hz", u.spectral()) - luminosity = u.Quantity( - formal_integral( - self, - frequency, - N), - 'erg' - ) * (frequency[1] - frequency[0]) + luminosity = u.Quantity(self.formal_integral(frequency, N), "erg") * ( + frequency[1] - frequency[0] + ) # Ugly hack to convert to 'bin edges' frequency = u.Quantity( - np.concatenate([ + np.concatenate( + [ frequency.value, - [ - frequency.value[-1] + np.diff(frequency.value)[-1] - ]]), - frequency.unit) - - return TARDISSpectrum( - frequency, - luminosity - ) + [frequency.value[-1] + np.diff(frequency.value)[-1]], + ] + ), + frequency.unit, + ) + + return TARDISSpectrum(frequency, luminosity) def make_source_function(self): """ @@ -101,120 +125,511 @@ def make_source_function(self): ------- Numpy array containing ( 1 - exp(-tau_ul) ) S_ul ordered by wavelength of the transition u -> l """ + model = self.model - plasma = self.plasma runner = self.runner - atomic_data = self.plasma.atomic_data - macro_ref = atomic_data.macro_atom_references - macro_data = atomic_data.macro_atom_data - no_lvls = len(atomic_data.levels) + macro_ref = self.atomic_data.macro_atom_references + macro_data = self.atomic_data.macro_atom_data + + no_lvls = len(self.atomic_data.levels) no_shells = len(model.w) - if runner.line_interaction_type == 'macroatom': + if runner.line_interaction_type == "macroatom": internal_jump_mask = (macro_data.transition_type >= 0).values ma_int_data = macro_data[internal_jump_mask] - internal = plasma.transition_probabilities[internal_jump_mask] + internal = self.original_plasma.transition_probabilities[ + internal_jump_mask + ] source_level_idx = ma_int_data.source_level_idx.values - destination_level_idx = ma_int_data.destination_level_idx.values + destination_level_idx = ma_int_data.destination_level_idx.values - Edotlu_norm_factor = (1 / (runner.time_of_simulation * model.volume)) - exptau = 1 - np.exp(- plasma.tau_sobolevs) + Edotlu_norm_factor = 1 / (runner.time_of_simulation * model.volume) + exptau = 1 - np.exp(-self.plasma.tau_sobolev) Edotlu = Edotlu_norm_factor * exptau * runner.Edotlu_estimator # The following may be achieved by calling the appropriate plasma # functions - Jbluelu_norm_factor = (const.c.cgs * model.time_explosion / - (4 * np.pi * runner.time_of_simulation * - model.volume)).to("1/(cm^2 s)").value + Jbluelu_norm_factor = ( + ( + const.c.cgs + * model.time_explosion + / (4 * np.pi * runner.time_of_simulation * model.volume) + ) + .to("1/(cm^2 s)") + .value + ) # Jbluelu should already by in the correct order, i.e. by wavelength of # the transition l->u Jbluelu = runner.j_blue_estimator * Jbluelu_norm_factor - upper_level_index = atomic_data.lines.index.droplevel('level_number_lower') - e_dot_lu = pd.DataFrame(Edotlu, index=upper_level_index) - e_dot_u = e_dot_lu.groupby(level=[0, 1, 2]).sum() + upper_level_index = self.atomic_data.lines.index.droplevel( + "level_number_lower" + ) + e_dot_lu = pd.DataFrame(Edotlu, index=upper_level_index) + e_dot_u = e_dot_lu.groupby(level=[0, 1, 2]).sum() e_dot_u_src_idx = macro_ref.loc[e_dot_u.index].references_idx.values - if runner.line_interaction_type == 'macroatom': + if runner.line_interaction_type == "macroatom": C_frame = pd.DataFrame( columns=np.arange(no_shells), index=macro_ref.index ) q_indices = (source_level_idx, destination_level_idx) for shell in range(no_shells): Q = sp.coo_matrix( - (internal[shell], q_indices), shape=(no_lvls, no_lvls) + (internal[shell], q_indices), shape=(no_lvls, no_lvls) ) inv_N = sp.identity(no_lvls) - Q e_dot_u_vec = np.zeros(no_lvls) e_dot_u_vec[e_dot_u_src_idx] = e_dot_u[shell].values C_frame[shell] = sp.linalg.spsolve(inv_N.T, e_dot_u_vec) - e_dot_u.index.names = ['atomic_number', 'ion_number', 'source_level_number'] # To make the q_ul e_dot_u product work, could be cleaner - transitions = atomic_data.macro_atom_data[atomic_data.macro_atom_data.transition_type == -1].copy() - transitions_index = transitions.set_index(['atomic_number', 'ion_number', 'source_level_number']).index.copy() - tmp = plasma.transition_probabilities[(atomic_data.macro_atom_data.transition_type == -1).values] + e_dot_u.index.names = [ + "atomic_number", + "ion_number", + "source_level_number", + ] # To make the q_ul e_dot_u product work, could be cleaner + transitions = self.original_plasma.atomic_data.macro_atom_data[ + self.original_plasma.atomic_data.macro_atom_data.transition_type + == -1 + ].copy() + transitions_index = transitions.set_index( + ["atomic_number", "ion_number", "source_level_number"] + ).index.copy() + tmp = self.original_plasma.transition_probabilities[ + (self.atomic_data.macro_atom_data.transition_type == -1).values + ] q_ul = tmp.set_index(transitions_index) - t = model.time_explosion.value - lines = atomic_data.lines.set_index('line_id') - wave = lines.wavelength_cm.loc[transitions.transition_line_id].values.reshape(-1,1) - if runner.line_interaction_type == 'macroatom': + t = model.time_explosion.value + lines = self.atomic_data.lines.set_index("line_id") + wave = lines.wavelength_cm.loc[ + transitions.transition_line_id + ].values.reshape(-1, 1) + if runner.line_interaction_type == "macroatom": e_dot_u = C_frame.loc[e_dot_u.index] - att_S_ul = (wave * (q_ul * e_dot_u) * t / (4 * np.pi)) + att_S_ul = wave * (q_ul * e_dot_u) * t / (4 * np.pi) - result = pd.DataFrame(att_S_ul.values, index=transitions.transition_line_id.values) + result = pd.DataFrame( + att_S_ul.values, index=transitions.transition_line_id.values + ) att_S_ul = result.loc[lines.index.values].values # Jredlu should already by in the correct order, i.e. by wavelength of # the transition l->u (similar to Jbluelu) - Jredlu = Jbluelu * np.exp(-plasma.tau_sobolevs.values) + att_S_ul + Jredlu = Jbluelu * np.exp(-self.plasma.tau_sobolev) + att_S_ul if self.interpolate_shells > 0: - att_S_ul, Jredlu, Jbluelu, e_dot_u = self.interpolate_integrator_quantities( - att_S_ul, Jredlu, Jbluelu, e_dot_u) + ( + att_S_ul, + Jredlu, + Jbluelu, + e_dot_u, + ) = self.interpolate_integrator_quantities( + att_S_ul, Jredlu, Jbluelu, e_dot_u + ) else: runner.r_inner_i = runner.r_inner_cgs runner.r_outer_i = runner.r_outer_cgs - runner.tau_sobolevs_integ = plasma.tau_sobolevs.values - runner.electron_densities_integ = plasma.electron_densities.values + runner.tau_sobolevs_integ = self.plasma.tau_sobolev + runner.electron_densities_integ = self.plasma.electron_density return att_S_ul, Jredlu, Jbluelu, e_dot_u - def interpolate_integrator_quantities(self, att_S_ul, Jredlu, - Jbluelu, e_dot_u): + def interpolate_integrator_quantities( + self, att_S_ul, Jredlu, Jbluelu, e_dot_u + ): runner = self.runner plasma = self.plasma nshells = self.interpolate_shells - r_middle = (runner.r_inner_cgs + runner.r_outer_cgs) / 2. + r_middle = (runner.r_inner_cgs + runner.r_outer_cgs) / 2.0 r_integ = np.linspace( - runner.r_inner_cgs[0], runner.r_outer_cgs[-1], nshells + runner.r_inner_cgs[0], runner.r_outer_cgs[-1], nshells ) runner.r_inner_i = r_integ[:-1] runner.r_outer_i = r_integ[1:] - r_middle_integ = (r_integ[:-1] + r_integ[1:]) / 2. + r_middle_integ = (r_integ[:-1] + r_integ[1:]) / 2.0 runner.electron_densities_integ = interp1d( - r_middle, plasma.electron_densities, - fill_value='extrapolate', kind='nearest')(r_middle_integ) + r_middle, + plasma.electron_density, + fill_value="extrapolate", + kind="nearest", + )(r_middle_integ) # Assume tau_sobolevs to be constant within a shell # (as in the MC simulation) runner.tau_sobolevs_integ = interp1d( - r_middle, plasma.tau_sobolevs, - fill_value='extrapolate', kind='nearest')(r_middle_integ) - att_S_ul = interp1d( - r_middle, att_S_ul, fill_value='extrapolate')(r_middle_integ) - Jredlu = interp1d( - r_middle, Jredlu, fill_value='extrapolate')(r_middle_integ) - Jbluelu = interp1d( - r_middle, Jbluelu, fill_value='extrapolate')(r_middle_integ) - e_dot_u = interp1d( - r_middle, e_dot_u, fill_value='extrapolate')(r_middle_integ) + r_middle, + plasma.tau_sobolev, + fill_value="extrapolate", + kind="nearest", + )(r_middle_integ) + att_S_ul = interp1d(r_middle, att_S_ul, fill_value="extrapolate")( + r_middle_integ + ) + Jredlu = interp1d(r_middle, Jredlu, fill_value="extrapolate")( + r_middle_integ + ) + Jbluelu = interp1d(r_middle, Jbluelu, fill_value="extrapolate")( + r_middle_integ + ) + e_dot_u = interp1d(r_middle, e_dot_u, fill_value="extrapolate")( + r_middle_integ + ) # Set negative values from the extrapolation to zero - att_S_ul = att_S_ul.clip(0.) - Jbluelu = Jbluelu.clip(0.) - Jredlu = Jredlu.clip(0.) - e_dot_u = e_dot_u.clip(0.) + att_S_ul = att_S_ul.clip(0.0) + Jbluelu = Jbluelu.clip(0.0) + Jredlu = Jredlu.clip(0.0) + e_dot_u = e_dot_u.clip(0.0) return att_S_ul, Jredlu, Jbluelu, e_dot_u + + def formal_integral(self, nu, N): + # TODO: get rid of storage later on + + res = self.make_source_function() + + att_S_ul = res[0].flatten(order="F") + Jred_lu = res[1].flatten(order="F") + Jblue_lu = res[2].flatten(order="F") + L = self._formal_integral( + self.model.t_inner.value, + nu, + nu.shape[0], + att_S_ul, + Jred_lu, + Jblue_lu, + N, + ) + return np.array(L, np.NPY_DOUBLE, nu.shape[0]) + + def _formal_integral( + self, iT, inu, inu_size, att_S_ul, Jred_lu, Jblue_lu, N + ): + # todo: add all the original todos + # Initialize the output which is shared among threads + L = np.zeros(inu_size) + # global read-only values + size_line = len(self.plasma.line_list_nu) + size_shell = self.model.no_of_shells # check + size_tau = size_line * size_shell + finished_nus = 0 + + R_ph = self.runner.r_inner_i[0] + R_max = self.runner.r_outer_i[size_shell - 1] + pp = np.zeros(N) # check + exp_tau = np.zeros(size_tau) + # TODO: multiprocessing + offset = 0 + size_z = 0 + z = np.zeros(2 * self.model.no_of_shells) + idx_nu_start = 0 + direction = 0 + I_nu = np.zeros(N) + shell_id = np.zeros(2 * self.model.no_of_shells) + # instantiate more variables here, maybe? + + # prepare exp_tau + exp_tau = np.exp(-self.plasma.tau_sobolev) # check + pp = calculate_p_values(R_max, N, pp) + + # done with instantiation + # now loop over wavelength in spectrum + for nu_idx in range(inu_size): + nu = inu[nu_idx] + # now loop over discrete values along line + for p_idx in range(1, N): + escat_contrib = 0 + p = pp[p_idx] + + # initialize z intersections for p values + size_z = self.populate_z(p, z, shell_id) # check returns + + # initialize I_nu + if p <= R_ph: + I_nu[p_idx] = intensity_black_body(nu * z[0], iT) + else: + I_nu[p_idx] = 0 + + # find first contributing lines + nu_start = nu * z[0] + nu_end = nu * z[1] + idx_nu_start = line_search( + self.plasma.line_list_nu, nu_start, size_line, idx_nu_start + ) + offset = shell_id[0] * size_line + + # start tracking accumulated e-scattering optical depth + zstart = self.model.time_explosion / C_INV * (1.0 - z[0]) + + # Initialize "pointers" + pline = self.plasma.line_list_nu + idx_nu_start + pexp_tau = exp_tau + offset + idx_nu_start + patt_S_ul = att_S_ul + offset + idx_nu_start + pJred_lu = Jred_lu + offset + idx_nu_start + pJblue_lu = Jblue_lu + offset + idx_nu_start + + # flag for first contribution to integration on current p-ray + first = 1 + + # loop over all interactions + for i in range(size_z - 1): + escat_op = ( + self.plasma.electron_density[int(shell_id[i])] + * SIGMA_THOMSON + ) + nu_end = nu * z[i + 1] + while np.all( + pline < self.plasma.line_list_nu + size_line + ): # check all condition + # increment all pointers simulatenously + pline += 1 + pexp_tau += 1 + patt_S_ul += 1 + pJblue_lu += 1 + + if pline[0] < nu_end.value: + break + + # calculate e-scattering optical depth to next resonance point + zend = ( + self.model.time_explosion + / C_INV + * (1.0 - pline[0] / nu.value) + ) # check + + if first == 1: + # first contribution to integration + # NOTE: this treatment of I_nu_b (given + # by boundary conditions) is not in Lucy 1999; + # should be re-examined carefully + escat_contrib += ( + (zend - zstart) + * escat_op + * (pJblue_lu[0] - I_nu[p_idx]) + ) + first = 0 + else: + # Account for e-scattering, c.f. Eqs 27, 28 in Lucy 1999 + Jkkp = 0.5 * (pJred_lu[0] + pJblue_lu[0]) + escat_contrib += ( + (zend - zstart) + * escat_op + * (Jkkp - I_nu[p_idx]) + ) + # this introduces the necessary ffset of one element between + # pJblue_lu and pJred_lu + pJred_lu += 1 + # pdb.set_trace() + I_nu[p_idx] = I_nu[p_idx] + escat_contrib.value + # // Lucy 1999, Eq 26 + I_nu[p_idx] = ( + I_nu[p_idx] * (pexp_tau[0][0]) + patt_S_ul[0] + ) # check about taking about asterisks beforehand elsewhere + + # // reset e-scattering opacity + escat_contrib = 0 + zstart = zend + # calculate e-scattering optical depth to grid cell boundary + + Jkkp = 0.5 * (pJred_lu[0] + pJblue_lu[0]) + zend = ( + self.model.time_explosion / C_INV * (1.0 - nu_end / nu) + ) # check + escat_contrib += ( + (zend - zstart) * escat_op * (Jkkp - I_nu[p_idx]) + ) + zstart = zend + + if i < size_z - 1: + # advance pointers + direction = shell_id[i + 1] - shell_id[i] + pexp_tau += direction * size_line + patt_S_ul += direction * size_line + pJred_lu += direction * size_line + pJblue_lu += direction * size_line + I_nu[p_idx] *= p + L[nu_idx] = ( + 8 * M_PI * M_PI * trapezoid_integration(I_nu, R_max / N, N) + ) + # something pragma op atomic + return L + + # @njit(**njit_dict) + def populate_z(self, p, oz, oshell_id): + """Calculate p line intersections + + This function calculates the intersection points of the p-line with + each shell + + Inputs: + :p: (double) distance of the integration line to the center + :oz: (array of doubles) will be set with z values. the array is truncated + by the value `1`. + :oshell_id: (int64) will be set with the corresponding shell_ids + """ + # abbreviations + r = self.runner.r_outer_i + N = self.model.no_of_shells # check + print(N) + inv_t = 1 / self.model.time_explosion + z = 0 + offset = N + + if p <= self.runner.r_inner_i[0]: + # intersect the photosphere + for i in range(N): + oz[i] = 1 - self.calculate_z(r[i], p, inv_t) + oshell_id[i] = i + return N + else: + # no intersection with photosphere + # that means we intersect each shell twice + for i in range(N): + z = self.calculate_z(r[i], p, inv_t) + if z == 0: + continue + if offset == N: + offset = i + # calculate the index in the resulting array + i_low = N - i - 1 # the far intersection with the shell + i_up = ( + N + i - 2 * offset + ) # the nearer intersection with the shell + + # setting the arrays; check return them? + oz[i_low] = 1 + z + oshell_id[i_low] = i + oz[i_up] = 1 - z + oshell_id[i_up] = i + return 2 * (N - offset) + + # @njit(**njit_dict) + def calculate_z(self, r, p, inv_t): + """Calculate distance to p line + + Calculate half of the length of the p-line inside a shell + of radius r in terms of unit length (c * t_exp). + If shell and p-line do not intersect, return 0. + + Inputs: + :r: (double) radius of the shell + :p: (double) distance of the p-line to the center of the supernova + :inv_t: (double) inverse time_explosio is needed to norm to unit-length + """ + if r > p: + return np.sqrt(r * r - p * p) * C_INV * inv_t.value + else: + return 0 + + +class BoundsError(ValueError): + pass + + +@njit(**njit_dict) +def line_search(nu, nu_insert, number_of_lines, result): + """ + Insert a value in to an array of line frequencies + + Inputs: + :nu: (array) line frequencies + :nu_insert: (int) value of nu key + :number_of_lines: (int) number of lines in the line list + + Outputs: + index of the next line ot the red. + If the key value is redder + than the reddest line returns number_of_lines. + """ + # TODO: fix the TARDIS_ERROR_OK + # tardis_error_t ret_val = TARDIS_ERROR_OK # check + imin = 0 + imax = number_of_lines - 1 + if nu_insert > nu[imin]: + result = imin + elif nu_insert < nu[imax]: + result = imax + 1 + else: + result = reverse_binary_search(nu, nu_insert, imin, imax, result) + result = result + 1 + return result + + +@njit(**njit_dict) +def reverse_binary_search(x, x_insert, imin, imax, result): + """Look for a place to insert a value in an inversely sorted float array. + + Inputs: + :x: (array) an inversely (largest to lowest) sorted float array + :x_insert: (value) a value to insert + :imin: (int) lower bound + :imax: (int) upper bound + + Outputs: + index of the next boundary to the left + """ + # ret_val = TARDIS_ERROR_OK # check + if x_insert > x[imin] or x_insert < x[imax]: + raise BoundsError # check + else: + imid = (imin + imax) >> 1 + while imax - imin > 2: + if x[imid] < x_insert: + imax = imid + 1 + else: + imin = imid + imid = (imin + imax) >> 1 + if imax - imin == 2 and x_insert < x[imin + 1]: + result = imin + 1 + else: + result = imin + return result + + +@njit(**njit_dict) +def binary_search(x, x_insert, imin, imax, result): + # TODO: actually return result + if x_insert < x[imin] or x_insert > x[imax]: + raise BoundsError + else: + while imax >= imin: + imid = (imin + imax) / 2 + if x[imid] == x_insert: + result = imid + break + elif x[imid] < x_insert: + imin = imid + 1 + else: + imax = imid - 1 + if imax - imid == 2 and x_insert < x[imin + 1]: + result = imin + else: + result = imin # check + return result + + +@njit(**njit_dict) +def trapezoid_integration(array, h, N): + # TODO: replace with np.trapz? + result = (array[0] + array[N - 1]) / 2 + for idx in range(1, N - 1): + result += array[idx] + return result * h + + +@njit +def intensity_black_body(nu, T): + if nu == 0: + return np.nan # to avoid ZeroDivisionError + beta_rad = 1 / (KB_CGS * T) + coefficient = 2 * H_CGS * C_INV * C_INV + return coefficient * nu * nu * nu / (np.exp(H_CGS * nu * beta_rad) - 1) + + +@njit(**njit_dict) +def calculate_p_values(R_max, N, opp): + for i in range(N): + opp[i] = R_max / (N - 1) * (i) + return opp diff --git a/tardis/montecarlo/montecarlo.pyx b/tardis/montecarlo/montecarlo.pyx index e3ac6e42f58..4b398c60425 100644 --- a/tardis/montecarlo/montecarlo.pyx +++ b/tardis/montecarlo/montecarlo.pyx @@ -123,7 +123,6 @@ cdef extern from "src/cmontecarlo.h": double *tau_bias int enable_biasing - void montecarlo_main_loop(storage_model_t * storage, int_type_t virtual_packet_flag, int nthreads, unsigned long seed) cdef extern from "src/integrator.h": double *_formal_integral( @@ -246,14 +245,14 @@ cdef initialize_storage_model(model, plasma, runner, storage_model_t *storage): storage.spectrum_start_nu = runner.spectrum_frequency.to('Hz').value.min() storage.spectrum_end_nu = runner.spectrum_frequency.to('Hz').value.max() # TODO: Linspace handling for virtual_spectrum_range - storage.spectrum_virt_start_nu = runner.virtual_spectrum_range.stop.to('Hz', units.spectral()).value - storage.spectrum_virt_end_nu = runner.virtual_spectrum_range.start.to('Hz', units.spectral()).value + storage.spectrum_virt_start_nu = runner.virtual_spectrum_spawn_range.end.to('Hz', units.spectral()).value + storage.spectrum_virt_end_nu = runner.virtual_spectrum_spawn_range.start.to('Hz', units.spectral()).value storage.spectrum_delta_nu = runner.spectrum_frequency.to('Hz').value[1] - runner.spectrum_frequency.to('Hz').value[0] storage.spectrum_virt_nu = PyArray_DATA( runner._montecarlo_virtual_luminosity.value) - storage.sigma_thomson = runner.sigma_thomson.cgs.value + storage.sigma_thomson = runner.sigma_thomson storage.inverse_sigma_thomson = 1.0 / storage.sigma_thomson storage.reflective_inner_boundary = runner.enable_reflective_inner_boundary storage.inner_boundary_albedo = runner.inner_boundary_albedo @@ -270,7 +269,7 @@ cdef initialize_storage_model(model, plasma, runner, storage_model_t *storage): runner.tau_bias[:-1] = ( ((runner.r_outer_cgs - runner.r_inner_cgs) * plasma.electron_densities.values * - runner.sigma_thomson.cgs.value)[::-1].cumsum()[::-1] + runner.sigma_thomson)[::-1].cumsum()[::-1] ) storage.tau_bias = PyArray_DATA(runner.tau_bias) @@ -278,59 +277,6 @@ cdef initialize_storage_model(model, plasma, runner, storage_model_t *storage): cdef np.ndarray[double, ndim=1] t_electrons = plasma.t_electrons storage.t_electrons = t_electrons.data -def montecarlo_radial1d(model, plasma, runner, int_type_t virtual_packet_flag=0, - int nthreads=4,last_run=False): - """ - Parameters - ---------- - model : `tardis.model_radial_oned.ModelRadial1D` - complete model - param photon_packets : PacketSource object - photon packets - - Returns - ------- - output_nus : `numpy.ndarray` - output_energies : `numpy.ndarray` - - TODO - np.ndarray[double, ndim=1] line_list_nu, - np.ndarray[double, ndim=2] tau_lines, - np.ndarray[double, ndim=1] ne, - double packet_energy, - np.ndarray[double, ndim=2] p_transition, - np.ndarray[int_type_t, ndim=1] type_transition, - np.ndarray[int_type_t, ndim=1] target_level_id, - np.ndarray[int_type_t, ndim=1] target_line_id, - np.ndarray[int_type_t, ndim=1] unroll_reference, - np.ndarray[int_type_t, ndim=1] line2level, - int_type_t log_packets, - int_type_t do_scatter - """ - - cdef storage_model_t storage - - initialize_storage_model(model, plasma, runner, &storage) - - montecarlo_main_loop(&storage, virtual_packet_flag, nthreads, runner.seed) - runner.virt_logging = LOG_VPACKETS - if LOG_VPACKETS != 0: - runner.virt_packet_nus = c_array_to_numpy(storage.virt_packet_nus, np.NPY_DOUBLE, storage.virt_packet_count) - runner.virt_packet_energies = c_array_to_numpy(storage.virt_packet_energies, np.NPY_DOUBLE, storage.virt_packet_count) - runner.virt_packet_last_interaction_in_nu = c_array_to_numpy(storage.virt_packet_last_interaction_in_nu, np.NPY_DOUBLE, storage.virt_packet_count) - runner.virt_packet_last_interaction_type = c_array_to_numpy(storage.virt_packet_last_interaction_type, np.NPY_INT64, storage.virt_packet_count) - runner.virt_packet_last_line_interaction_in_id = c_array_to_numpy(storage.virt_packet_last_line_interaction_in_id, np.NPY_INT64, - storage.virt_packet_count) - runner.virt_packet_last_line_interaction_out_id = c_array_to_numpy(storage.virt_packet_last_line_interaction_out_id, np.NPY_INT64, - storage.virt_packet_count) - else: - runner.virt_packet_nus = np.zeros(0) - runner.virt_packet_energies = np.zeros(0) - runner.virt_packet_last_interaction_in_nu = np.zeros(0) - runner.virt_packet_last_interaction_type = np.zeros(0) - runner.virt_packet_last_line_interaction_in_id = np.zeros(0) - runner.virt_packet_last_line_interaction_out_id = np.zeros(0) - # This will be a method of the Simulation object def formal_integral(self, nu, N): diff --git a/tardis/montecarlo/montecarlo_configuration.py b/tardis/montecarlo/montecarlo_configuration.py new file mode 100644 index 00000000000..438d9b15f19 --- /dev/null +++ b/tardis/montecarlo/montecarlo_configuration.py @@ -0,0 +1,15 @@ +from tardis import constants as const + +full_relativity = True +single_packet_seed = -1 +temporary_v_packet_bins = 0 +number_of_vpackets = 0 +montecarlo_seed = 0 +line_interaction_type = None +packet_seeds = [] +disable_electron_scattering = False +disable_line_scattering = False +survival_probability = 0.0 +tau_russian = 10.0 +LEGACY_MODE_ENABLED = False +VPACKET_LOGGING = False diff --git a/tardis/montecarlo/montecarlo_numba/__init__.py b/tardis/montecarlo/montecarlo_numba/__init__.py new file mode 100644 index 00000000000..53a7f130cd3 --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/__init__.py @@ -0,0 +1,8 @@ +from llvmlite import binding + +binding.set_option("tmp", "-non-global-value-max-name-size=2048") +njit_dict = {"fastmath": True, "error_model": "numpy"} + +from tardis.montecarlo.montecarlo_numba.r_packet import RPacket +from tardis.montecarlo.montecarlo_numba.base import montecarlo_radial1d +from tardis.montecarlo.montecarlo_numba.numba_interface import PacketCollection diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py new file mode 100644 index 00000000000..763399ee913 --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/base.py @@ -0,0 +1,226 @@ +from numba import prange, njit, jit +import logging +import numpy as np + +from tardis.montecarlo.montecarlo_numba.r_packet import ( + RPacket, + PacketStatus, + MonteCarloException, +) +from tardis.montecarlo.montecarlo_numba.numba_interface import ( + PacketCollection, + VPacketCollection, + NumbaModel, + numba_plasma_initialize, + Estimators, + configuration_initialize, +) + +from tardis.montecarlo import ( + montecarlo_configuration as montecarlo_configuration, +) + +from tardis.montecarlo.montecarlo_numba.single_packet_loop import ( + single_packet_loop, +) +from tardis.montecarlo.montecarlo_numba import njit_dict + + +def montecarlo_radial1d(model, plasma, runner): + packet_collection = PacketCollection( + runner.input_nu, + runner.input_mu, + runner.input_energy, + runner._output_nu, + runner._output_energy, + ) + + numba_model = NumbaModel( + runner.r_inner_cgs, + runner.r_outer_cgs, + model.time_explosion.to("s").value, + ) + numba_plasma = numba_plasma_initialize(plasma, runner.line_interaction_type) + estimators = Estimators( + runner.j_estimator, + runner.nu_bar_estimator, + runner.j_blue_estimator, + runner.Edotlu_estimator, + ) + packet_seeds = montecarlo_configuration.packet_seeds + + number_of_vpackets = montecarlo_configuration.number_of_vpackets + + ( + v_packets_energy_hist, + last_interaction_type, + virt_packet_nus, + virt_packet_energies, + virt_packet_last_interaction_in_nu, + virt_packet_last_interaction_type, + virt_packet_last_line_interaction_in_id, + virt_packet_last_line_interaction_out_id, + ) = montecarlo_main_loop( + packet_collection, + numba_model, + numba_plasma, + estimators, + runner.spectrum_frequency.value, + number_of_vpackets, + packet_seeds, + ) + + runner._montecarlo_virtual_luminosity.value[:] = v_packets_energy_hist + runner.last_interaction_type = last_interaction_type + + if montecarlo_configuration.VPACKET_LOGGING and number_of_vpackets > 0: + runner.virt_packet_nus = np.concatenate( + np.array(virt_packet_nus) + ).ravel() + runner.virt_packet_energies = np.concatenate( + np.array(virt_packet_energies) + ).ravel() + runner.virt_packet_last_interaction_in_nu = np.array( + virt_packet_last_interaction_in_nu + ) + runner.virt_packet_last_interaction_type = np.array( + virt_packet_last_interaction_type + ) + runner.virt_packet_last_line_interaction_in_id = np.array( + virt_packet_last_line_interaction_in_id + ) + runner.virt_packet_last_line_interaction_out_id = np.array( + virt_packet_last_line_interaction_out_id + ) + + +@njit(**njit_dict, nogil=True) +def montecarlo_main_loop( + packet_collection, + numba_model, + numba_plasma, + estimators, + spectrum_frequency, + number_of_vpackets, + packet_seeds, +): + """ + This is the main loop of the MonteCarlo routine that generates packets + and sends them through the ejecta. + + Parameters + ---------- + packet_collection: PacketCollection + numba_model: NumbaModel + estimators: NumbaEstimators + spectrum_frequency: astropy.units.Quantity + frequency bins + number_of_vpackets: int + VPackets released per interaction + packet_seeds: numpy.array + """ + output_nus = np.empty_like(packet_collection.packets_output_nu) + last_interaction_types = ( + np.ones_like(packet_collection.packets_output_nu) * -1 + ) + output_energies = np.empty_like(packet_collection.packets_output_nu) + + v_packets_energy_hist = np.zeros_like(spectrum_frequency) + delta_nu = spectrum_frequency[1] - spectrum_frequency[0] + + virt_packet_nus = [] + virt_packet_energies = [] + virt_packet_last_interaction_in_nu = [] + virt_packet_last_interaction_type = [] + virt_packet_last_line_interaction_in_id = [] + virt_packet_last_line_interaction_out_id = [] + + print("Running post-merge numba montecarlo (with C close lines)!") + for i in prange(len(output_nus)): + if montecarlo_configuration.single_packet_seed != -1: + seed = packet_seeds[montecarlo_configuration.single_packet_seed] + np.random.seed(seed) + else: + seed = packet_seeds[i] + np.random.seed(seed) + r_packet = RPacket( + numba_model.r_inner[0], + packet_collection.packets_input_mu[i], + packet_collection.packets_input_nu[i], + packet_collection.packets_input_energy[i], + seed, + i, + 0, + ) + vpacket_collection = VPacketCollection( + r_packet.index, + spectrum_frequency, + montecarlo_configuration.v_packet_spawn_start_frequency, + montecarlo_configuration.v_packet_spawn_end_frequency, + number_of_vpackets, + montecarlo_configuration.temporary_v_packet_bins, + ) + loop = single_packet_loop( + r_packet, numba_model, numba_plasma, estimators, vpacket_collection + ) + # if loop and 'stop' in loop: + # raise MonteCarloException + + output_nus[i] = r_packet.nu + + if r_packet.status == PacketStatus.REABSORBED: + output_energies[i] = -r_packet.energy + last_interaction_types[i] = 0 + elif r_packet.status == PacketStatus.EMITTED: + output_energies[i] = r_packet.energy + last_interaction_types[i] = 1 + + vpackets_nu = vpacket_collection.nus[: vpacket_collection.idx] + vpackets_energy = vpacket_collection.energies[: vpacket_collection.idx] + + v_packets_idx = np.floor( + (vpackets_nu - spectrum_frequency[0]) / delta_nu + ).astype(np.int64) + # if we're only in a single-packet mode + # if montecarlo_configuration.single_packet_seed == -1: + # break + for j, idx in enumerate(v_packets_idx): + if (vpackets_nu[j] < spectrum_frequency[0]) or ( + vpackets_nu[j] > spectrum_frequency[-1] + ): + continue + v_packets_energy_hist[idx] += vpackets_energy[j] + + if montecarlo_configuration.VPACKET_LOGGING: + virt_packet_nus.append( + vpacket_collection.nus[: vpacket_collection.idx] + ) + virt_packet_energies.append( + vpacket_collection.energies[: vpacket_collection.idx] + ) + virt_packet_last_interaction_in_nu.append( + vpacket_collection.last_interaction_in_nu + ) + virt_packet_last_interaction_type.append( + vpacket_collection.last_interaction_type + ) + virt_packet_last_line_interaction_in_id.append( + vpacket_collection.last_interaction_in_id + ) + virt_packet_last_line_interaction_out_id.append( + vpacket_collection.last_interaction_out_id + ) + + packet_collection.packets_output_energy[:] = output_energies[:] + packet_collection.packets_output_nu[:] = output_nus[:] + + return ( + v_packets_energy_hist, + last_interaction_types, + virt_packet_nus, + virt_packet_energies, + virt_packet_last_interaction_in_nu, + virt_packet_last_interaction_type, + virt_packet_last_line_interaction_in_id, + virt_packet_last_line_interaction_out_id, + ) diff --git a/tardis/montecarlo/montecarlo_numba/interaction.py b/tardis/montecarlo/montecarlo_numba/interaction.py new file mode 100644 index 00000000000..2f3e72717d3 --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/interaction.py @@ -0,0 +1,120 @@ +from numba import njit +from tardis.montecarlo.montecarlo_numba import njit_dict +from tardis.montecarlo.montecarlo_numba.numba_interface import ( + LineInteractionType, +) + +from tardis.montecarlo import ( + montecarlo_configuration as montecarlo_configuration, +) +from tardis.montecarlo.montecarlo_numba.r_packet import ( + get_doppler_factor, + get_inverse_doppler_factor, + get_random_mu, + angle_aberration_CMF_to_LF, + test_for_close_line, +) +from tardis.montecarlo.montecarlo_numba.macro_atom import macro_atom + + +@njit(**njit_dict) +def thomson_scatter(r_packet, time_explosion): + """ + Thomson scattering — no longer line scattering + 2) get the doppler factor at that position with the old angle + 3) convert the current energy and nu into the comoving + frame with the old mu + 4) Scatter and draw new mu - update mu + 5) Transform the comoving energy and nu back using the new mu + + Parameters + ---------- + r_packet : RPacket + time_explosion: float + time since explosion in seconds + + """ + old_doppler_factor = get_doppler_factor( + r_packet.r, r_packet.mu, time_explosion + ) + comov_nu = r_packet.nu * old_doppler_factor + comov_energy = r_packet.energy * old_doppler_factor + r_packet.mu = get_random_mu() + inverse_new_doppler_factor = get_inverse_doppler_factor( + r_packet.r, r_packet.mu, time_explosion + ) + + r_packet.nu = comov_nu * inverse_new_doppler_factor + r_packet.energy = comov_energy * inverse_new_doppler_factor + if montecarlo_configuration.full_relativity: + r_packet.mu = angle_aberration_CMF_to_LF( + r_packet, time_explosion, r_packet.mu + ) + + +@njit(**njit_dict) +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 + """ + + old_doppler_factor = get_doppler_factor( + r_packet.r, r_packet.mu, time_explosion + ) + r_packet.mu = get_random_mu() + + inverse_new_doppler_factor = get_inverse_doppler_factor( + r_packet.r, r_packet.mu, time_explosion + ) + + comov_energy = r_packet.energy * old_doppler_factor + r_packet.energy = comov_energy * inverse_new_doppler_factor + + if line_interaction_type == LineInteractionType.SCATTER: + line_emission( + r_packet, r_packet.next_line_id, time_explosion, numba_plasma + ) + else: # includes both macro atom and downbranch - encoded in the transition probabilities + emission_line_id = macro_atom(r_packet, numba_plasma) + line_emission(r_packet, emission_line_id, time_explosion, numba_plasma) + + +@njit(**njit_dict) +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.last_line_interaction_out_id = emission_line_id + + if emission_line_id != r_packet.next_line_id: + pass + inverse_doppler_factor = get_inverse_doppler_factor( + r_packet.r, r_packet.mu, time_explosion + ) + r_packet.nu = ( + numba_plasma.line_list_nu[emission_line_id] * inverse_doppler_factor + ) + r_packet.next_line_id = emission_line_id + 1 + nu_line = numba_plasma.line_list_nu[emission_line_id] + + if emission_line_id != (len(numba_plasma.line_list_nu) - 1): + test_for_close_line( + r_packet, emission_line_id + 1, nu_line, numba_plasma + ) + + if montecarlo_configuration.full_relativity: + r_packet.mu = angle_aberration_CMF_to_LF( + r_packet, time_explosion, r_packet.mu + ) diff --git a/tardis/montecarlo/montecarlo_numba/macro_atom.py b/tardis/montecarlo/montecarlo_numba/macro_atom.py new file mode 100644 index 00000000000..8f8147217e6 --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/macro_atom.py @@ -0,0 +1,74 @@ +import numpy as np +from enum import IntEnum + +from numba import njit +from tardis.montecarlo.montecarlo_numba import njit_dict + + +class MacroAtomError(ValueError): + pass + + +class MacroAtomTransitionType(IntEnum): + INTERNAL_UP = 1 + INTERNAL_DOWN = 0 + BB_EMISSION = -1 + BF_EMISSION = -2 + FF_EMISSION = -3 + ADIABATIC_COOLING = -4 + + +@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 + + Returns + ------- + + """ + activation_level_id = numba_plasma.line2macro_level_upper[ + r_packet.next_line_id + ] + current_transition_type = 0 + + while current_transition_type >= 0: + probability = 0.0 + probability_event = np.random.random() + + block_start = numba_plasma.macro_block_references[activation_level_id] + block_end = numba_plasma.macro_block_references[activation_level_id + 1] + + # looping through the transition probabilities + for transition_id in range(block_start, block_end): + + transition_probability = numba_plasma.transition_probabilities[ + transition_id, r_packet.current_shell_id + ] + + probability += transition_probability + + if probability > probability_event: + activation_level_id = numba_plasma.destination_level_id[ + transition_id + ] + current_transition_type = numba_plasma.transition_type[ + transition_id + ] + break + + else: + raise MacroAtomError( + "MacroAtom ran out of the block. This should not happen as " + "the sum of probabilities is normalized to 1 and " + "the probability_event should be less than 1" + ) + + if current_transition_type == MacroAtomTransitionType.BB_EMISSION: + return numba_plasma.transition_line_id[transition_id] + else: + raise MacroAtomError("MacroAtom currently only allows BB transitions") diff --git a/tardis/montecarlo/montecarlo_numba/montecarlo_logger.py b/tardis/montecarlo/montecarlo_numba/montecarlo_logger.py new file mode 100644 index 00000000000..dbdea15cde2 --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/montecarlo_logger.py @@ -0,0 +1,67 @@ +import logging +from functools import wraps + +DEBUG_MODE = False +LOG_FILE = "montecarlo_log.log" +BUFFER = 1 +ticker = 1 + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) +logger.handlers = [] + +if LOG_FILE: + logger.propagate = False + console_handler = logging.FileHandler(LOG_FILE) +else: + console_handler = logging.StreamHandler() + +console_handler.setLevel(logging.DEBUG) +console_formatter = logging.Formatter("%(name)s - %(levelname)s - %(message)s") +console_handler.setFormatter(console_formatter) +logger.addHandler(console_handler) + + +def log_decorator(func): + """ + Decorator to log functions while in debug mode, i.e., when + `debug_montecarlo` is True in the config. Works for + `@jit'd and `@njit`'d functions, but with a significant speed + penalty. + + TODO: in nopython mode: do I need a context manager? + + Input: + func : (function) function to be logged. + + Output: + wrapper : (function) wrapper to the function being logged. + """ + + # to ensure that calling `help` on the decorated function still works + # @wraps(func) + def wrapper(*args, **kwargs): + """ + Wrapper to the log_decorator. + + When called, it has the side effect of + logging every `BUFFER` input and output to `func`, if `DEBUG_MODE` is + `True`. + + Input: + *args : arguments to be passed to `func`. + **kwargs : keyword arguments to be passed to `func`. + + Output: + result : result of calling `func` on `*args` and `**kwargs`. + """ + result = func(*args, **kwargs) + if DEBUG_MODE: + global ticker + ticker += 1 + if ticker % BUFFER == 0: # without a buffer, performance suffers + logger.debug(f"Func: {func.__name__}. Input: {(args, kwargs)}") + logger.debug(f"Output: {result}.") + return result + + return wrapper diff --git a/tardis/montecarlo/montecarlo_numba/numba_config.py b/tardis/montecarlo/montecarlo_numba/numba_config.py new file mode 100644 index 00000000000..a99a40a428d --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/numba_config.py @@ -0,0 +1,8 @@ +from tardis import constants as const + +SIGMA_THOMSON = const.sigma_T.to("cm^2").value +CLOSE_LINE_THRESHOLD = 1e-7 +C_SPEED_OF_LIGHT = const.c.to("cm/s").value +MISS_DISTANCE = 1e99 + +ENABLE_FULL_RELATIVITY = False diff --git a/tardis/montecarlo/montecarlo_numba/numba_interface.py b/tardis/montecarlo/montecarlo_numba/numba_interface.py new file mode 100644 index 00000000000..7c5945344e4 --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/numba_interface.py @@ -0,0 +1,321 @@ +from enum import IntEnum + +from numba import float64, int64, boolean +from numba.experimental import jitclass +import numpy as np + +from astropy import units as u +from tardis import constants as const + +from tardis.montecarlo import ( + montecarlo_configuration as montecarlo_configuration, +) + + +C_SPEED_OF_LIGHT = const.c.to("cm/s").value + +numba_model_spec = [ + ("r_inner", float64[:]), + ("r_outer", float64[:]), + ("time_explosion", float64), +] + + +@jitclass(numba_model_spec) +class NumbaModel(object): + def __init__(self, r_inner, r_outer, time_explosion): + """ + Model for the Numba mode + + Parameters + ---------- + r_inner: numpy.ndarray + r_outer: numpy.ndarray + time_explosion: float + """ + self.r_inner = r_inner + self.r_outer = r_outer + self.time_explosion = time_explosion + + +numba_plasma_spec = [ + ("electron_density", float64[:]), + ("line_list_nu", float64[:]), + ("tau_sobolev", float64[:, :]), + ("transition_probabilities", float64[:, :]), + ("line2macro_level_upper", int64[:]), + ("macro_block_references", int64[:]), + ("transition_type", int64[:]), + ("destination_level_id", int64[:]), + ("transition_line_id", int64[:]), +] + + +@jitclass(numba_plasma_spec) +class NumbaPlasma(object): + def __init__( + self, + electron_density, + line_list_nu, + tau_sobolev, + transition_probabilities, + line2macro_level_upper, + macro_block_references, + transition_type, + destination_level_id, + transition_line_id, + ): + """ + 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 + """ + + self.electron_density = electron_density + self.line_list_nu = line_list_nu + self.tau_sobolev = tau_sobolev + + #### Macro Atom transition probabilities + self.transition_probabilities = transition_probabilities + self.line2macro_level_upper = line2macro_level_upper + + self.macro_block_references = macro_block_references + self.transition_type = transition_type + + # Destination level is not needed and/or generated for downbranch + self.destination_level_id = destination_level_id + self.transition_line_id = transition_line_id + + +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 + """ + electron_densities = plasma.electron_densities.values + line_list_nu = plasma.atomic_data.lines.nu.values + tau_sobolev = np.ascontiguousarray( + plasma.tau_sobolevs.values.copy(), dtype=np.float64 + ) + if montecarlo_configuration.disable_line_scattering: + tau_sobolev *= 0 + + if line_interaction_type == "scatter": + # to adhere to data types, we must have an array of minimum size 1 + array_size = 1 + transition_probabilities = np.zeros( + (array_size, array_size), dtype=np.float64 + ) # to adhere to data types + line2macro_level_upper = np.zeros(array_size, dtype=np.int64) + macro_block_references = np.zeros(array_size, dtype=np.int64) + transition_type = np.zeros(array_size, dtype=np.int64) + destination_level_id = np.zeros(array_size, dtype=np.int64) + transition_line_id = np.zeros(array_size, dtype=np.int64) + else: + transition_probabilities = np.ascontiguousarray( + plasma.transition_probabilities.values.copy(), dtype=np.float64 + ) + line2macro_level_upper = ( + plasma.atomic_data.lines_upper2macro_reference_idx + ) + macro_block_references = plasma.atomic_data.macro_atom_references[ + "block_references" + ].values + transition_type = plasma.atomic_data.macro_atom_data[ + "transition_type" + ].values + + # Destination level is not needed and/or generated for downbranch + destination_level_id = plasma.atomic_data.macro_atom_data[ + "destination_level_idx" + ].values + transition_line_id = plasma.atomic_data.macro_atom_data[ + "lines_idx" + ].values + + return NumbaPlasma( + electron_densities, + line_list_nu, + tau_sobolev, + transition_probabilities, + line2macro_level_upper, + macro_block_references, + transition_type, + destination_level_id, + transition_line_id, + ) + + +packet_collection_spec = [ + ("packets_input_nu", float64[:]), + ("packets_input_mu", float64[:]), + ("packets_input_energy", float64[:]), + ("packets_output_nu", float64[:]), + ("packets_output_energy", float64[:]), +] + + +@jitclass(packet_collection_spec) +class PacketCollection(object): + def __init__( + self, + packets_input_nu, + packets_input_mu, + packets_input_energy, + packets_output_nu, + packets_output_energy, + ): + self.packets_input_nu = packets_input_nu + self.packets_input_mu = packets_input_mu + self.packets_input_energy = packets_input_energy + self.packets_output_nu = packets_output_nu + self.packets_output_energy = packets_output_energy + + +vpacket_collection_spec = [ + ("rpacket_index", int64), + ("spectrum_frequency", float64[:]), + ("v_packet_spawn_start_frequency", float64), + ("v_packet_spawn_end_frequency", float64), + ("nus", float64[:]), + ("energies", float64[:]), + ("idx", int64), + ("number_of_vpackets", int64), + ("length", int64), + ("last_interaction_in_nu", float64), + ("last_interaction_type", int64), + ("last_interaction_in_id", int64), + ("last_interaction_out_id", int64), +] + + +@jitclass(vpacket_collection_spec) +class VPacketCollection(object): + def __init__( + self, + rpacket_index, + spectrum_frequency, + v_packet_spawn_start_frequency, + v_packet_spawn_end_frequency, + number_of_vpackets, + temporary_v_packet_bins, + ): + self.spectrum_frequency = spectrum_frequency + self.v_packet_spawn_start_frequency = v_packet_spawn_start_frequency + self.v_packet_spawn_end_frequency = v_packet_spawn_end_frequency + self.nus = np.empty(temporary_v_packet_bins, dtype=np.float64) + self.energies = np.empty(temporary_v_packet_bins, dtype=np.float64) + self.number_of_vpackets = number_of_vpackets + self.last_interaction_in_nu = 0.0 + self.last_interaction_type = -1 + self.last_interaction_in_id = -1 + self.last_interaction_out_id = -1 + self.idx = 0 + self.rpacket_index = rpacket_index + self.length = temporary_v_packet_bins + + def set_properties( + self, + nu, + energy, + last_interaction_in_nu, + last_interaction_type, + last_interaction_in_id, + last_interaction_out_id, + ): + if self.idx >= self.length: + temp_length = self.length * 2 + self.number_of_vpackets + temp_nus = np.empty(temp_length, dtype=np.float64) + temp_energies = np.empty(temp_length, dtype=np.float64) + temp_nus[: self.length] = self.nus + temp_energies[: self.length] = self.energies + + self.nus = temp_nus + self.energies = temp_energies + self.length = temp_length + + self.nus[self.idx] = nu + self.energies[self.idx] = energy + self.last_interaction_type = last_interaction_type + self.last_interaction_in_nu = last_interaction_in_nu + self.last_interaction_in_id = last_interaction_in_id + self.last_interaction_out_id = last_interaction_out_id + self.idx += 1 + + +estimators_spec = [ + ("j_estimator", float64[:]), + ("nu_bar_estimator", float64[:]), + ("j_blue_estimator", float64[:, :]), + ("Edotlu_estimator", float64[:, :]), +] + + +@jitclass(estimators_spec) +class Estimators(object): + def __init__( + self, j_estimator, nu_bar_estimator, j_blue_estimator, Edotlu_estimator + ): + self.j_estimator = j_estimator + self.nu_bar_estimator = nu_bar_estimator + self.j_blue_estimator = j_blue_estimator + self.Edotlu_estimator = Edotlu_estimator + + +def configuration_initialize(runner, number_of_vpackets): + if runner.line_interaction_type == "macroatom": + montecarlo_configuration.line_interaction_type = ( + LineInteractionType.MACROATOM + ) + elif runner.line_interaction_type == "downbranch": + montecarlo_configuration.line_interaction_type = ( + LineInteractionType.DOWNBRANCH + ) + elif runner.line_interaction_type == "scatter": + montecarlo_configuration.line_interaction_type = ( + LineInteractionType.SCATTER + ) + else: + raise ValueError( + f'Line interaction type must be one of "macroatom",' + f'"downbranch", or "scatter" but is ' + f"{runner.line_interaction_type}" + ) + montecarlo_configuration.number_of_vpackets = number_of_vpackets + montecarlo_configuration.temporary_v_packet_bins = number_of_vpackets + montecarlo_configuration.full_relativity = runner.enable_full_relativity + montecarlo_configuration.montecarlo_seed = runner.seed + montecarlo_configuration.single_packet_seed = runner.single_packet_seed + montecarlo_configuration.v_packet_spawn_start_frequency = ( + runner.virtual_spectrum_spawn_range.end.to( + u.Hz, equivalencies=u.spectral() + ).value + ) + montecarlo_configuration.v_packet_spawn_end_frequency = ( + runner.virtual_spectrum_spawn_range.start.to( + u.Hz, equivalencies=u.spectral() + ).value + ) + montecarlo_configuration.VPACKET_LOGGING = runner.virt_logging + + +# class TrackRPacket(object): +class LineInteractionType(IntEnum): + SCATTER = 0 + DOWNBRANCH = 1 + MACROATOM = 2 diff --git a/tardis/montecarlo/montecarlo_numba/r_packet.py b/tardis/montecarlo/montecarlo_numba/r_packet.py new file mode 100644 index 00000000000..18266daa976 --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/r_packet.py @@ -0,0 +1,624 @@ +import numpy as np +from enum import IntEnum +from numba import int64, float64, boolean +from numba import njit +from numba.experimental import jitclass + +import math +from tardis.montecarlo.montecarlo_numba import njit_dict, numba_config +from tardis.montecarlo import ( + montecarlo_configuration as montecarlo_configuration, +) +from tardis.montecarlo.montecarlo_numba.numba_config import ( + CLOSE_LINE_THRESHOLD, + C_SPEED_OF_LIGHT, + MISS_DISTANCE, + SIGMA_THOMSON, +) + + +class MonteCarloException(ValueError): + pass + + +class PacketStatus(IntEnum): + IN_PROCESS = 0 + EMITTED = 1 + REABSORBED = 2 + + +class InteractionType(IntEnum): + BOUNDARY = 1 + LINE = 2 + ESCATTERING = 3 + + +rpacket_spec = [ + ("r", float64), + ("mu", float64), + ("nu", float64), + ("energy", float64), + ("next_line_id", int64), + ("current_shell_id", int64), + ("status", int64), + ("seed", int64), + ("index", int64), + ("is_close_line", boolean), + ("last_interaction_type", int64), + ("last_line_interaction_out_id", int64), +] + + +@jitclass(rpacket_spec) +class RPacket(object): + def __init__(self, r, mu, nu, energy, seed, index=0, is_close_line=False): + self.r = r + self.mu = mu + self.nu = nu + self.energy = energy + self.current_shell_id = 0 + self.status = PacketStatus.IN_PROCESS + self.seed = seed + self.index = index + self.is_close_line = is_close_line + self.last_interaction_type = -1 + self.last_line_interaction_out_id = -1 + + def initialize_line_id(self, numba_plasma, numba_model): + inverse_line_list_nu = numba_plasma.line_list_nu[::-1] + doppler_factor = get_doppler_factor( + self.r, self.mu, numba_model.time_explosion + ) + comov_nu = self.nu * doppler_factor + next_line_id = len(numba_plasma.line_list_nu) - np.searchsorted( + inverse_line_list_nu, comov_nu + ) + if next_line_id == len(numba_plasma.line_list_nu): + next_line_id -= 1 + self.next_line_id = next_line_id + + +@njit(**njit_dict) +def calculate_distance_boundary(r, mu, r_inner, r_outer): + """ + Calculate distance to shell boundary in cm. + + Parameters: + ------------ + + r: float + radial coordinate of the RPacket + mu: float + cosine of the direction of movement + r_inner: float + inner radius of current shell + r_outer: float + outer radius of current shell + """ + + delta_shell = 0 + if mu > 0.0: + # direction outward + distance = math.sqrt(r_outer * r_outer + ((mu * mu - 1.0) * r * r)) - ( + r * mu + ) + delta_shell = 1 + else: + # going inward + check = r_inner * r_inner + (r * r * (mu * mu - 1.0)) + + if check >= 0.0: + # hit inner boundary + distance = -r * mu - math.sqrt(check) + delta_shell = -1 + else: + # miss inner boundary + distance = math.sqrt( + r_outer * r_outer + ((mu * mu - 1.0) * r * r) + ) - (r * mu) + delta_shell = 1 + + return distance, delta_shell + + +# @log_decorator +#'float64(RPacket, float64, int64, float64, float64)' +@njit(**njit_dict) +def calculate_distance_line( + r_packet, comov_nu, is_last_line, nu_line, time_explosion +): + """ + Calculate distance until RPacket is in resonance with the next line + Parameters + ---------- + r_packet: RPacket + comov_nu: float + comoving frequency at the CURRENT position of the RPacket + nu_line: float + line to check the distance to + time_explosion: float + time since explosion in seconds + is_last_line: bool + return MISS_DISTANCE if at the end of the line list + + Returns + ------- + + """ + + nu = r_packet.nu + + if is_last_line: + return MISS_DISTANCE + + nu_diff = comov_nu - nu_line + + # for numerical reasons, if line is too close, we set the distance to 0. + if r_packet.is_close_line: + nu_diff = 0.0 + r_packet.is_close_line = False + + if nu_diff >= 0: + distance = (nu_diff / nu) * C_SPEED_OF_LIGHT * time_explosion + else: + print("WARNING: nu difference is less than 0.0") + raise MonteCarloException( + "nu difference is less than 0.0; for more" + " information, see print statement beforehand" + ) + + if numba_config.ENABLE_FULL_RELATIVITY: + return calculate_distance_line_full_relativity( + nu_line, nu, time_explosion, r_packet + ) + return distance + + +@njit(**njit_dict) +def calculate_distance_line_full_relativity( + nu_line, nu, time_explosion, r_packet +): + # distance = - mu * r + (ct - nu_r * nu_r * sqrt(ct * ct - (1 + r * r * (1 - mu * mu) * (1 + pow(nu_r, -2))))) / (1 + nu_r * nu_r); + nu_r = nu_line / nu + ct = C_SPEED_OF_LIGHT * time_explosion + distance = -r_packet.mu * r_packet.r + ( + ct + - nu_r + * nu_r + * math.sqrt( + ct * ct + - ( + 1 + + r_packet.r + * r_packet.r + * (1 - r_packet.mu * r_packet.mu) + * (1 + 1.0 / (nu_r * nu_r)) + ) + ) + ) / (1 + nu_r * nu_r) + return distance + + +@njit(**njit_dict) +def calculate_distance_electron(electron_density, tau_event): + """ + Calculate distance to Thomson Scattering + + electron_density: float + tau_event: float + """ + # add full_relativity here + return tau_event / (electron_density * numba_config.SIGMA_THOMSON) + + +@njit(**njit_dict) +def calculate_tau_electron(electron_density, distance): + """ + Calculate tau for Thomson scattering + + Parameters: + ------------ + + electron_density: float + distance: float + """ + return electron_density * numba_config.SIGMA_THOMSON * distance + + +@njit(**njit_dict) +def get_doppler_factor(r, mu, time_explosion): + inv_c = 1 / C_SPEED_OF_LIGHT + inv_t = 1 / time_explosion + beta = r * inv_t * inv_c + if not numba_config.ENABLE_FULL_RELATIVITY: + return get_doppler_factor_partial_relativity(mu, beta) + else: + return get_doppler_factor_full_relativity(mu, beta) + + +@njit(**njit_dict) +def get_doppler_factor_partial_relativity(mu, beta): + return 1.0 - mu * beta + + +@njit(**njit_dict) +def get_doppler_factor_full_relativity(mu, beta): + return (1.0 - mu * beta) / math.sqrt(1 - beta * beta) + + +@njit(**njit_dict) +def get_inverse_doppler_factor(r, mu, time_explosion): + """ + Calculate doppler factor for frame transformation + + Parameters: + ------------ + + r: float + mu: float + time_explosion: float + """ + inv_c = 1 / C_SPEED_OF_LIGHT + inv_t = 1 / time_explosion + beta = r * inv_t * inv_c + if not numba_config.ENABLE_FULL_RELATIVITY: + return get_inverse_doppler_factor_partial_relativity(mu, beta) + else: + return get_inverse_doppler_factor_full_relativity(mu, beta) + + +@njit(**njit_dict) +def get_inverse_doppler_factor_partial_relativity(mu, beta): + return 1.0 / (1.0 - mu * beta) + + +@njit(**njit_dict) +def get_inverse_doppler_factor_full_relativity(mu, beta): + return (1.0 + mu * beta) / math.sqrt(1 - beta * beta) + + +@njit(**njit_dict) +def get_random_mu(): + return 2.0 * np.random.random() - 1.0 + + +@njit(**njit_dict) +def update_line_estimators( + estimators, r_packet, cur_line_id, distance_trace, time_explosion +): + """ + Function to update the line estimators + + Parameters + ---------- + estimators: Estimators + r_packet: RPacket + cur_line_id: int + distance_trace: float + time_explosion: float + + """ + + """ Actual calculation - simplified below + r_interaction = math.sqrt(r_packet.r**2 + distance_trace**2 + + 2 * r_packet.r * distance_trace * r_packet.mu) + mu_interaction = (r_packet.mu * r_packet.r + distance_trace) / r_interaction + doppler_factor = 1.0 - mu_interaction * r_interaction / + ( time_explosion * C) + """ + + if not numba_config.ENABLE_FULL_RELATIVITY: + energy = calc_packet_energy(r_packet, distance_trace, time_explosion) + else: + energy = calc_packet_energy_full_relativity(r_packet) + + estimators.j_blue_estimator[cur_line_id, r_packet.current_shell_id] += ( + energy / r_packet.nu + ) + estimators.Edotlu_estimator[ + cur_line_id, r_packet.current_shell_id + ] += energy + + +@njit(**njit_dict) +def calc_packet_energy_full_relativity(r_packet): + # accurate to 1 / gamma - according to C. Vogl + return r_packet.energy + + +@njit(**njit_dict) +def calc_packet_energy(r_packet, distance_trace, time_explosion): + doppler_factor = 1.0 - ( + (distance_trace + r_packet.mu * r_packet.r) + / (time_explosion * C_SPEED_OF_LIGHT) + ) + energy = r_packet.energy * doppler_factor + return energy + + +@njit(**njit_dict) +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 + + Returns + ------- + + """ + + r_inner = numba_model.r_inner[r_packet.current_shell_id] + r_outer = numba_model.r_outer[r_packet.current_shell_id] + + distance_boundary, delta_shell = 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 + + # defining taus + tau_event = -np.log(np.random.random()) + tau_trace_line_combined = 0.0 + + # e scattering initialization + + cur_electron_density = numba_plasma.electron_density[ + r_packet.current_shell_id + ] + distance_electron = calculate_distance_electron( + cur_electron_density, tau_event + ) + + # Calculating doppler factor + doppler_factor = get_doppler_factor( + r_packet.r, r_packet.mu, numba_model.time_explosion + ) + comov_nu = r_packet.nu * doppler_factor + + cur_line_id = start_line_id # initializing varibale for Numba + # - do not remove + last_line_id = len(numba_plasma.line_list_nu) - 1 + + for cur_line_id in range(start_line_id, len(numba_plasma.line_list_nu)): + + # Going through the lines + nu_line = numba_plasma.line_list_nu[cur_line_id] + nu_line_last_interaction = numba_plasma.line_list_nu[cur_line_id - 1] + + # Getting the tau for the next line + tau_trace_line = numba_plasma.tau_sobolev[ + cur_line_id, r_packet.current_shell_id + ] + + # Adding it to the tau_trace_line_combined + tau_trace_line_combined += tau_trace_line + + # Calculating the distance until the current photons co-moving nu + # redshifts to the line frequency + is_last_line = cur_line_id == last_line_id + + distance_trace = calculate_distance_line( + r_packet, + comov_nu, + is_last_line, + nu_line, + numba_model.time_explosion, + ) + + # calculating the tau electron of how far the trace has progressed + tau_trace_electron = calculate_tau_electron( + cur_electron_density, distance_trace + ) + + # calculating the trace + tau_trace_combined = tau_trace_line_combined + tau_trace_electron + + if ( + (distance_boundary <= distance_trace) + and (distance_boundary <= distance_electron) + ) and distance_trace != 0.0: + interaction_type = InteractionType.BOUNDARY # BOUNDARY + r_packet.next_line_id = cur_line_id + distance = distance_boundary + break + + if ( + (distance_electron < distance_trace) + and (distance_electron < distance_boundary) + ) and distance_trace != 0.0: + interaction_type = InteractionType.ESCATTERING + # print('scattering') + distance = distance_electron + r_packet.next_line_id = cur_line_id + break + + # Updating the J_b_lu and E_dot_lu + # This means we are still looking for line interaction and have not + # been kicked out of the path by boundary or electron interaction + + update_line_estimators( + estimators, + r_packet, + cur_line_id, + distance_trace, + numba_model.time_explosion, + ) + + if ( + tau_trace_combined > tau_event + and not montecarlo_configuration.disable_line_scattering + ): + interaction_type = InteractionType.LINE # Line + r_packet.next_line_id = cur_line_id + distance = distance_trace + break + + if not is_last_line: + test_for_close_line( + r_packet, cur_line_id + 1, nu_line, numba_plasma + ) + + # Recalculating distance_electron using tau_event - + # tau_trace_line_combined + distance_electron = calculate_distance_electron( + cur_electron_density, tau_event - tau_trace_line_combined + ) + + else: # Executed when no break occurs in the for loop + # We are beyond the line list now and the only next thing is to see + # if we are interacting with the boundary or electron scattering + if cur_line_id == (len(numba_plasma.line_list_nu) - 1): + # Treatment for last line + cur_line_id += 1 + if distance_electron < distance_boundary: + distance = distance_electron + interaction_type = InteractionType.ESCATTERING + # print('scattering') + else: + distance = distance_boundary + interaction_type = InteractionType.BOUNDARY + + # r_packet.next_line_id = cur_line_id + + r_packet.last_interaction_type = interaction_type + + return distance, interaction_type, delta_shell + + +@njit(**njit_dict) +def move_r_packet(r_packet, distance, time_explosion, numba_estimator): + """ + Move packet a distance and recalculate the new angle mu + + Parameters + ---------- + + r_packet: tardis.montecarlo.montecarlo_numba.r_packet.RPacket + r_packet objects + time_explosion: float + time since explosion in s + numba_estimator: tardis.montecarlo.montecarlo_numba.numba_interface.NumbaEstimator + Estimators object + distance : float + distance in cm + """ + + doppler_factor = get_doppler_factor(r_packet.r, r_packet.mu, time_explosion) + + r = r_packet.r + if distance > 0.0: + new_r = np.sqrt( + r * r + distance * distance + 2.0 * r * distance * r_packet.mu + ) + r_packet.mu = (r_packet.mu * r + distance) / new_r + r_packet.r = new_r + + comov_nu = r_packet.nu * doppler_factor + comov_energy = r_packet.energy * doppler_factor + + if not numba_config.ENABLE_FULL_RELATIVITY: + set_estimators( + r_packet, distance, numba_estimator, comov_nu, comov_energy + ) + + else: + distance = distance * doppler_factor + set_estimators_full_relativity( + r_packet, + distance, + numba_estimator, + comov_nu, + comov_energy, + doppler_factor, + ) + + +@njit(**njit_dict) +def set_estimators(r_packet, distance, numba_estimator, comov_nu, comov_energy): + """ + Updating the estimators + """ + numba_estimator.j_estimator[r_packet.current_shell_id] += ( + comov_energy * distance + ) + numba_estimator.nu_bar_estimator[r_packet.current_shell_id] += ( + comov_energy * distance * comov_nu + ) + + +@njit(**njit_dict) +def set_estimators_full_relativity( + r_packet, distance, numba_estimator, comov_nu, comov_energy, doppler_factor +): + numba_estimator.j_estimator[r_packet.current_shell_id] += ( + comov_energy * distance * doppler_factor + ) + numba_estimator.nu_bar_estimator[r_packet.current_shell_id] += ( + comov_energy * distance * comov_nu * doppler_factor + ) + + +@njit(**njit_dict) +def move_packet_across_shell_boundary(packet, delta_shell, no_of_shells): + """ + Move packet across shell boundary - realizing if we are still in the simulation or have + moved out through the inner boundary or outer boundary and updating packet + status. + + Parameters + ---------- + distance : float + distance to move to shell boundary + + delta_shell: int + is +1 if moving outward or -1 if moving inward + + no_of_shells: int + number of shells in TARDIS simulation + """ + next_shell_id = packet.current_shell_id + delta_shell + + if next_shell_id >= no_of_shells: + packet.status = PacketStatus.EMITTED + elif next_shell_id < 0: + packet.status = PacketStatus.REABSORBED + else: + packet.current_shell_id = next_shell_id + + +@njit(**njit_dict) +def angle_aberration_CMF_to_LF(r_packet, time_explosion, mu): + """ + Converts angle aberration from comoving frame to + laboratory frame. + """ + ct = C_SPEED_OF_LIGHT * time_explosion + beta = r_packet.r / (ct) + return (r_packet.mu + beta) / (1.0 + beta * mu) + + +@njit(**njit_dict) +def angle_aberration_LF_to_CMF(r_packet, time_explosion, mu): + """ + + c code: + double beta = rpacket_get_r (packet) * storage->inverse_time_explosion * INVERSE_C; + return (mu - beta) / (1.0 - beta * mu); + """ + ct = C_SPEED_OF_LIGHT * time_explosion + beta = r_packet.r / (ct) + return (mu - beta) / (1.0 - beta * mu) + + +@njit(**njit_dict) +def test_for_close_line(r_packet, line_id, nu_line, numba_plasma): + r_packet.is_close_line = abs( + numba_plasma.line_list_nu[line_id] - nu_line + ) < (nu_line * CLOSE_LINE_THRESHOLD) diff --git a/tardis/montecarlo/montecarlo_numba/single_packet_loop.py b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py new file mode 100644 index 00000000000..2dd80e66fff --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/single_packet_loop.py @@ -0,0 +1,153 @@ +from numba import njit +import numpy as np + +from tardis.montecarlo.montecarlo_numba.r_packet import ( + InteractionType, + PacketStatus, + get_inverse_doppler_factor, + trace_packet, + move_packet_across_shell_boundary, + move_r_packet, + MonteCarloException, +) +from tardis.montecarlo.montecarlo_numba.interaction import ( + thomson_scatter, + line_scatter, +) +from tardis.montecarlo.montecarlo_numba.numba_interface import ( + LineInteractionType, +) +from tardis.montecarlo import ( + montecarlo_configuration as montecarlo_configuration, +) + +from tardis.montecarlo.montecarlo_numba.vpacket import trace_vpacket_volley + +from tardis import constants as const + +C_SPEED_OF_LIGHT = const.c.to("cm/s").value + +from tardis.montecarlo.montecarlo_numba.montecarlo_logger import log_decorator +from tardis.montecarlo.montecarlo_numba import montecarlo_logger as mc_logger + +# @log_decorator +@njit +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 + + Returns + ------- + : None + + This function does not return anything but changes the r_packet object + and if virtual packets are requested - also updates the vpacket_collection + """ + + line_interaction_type = montecarlo_configuration.line_interaction_type + + if montecarlo_configuration.full_relativity: + set_packet_props_full_relativity(r_packet, numba_model) + else: + set_packet_props_partial_relativity(r_packet, numba_model) + r_packet.initialize_line_id(numba_plasma, numba_model) + + trace_vpacket_volley( + r_packet, vpacket_collection, numba_model, numba_plasma + ) + + if mc_logger.DEBUG_MODE: + r_packet_track_nu = [r_packet.nu] + r_packet_track_mu = [r_packet.mu] + r_packet_track_r = [r_packet.r] + r_packet_track_interaction = [InteractionType.BOUNDARY] + r_packet_track_distance = [0.0] + + while r_packet.status == PacketStatus.IN_PROCESS: + distance, interaction_type, delta_shell = trace_packet( + r_packet, numba_model, numba_plasma, estimators + ) + + if interaction_type == InteractionType.BOUNDARY: + move_r_packet( + r_packet, distance, numba_model.time_explosion, estimators + ) + move_packet_across_shell_boundary( + r_packet, delta_shell, len(numba_model.r_inner) + ) + + elif interaction_type == InteractionType.LINE: + move_r_packet( + r_packet, distance, numba_model.time_explosion, estimators + ) + line_scatter( + r_packet, + numba_model.time_explosion, + line_interaction_type, + numba_plasma, + ) + trace_vpacket_volley( + r_packet, vpacket_collection, numba_model, numba_plasma + ) + + elif interaction_type == InteractionType.ESCATTERING: + move_r_packet( + r_packet, distance, numba_model.time_explosion, estimators + ) + thomson_scatter(r_packet, numba_model.time_explosion) + + trace_vpacket_volley( + r_packet, vpacket_collection, numba_model, numba_plasma + ) + if mc_logger.DEBUG_MODE: + r_packet_track_nu.append(r_packet.nu) + r_packet_track_mu.append(r_packet.mu) + r_packet_track_r.append(r_packet.r) + r_packet_track_interaction.append(interaction_type) + r_packet_track_distance.append(distance) + + if mc_logger.DEBUG_MODE: + return ( + r_packet_track_nu, + r_packet_track_mu, + r_packet_track_r, + r_packet_track_interaction, + r_packet_track_distance, + ) + + # check where else initialize line ID happens! + + +@njit +def set_packet_props_partial_relativity(r_packet, numba_model): + inverse_doppler_factor = get_inverse_doppler_factor( + r_packet.r, + r_packet.mu, + numba_model.time_explosion, + ) + r_packet.nu *= inverse_doppler_factor + r_packet.energy *= inverse_doppler_factor + + +@njit +def set_packet_props_full_relativity(r_packet, numba_model): + beta = (r_packet.r / numba_model.time_explosion) / C_SPEED_OF_LIGHT + + inverse_doppler_factor = get_inverse_doppler_factor( + r_packet.r, + r_packet.mu, + numba_model.time_explosion, + ) + + r_packet.nu *= inverse_doppler_factor + r_packet.energy *= inverse_doppler_factor + r_packet.mu = (r_packet.mu + beta) / (1 + beta * r_packet.mu) diff --git a/tardis/montecarlo/montecarlo_numba/tests/conftest.py b/tardis/montecarlo/montecarlo_numba/tests/conftest.py new file mode 100644 index 00000000000..20223587ed4 --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/tests/conftest.py @@ -0,0 +1,149 @@ +from copy import deepcopy + +import pytest +import numpy as np +from numba import njit + +from tardis.simulation import Simulation +from tardis.montecarlo.montecarlo_numba import RPacket, PacketCollection + + +from tardis.montecarlo.montecarlo_numba.numba_interface import ( + numba_plasma_initialize, + NumbaModel, + Estimators, + VPacketCollection, +) + + +@pytest.fixture(scope="package") +def nb_simulation_verysimple(config_verysimple, atomic_dataset): + atomic_data = deepcopy(atomic_dataset) + sim = Simulation.from_config(config_verysimple, atom_data=atomic_data) + sim.iterate(10) + return sim + + +@pytest.fixture() +def verysimple_collection(nb_simulation_verysimple): + runner = nb_simulation_verysimple.runner + return PacketCollection( + runner.input_nu, + runner.input_mu, + runner.input_energy, + runner._output_nu, + runner._output_energy, + ) + + +@pytest.fixture(scope="package") +def verysimple_numba_plasma(nb_simulation_verysimple): + return numba_plasma_initialize( + nb_simulation_verysimple.plasma, line_interaction_type="macroatom" + ) + + +@pytest.fixture(scope="package") +def verysimple_numba_model(nb_simulation_verysimple): + runner = nb_simulation_verysimple.runner + model = nb_simulation_verysimple.model + return NumbaModel( + runner.r_inner_cgs, + runner.r_outer_cgs, + model.time_explosion.to("s").value, + ) + + +@pytest.fixture(scope="package") +def verysimple_estimators(nb_simulation_verysimple): + runner = nb_simulation_verysimple.runner + return Estimators( + runner.j_estimator, + runner.nu_bar_estimator, + runner.j_blue_estimator, + runner.Edotlu_estimator, + ) + + +@pytest.fixture(scope="package") +def verysimple_vpacket_collection(nb_simulation_verysimple): + spectrum_frequency = ( + nb_simulation_verysimple.runner.spectrum_frequency.value + ) + return VPacketCollection( + rpacket_index=0, + spectrum_frequency=spectrum_frequency, + number_of_vpackets=0, + v_packet_spawn_start_frequency=0, + v_packet_spawn_end_frequency=np.inf, + temporary_v_packet_bins=0, + ) + + +@pytest.fixture(scope="package") +def verysimple_3vpacket_collection(nb_simulation_verysimple): + spectrum_frequency = ( + nb_simulation_verysimple.runner.spectrum_frequency.value + ) + return VPacketCollection( + rpacket_index=0, + spectrum_frequency=spectrum_frequency, + number_of_vpackets=3, + v_packet_spawn_start_frequency=0, + v_packet_spawn_end_frequency=np.inf, + temporary_v_packet_bins=0, + ) + + +@pytest.fixture(scope="package") +def verysimple_packet_collection(nb_simulation_verysimple): + runner = nb_simulation_verysimple.runner + return PacketCollection( + runner.input_nu, + runner.input_mu, + runner.input_energy, + runner._output_nu, + runner._output_energy, + ) + + +@pytest.fixture(scope="function") +def packet(verysimple_packet_collection): + return RPacket( + r=7.5e14, + nu=verysimple_packet_collection.packets_input_nu[0], + mu=verysimple_packet_collection.packets_input_mu[0], + energy=verysimple_packet_collection.packets_input_energy[0], + seed=1963, + index=0, + is_close_line=0, + ) + + +@pytest.fixture(scope="function") +def static_packet(): + return RPacket( + r=7.5e14, + nu=0.4, + mu=0.3, + energy=0.9, + seed=1963, + index=0, + is_close_line=0, + ) + + +@pytest.fixture() +def set_seed_fixture(): + def set_seed(value): + np.random.seed(value) + + return njit(set_seed) + + +@pytest.fixture() +def random_call_fixture(): + def random_call(): + np.random.random() + + return njit(random_call) diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_base.py b/tardis/montecarlo/montecarlo_numba/tests/test_base.py new file mode 100644 index 00000000000..33b4bb1d78c --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/tests/test_base.py @@ -0,0 +1,169 @@ +import pytest +import pandas as pd +import os +import numpy.testing as npt +import numpy as np +from copy import deepcopy +from astropy import units as u + +from tardis.montecarlo import ( + montecarlo_configuration as montecarlo_configuration, +) +import tardis.montecarlo.montecarlo_numba.base as base +from tardis.simulation import Simulation +import tardis.montecarlo.montecarlo_numba.r_packet as r_packet +import tardis.montecarlo.montecarlo_numba.single_packet_loop as spl +from tardis.montecarlo.montecarlo_numba.numba_interface import ( + PacketCollection, + VPacketCollection, + NumbaModel, + numba_plasma_initialize, + Estimators, + configuration_initialize, +) + + +@pytest.mark.xfail(reason="To be implemented") +def test_montecarlo_radial1d(): + assert False + + +def test_montecarlo_main_loop( + config_verysimple, + atomic_dataset, + tardis_ref_path, + tmpdir, + set_seed_fixture, + random_call_fixture, +): + + montecarlo_configuration.LEGACY_MODE_ENABLED = True + + # Load C data from refdata + C_fname = os.path.join(tardis_ref_path, "montecarlo_1e5_compare_data.h5") + expected_nu = pd.read_hdf( + C_fname, key="/simulation/runner/output_nu" + ).values + expected_energy = pd.read_hdf( + C_fname, key="/simulation/runner/output_energy" + ).values + expected_nu_bar_estimator = pd.read_hdf( + C_fname, key="/simulation/runner/nu_bar_estimator" + ).values + expected_j_estimator = pd.read_hdf( + C_fname, key="/simulation/runner/j_estimator" + ).values + + # Setup model config from verysimple + atomic_data = deepcopy(atomic_dataset) + config_verysimple.montecarlo.last_no_of_packets = 1e5 + config_verysimple.montecarlo.no_of_virtual_packets = 0 + config_verysimple.montecarlo.iterations = 1 + config_verysimple.montecarlo.single_packet_seed = 0 + del config_verysimple["config_dirname"] + + sim = Simulation.from_config(config_verysimple, atom_data=atomic_data) + + # Init model + numba_plasma = numba_plasma_initialize( + sim.plasma, line_interaction_type="macroatom" + ) + + runner = sim.runner + model = sim.model + + runner._initialize_geometry_arrays(model) + runner._initialize_estimator_arrays(numba_plasma.tau_sobolev.shape) + runner._initialize_packets(model.t_inner.value, 100000, 0) + + # Init parameters + montecarlo_configuration.v_packet_spawn_start_frequency = ( + runner.virtual_spectrum_spawn_range.end.to( + u.Hz, equivalencies=u.spectral() + ).value + ) + montecarlo_configuration.v_packet_spawn_end_frequency = ( + runner.virtual_spectrum_spawn_range.start.to( + u.Hz, equivalencies=u.spectral() + ).value + ) + montecarlo_configuration.temporary_v_packet_bins = 20000 + montecarlo_configuration.full_relativity = runner.enable_full_relativity + montecarlo_configuration.single_packet_seed = 0 + + # Init packet collection from runner + packet_collection = PacketCollection( + runner.input_nu, + runner.input_mu, + runner.input_energy, + runner._output_nu, + runner._output_energy, + ) + + # Init model from runner + numba_model = NumbaModel( + runner.r_inner_cgs, + runner.r_outer_cgs, + model.time_explosion.to("s").value, + ) + + # Init estimators from runner + estimators = Estimators( + runner.j_estimator, + runner.nu_bar_estimator, + runner.j_blue_estimator, + runner.Edotlu_estimator, + ) + + # Empty vpacket collection + vpacket_collection = VPacketCollection( + 0, np.array([0, 0], dtype=np.float64), 0, np.inf, 0, 0 + ) + + # output arrays + output_nus = np.empty_like(packet_collection.packets_output_nu) + output_energies = np.empty_like(packet_collection.packets_output_nu) + + # IMPORTANT: seeds RNG state within JIT + seed = 23111963 + set_seed_fixture(seed) + for i in range(len(packet_collection.packets_input_nu)): + # Generate packet + packet = r_packet.RPacket( + numba_model.r_inner[0], + packet_collection.packets_input_mu[i], + packet_collection.packets_input_nu[i], + packet_collection.packets_input_energy[i], + seed, + i, + 0, + ) + + # Loop packet + spl.single_packet_loop( + packet, numba_model, numba_plasma, estimators, vpacket_collection + ) + output_nus[i] = packet.nu + if packet.status == r_packet.PacketStatus.REABSORBED: + output_energies[i] = -packet.energy + elif packet.status == r_packet.PacketStatus.EMITTED: + output_energies[i] = packet.energy + + # RNG to match C + random_call_fixture() + + packet_collection.packets_output_energy[:] = output_energies[:] + packet_collection.packets_output_nu[:] = output_nus[:] + + actual_energy = packet_collection.packets_output_energy + actual_nu = packet_collection.packets_output_nu + actual_nu_bar_estimator = estimators.nu_bar_estimator + actual_j_estimator = estimators.j_estimator + + # Compare + npt.assert_allclose( + actual_nu_bar_estimator, expected_nu_bar_estimator, rtol=1e-13 + ) + npt.assert_allclose(actual_j_estimator, expected_j_estimator, rtol=1e-13) + npt.assert_allclose(actual_energy, expected_energy, rtol=1e-13) + npt.assert_allclose(actual_nu, expected_nu, rtol=1e-13) diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py new file mode 100644 index 00000000000..bcdaa8f1736 --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/tests/test_interaction.py @@ -0,0 +1,102 @@ +import pytest +import numpy.testing as npt +import numpy as np +import tardis.montecarlo.montecarlo_numba.interaction as interaction +from tardis.montecarlo.montecarlo_numba.numba_interface import ( + LineInteractionType, +) + + +def test_thomson_scatter(packet, verysimple_numba_model): + init_mu = packet.mu + init_nu = packet.nu + init_energy = packet.energy + time_explosion = verysimple_numba_model.time_explosion + + interaction.thomson_scatter(packet, time_explosion) + + assert np.abs(packet.mu - init_mu) > 1e-7 + assert np.abs(packet.nu - init_nu) > 1e-7 + assert np.abs(packet.energy - init_energy) > 1e-7 + + +@pytest.mark.parametrize( + "line_interaction_type", + [ + LineInteractionType.SCATTER, + LineInteractionType.DOWNBRANCH, + LineInteractionType.MACROATOM, + ], +) +def test_line_scatter( + line_interaction_type, + packet, + verysimple_numba_model, + verysimple_numba_plasma, +): + init_mu = packet.mu + init_nu = packet.nu + init_energy = packet.energy + packet.initialize_line_id(verysimple_numba_plasma, verysimple_numba_model) + time_explosion = verysimple_numba_model.time_explosion + + interaction.line_scatter( + packet, time_explosion, line_interaction_type, verysimple_numba_plasma + ) + + assert np.abs(packet.mu - init_mu) > 1e-7 + assert np.abs(packet.nu - init_nu) > 1e-7 + assert np.abs(packet.energy - init_energy) > 1e-7 + + +@pytest.mark.parametrize( + ["test_packet", "expected"], + [ + ( + { + "mu": 0.8599443103322428, + "emission_line_id": 1000, + "energy": 0.9114437898710559, + "nu": 0.0, + }, + {"mu": 0.8599443103322428, "energy": 0.9114437898710559}, + ), + ( + { + "mu": -0.6975116557422458, + "emission_line_id": 2000, + "energy": 0.8803098648913266, + }, + {"mu": -0.6975116557422458, "energy": 0.8803098648913266}, + ), + ( + { + "mu": -0.7115661419975774, + "emission_line_id": 0, + "energy": 0.8800385929341252, + }, + {"mu": -0.7115661419975774, "energy": 0.8800385929341252}, + ), + ], +) +def test_line_emission( + packet, + verysimple_numba_model, + verysimple_numba_plasma, + test_packet, + expected, +): + emission_line_id = test_packet["emission_line_id"] + packet.mu = test_packet["mu"] + packet.energy = test_packet["energy"] + packet.initialize_line_id(verysimple_numba_plasma, verysimple_numba_model) + + time_explosion = verysimple_numba_model.time_explosion + + interaction.line_emission( + packet, emission_line_id, time_explosion, verysimple_numba_plasma + ) + + assert packet.next_line_id == emission_line_id + 1 + npt.assert_almost_equal(packet.mu, expected["mu"]) + npt.assert_almost_equal(packet.energy, expected["energy"]) diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py b/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py new file mode 100644 index 00000000000..d8442be6153 --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/tests/test_macro_atom.py @@ -0,0 +1,23 @@ +import pytest +import tardis.montecarlo.montecarlo_numba.macro_atom as macro_atom +import numpy as np + + +@pytest.mark.parametrize( + ["seed", "expected"], + [(1963, 10015), (1, 9993), (2111963, 17296), (10000, 9993)], +) +def test_macro_atom( + static_packet, + verysimple_numba_plasma, + verysimple_numba_model, + set_seed_fixture, + seed, + expected, +): + set_seed_fixture(seed) + static_packet.initialize_line_id( + verysimple_numba_plasma, verysimple_numba_model + ) + result = macro_atom.macro_atom(static_packet, verysimple_numba_plasma) + assert result == expected diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_montecarlo_logger.py b/tardis/montecarlo/montecarlo_numba/tests/test_montecarlo_logger.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py b/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py new file mode 100644 index 00000000000..2e1eb9d7441 --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py @@ -0,0 +1,109 @@ +import pytest +import tardis.montecarlo.montecarlo_numba.numba_interface as numba_interface +import numpy.testing as npt +import numpy as np + + +@pytest.mark.parametrize("input_params", ["scatter", "macroatom", "downbranch"]) +def test_numba_plasma_initialize(nb_simulation_verysimple, input_params): + line_interaction_type = input_params + plasma = nb_simulation_verysimple.plasma + actual = numba_interface.numba_plasma_initialize( + plasma, line_interaction_type + ) + + npt.assert_allclose( + actual.electron_density, plasma.electron_densities.values + ) + npt.assert_allclose(actual.line_list_nu, plasma.atomic_data.lines.nu.values) + npt.assert_allclose(actual.tau_sobolev, plasma.tau_sobolevs.values) + if line_interaction_type == "scatter": + empty = np.zeros(1, dtype=np.int64) + npt.assert_allclose( + actual.transition_probabilities, np.zeros((1, 1), dtype=np.float64) + ) + npt.assert_allclose(actual.line2macro_level_upper, empty) + npt.assert_allclose(actual.macro_block_references, empty) + npt.assert_allclose(actual.transition_type, empty) + npt.assert_allclose(actual.destination_level_id, empty) + npt.assert_allclose(actual.transition_line_id, empty) + else: + npt.assert_allclose( + actual.transition_probabilities, + plasma.transition_probabilities.values, + ) + npt.assert_allclose( + actual.line2macro_level_upper, + plasma.atomic_data.lines_upper2macro_reference_idx, + ) + npt.assert_allclose( + actual.macro_block_references, + plasma.atomic_data.macro_atom_references["block_references"].values, + ) + npt.assert_allclose( + actual.transition_type, + plasma.atomic_data.macro_atom_data["transition_type"].values, + ) + npt.assert_allclose( + actual.destination_level_id, + plasma.atomic_data.macro_atom_data["destination_level_idx"].values, + ) + npt.assert_allclose( + actual.transition_line_id, + plasma.atomic_data.macro_atom_data["lines_idx"].values, + ) + + +@pytest.mark.xfail(reason="To be implemented") +def test_configuration_initialize(): + assert False + + +def test_VPacketCollection_set_properties(verysimple_3vpacket_collection): + + assert verysimple_3vpacket_collection.length == 0 + + nus = [3.0e15, 0.0, 1e15, 1e5] + energies = [0.4, 0.1, 0.6, 1e10] + last_interaction_in_nu = 3.0e15 + last_interaction_type = 1 + last_interaction_in_id = 100 + last_interaction_out_id = 1201 + + for (nu, energy) in zip(nus, energies): + verysimple_3vpacket_collection.set_properties( + nu, energy, last_interaction_in_nu, + last_interaction_type, + last_interaction_in_id, + last_interaction_out_id + ) + + npt.assert_array_equal( + verysimple_3vpacket_collection.nus[ + : verysimple_3vpacket_collection.idx + ], + nus, + ) + npt.assert_array_equal( + verysimple_3vpacket_collection.energies[ + : verysimple_3vpacket_collection.idx + ], + energies, + ) + npt.assert_array_equal( + verysimple_3vpacket_collection.last_interaction_in_nu, + last_interaction_in_nu, + ) + npt.assert_array_equal( + verysimple_3vpacket_collection.last_interaction_type, + last_interaction_type, + ) + npt.assert_array_equal( + verysimple_3vpacket_collection.last_interaction_in_id, + last_interaction_in_id, + ) + npt.assert_array_equal( + verysimple_3vpacket_collection.last_interaction_out_id, + last_interaction_out_id, + ) + assert verysimple_3vpacket_collection.length == 9 diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py new file mode 100644 index 00000000000..f68311f9359 --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py @@ -0,0 +1,406 @@ +import os +import pytest +import numpy as np +import pandas as pd +import tardis.montecarlo.formal_integral as formal_integral +import tardis.montecarlo.montecarlo_numba.r_packet as r_packet +import tardis.montecarlo.montecarlo_configuration as mc +import tardis.montecarlo.montecarlo_numba.numba_interface as numba_interface +from tardis import constants as const +from tardis.montecarlo.montecarlo_numba.numba_interface import Estimators +import tardis.montecarlo.montecarlo_numba.numba_config as numba_config +from tardis.montecarlo.montecarlo_numba import macro_atom + +C_SPEED_OF_LIGHT = const.c.to("cm/s").value +SIGMA_THOMSON = const.sigma_T.to("cm^2").value + +from numpy.testing import ( + assert_equal, + assert_almost_equal, + assert_array_equal, + assert_allclose, +) + + +@pytest.fixture(scope="function") +def model(): + return numba_interface.NumbaModel( + r_inner=np.array([6.912e14, 8.64e14], dtype=np.float64), + r_outer=np.array([8.64e14, 1.0368e15], dtype=np.float64), + time_explosion=5.2e7, + ) + + +@pytest.fixture(scope="function") +def estimators(): + return numba_interface.Estimators( + j_estimator=np.array([0.0, 0.0], dtype=np.float64), + nu_bar_estimator=np.array([0.0, 0.0], dtype=np.float64), + j_blue_estimator=np.array( + [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64 + ), + Edotlu_estimator=np.array( + [[0.0, 0.0, 1.0], [0.0, 0.0, 1.0]], dtype=np.float64 + ), + ) + + +@pytest.mark.parametrize( + ["packet_params", "expected_params"], + [ + ({"mu": 0.3, "r": 7.5e14}, {"d_boundary": 259376919351035.88}), + ({"mu": -0.3, "r": 7.5e13}, {"d_boundary": -664987228972291.5}), + ({"mu": -0.3, "r": 7.5e14}, {"d_boundary": 709376919351035.9}), + ], +) +def test_calculate_distance_boundary(packet_params, expected_params, model): + mu = packet_params["mu"] + r = packet_params["r"] + + d_boundary = r_packet.calculate_distance_boundary( + r, mu, model.r_inner[0], model.r_outer[0] + ) + + # Accuracy to within 0.1cm + assert_almost_equal(d_boundary[0], expected_params["d_boundary"], decimal=1) + + +# +# +# TODO: split this into two tests - one to assert errors and other for d_line +@pytest.mark.parametrize( + ["packet_params", "expected_params"], + [ + ( + {"nu_line": 0.1, "next_line_id": 0, "is_last_line": True}, + {"tardis_error": None, "d_line": 1e99}, + ), + ( + {"nu_line": 0.2, "next_line_id": 1, "is_last_line": False}, + {"tardis_error": None, "d_line": 7.792353908000001e17}, + ), + ( + {"nu_line": 0.5, "next_line_id": 1, "is_last_line": False}, + {"tardis_error": r_packet.MonteCarloException, "d_line": 0.0}, + ), + ( + {"nu_line": 0.6, "next_line_id": 0, "is_last_line": False}, + {"tardis_error": r_packet.MonteCarloException, "d_line": 0.0}, + ), + ], +) +def test_calculate_distance_line( + packet_params, expected_params, static_packet, model +): + nu_line = packet_params["nu_line"] + is_last_line = packet_params["is_last_line"] + + time_explosion = model.time_explosion + + doppler_factor = r_packet.get_doppler_factor( + static_packet.r, static_packet.mu, time_explosion + ) + comov_nu = static_packet.nu * doppler_factor + + d_line = 0 + obtained_tardis_error = None + try: + d_line = r_packet.calculate_distance_line( + static_packet, comov_nu, is_last_line, nu_line, time_explosion + ) + except r_packet.MonteCarloException: + obtained_tardis_error = r_packet.MonteCarloException + + assert_almost_equal(d_line, expected_params["d_line"]) + assert obtained_tardis_error == expected_params["tardis_error"] + + +@pytest.mark.parametrize( + ["electron_density", "tau_event"], [(1e-5, 1.0), (1e10, 1e10)] +) +def test_calculate_distance_electron(electron_density, tau_event): + actual = r_packet.calculate_distance_electron(electron_density, tau_event) + expected = tau_event / (electron_density * SIGMA_THOMSON) + + assert_almost_equal(actual, expected) + + +@pytest.mark.parametrize( + ["electron_density", "distance"], + [(1e-5, 1.0), (1e10, 1e10), (-1, 0), (-1e10, -1e10)], +) +def test_calculate_tau_electron(electron_density, distance): + actual = r_packet.calculate_tau_electron(electron_density, distance) + expected = electron_density * SIGMA_THOMSON * distance + + assert_almost_equal(actual, expected) + + +@pytest.mark.parametrize( + ["mu", "r", "inv_t_exp", "expected"], + [ + (0.3, 7.5e14, 1 / 5.2e7, 0.9998556693818854), + (-0.3, 0, 1 / 2.6e7, 1.0), + (0, 1, 1 / 2.6e7, 1.0), + ], +) +def test_get_doppler_factor(mu, r, inv_t_exp, expected): + # Set the params from test cases here + # TODO: add relativity tests + time_explosion = 1 / inv_t_exp + + # Perform any other setups just before this, they can be additional calls + # to other methods or introduction of some temporary variables + + obtained = r_packet.get_doppler_factor(r, mu, time_explosion) + + # Perform required assertions + assert_almost_equal(obtained, expected) + + +@pytest.mark.parametrize( + ["mu", "r", "inv_t_exp", "expected"], + [ + (0.3, 7.5e14, 1 / 5.2e7, 1 / 0.9998556693818854), + (-0.3, 0, 1 / 2.6e7, 1.0), + (0, 1, 1 / 2.6e7, 1.0), + ], +) +def test_get_inverse_doppler_factor(mu, r, inv_t_exp, expected): + # Set the params from test cases here + # TODO: add relativity tests + time_explosion = 1 / inv_t_exp + + # Perform any other setups just before this, they can be additional calls + # to other methods or introduction of some temporary variables + + obtained = r_packet.get_inverse_doppler_factor(r, mu, time_explosion) + + # Perform required assertions + assert_almost_equal(obtained, expected) + + +def test_get_random_mu(set_seed_fixture): + """ + Ensure that different calls results + """ + set_seed_fixture(1963) + + output1 = r_packet.get_random_mu() + assert output1 == 0.9136407866175174 + + +@pytest.mark.parametrize( + [ + "cur_line_id", + "distance_trace", + "time_explosion", + "expected_j_blue", + "expected_Edotlu", + ], + [ + ( + 0, + 1e12, + 5.2e7, + [[2.249673812803061, 0.0, 0.0], [0.0, 0.0, 0.0]], + [[2.249673812803061 * 0.4, 0.0, 1.0], [0.0, 0.0, 1.0]], + ), + ( + 0, + 0, + 5.2e7, + [[2.249675256109242, 0.0, 0.0], [0.0, 0.0, 0.0]], + [ + [2.249675256109242 * 0.4, 0.0, 1.0], + [0.0, 0.0, 1.0], + ], + ), + ( + 1, + 1e5, + 1e10, + [[0.0, 0.0, 0.0], [2.249998311331767, 0.0, 0.0]], + [[0.0, 0.0, 1.0], [2.249998311331767 * 0.4, 0.0, 1.0]], + ), + ], +) +def test_update_line_estimators( + estimators, + static_packet, + cur_line_id, + distance_trace, + time_explosion, + expected_j_blue, + expected_Edotlu, +): + r_packet.update_line_estimators( + estimators, static_packet, cur_line_id, distance_trace, time_explosion + ) + + assert_allclose(estimators.j_blue_estimator, expected_j_blue) + assert_allclose(estimators.Edotlu_estimator, expected_Edotlu) + + +@pytest.mark.xfail(reason='Need to fix estimator differences across runs') +# TODO set RNG consistently +def test_trace_packet( + packet, + verysimple_numba_model, + verysimple_numba_plasma, + verysimple_estimators, + set_seed_fixture, +): + + set_seed_fixture(1963) + packet.initialize_line_id(verysimple_numba_plasma, verysimple_numba_model) + distance, interaction_type, delta_shell = r_packet.trace_packet( + packet, + verysimple_numba_model, + verysimple_numba_plasma, + verysimple_estimators, + ) + + assert delta_shell == 1 + assert interaction_type == 3 + assert_almost_equal(distance, 22978745222176.88) + + +@pytest.mark.xfail(reason="bug in full relativity") +@pytest.mark.parametrize("ENABLE_FULL_RELATIVITY", [True, False]) +@pytest.mark.parametrize( + ["packet_params", "expected_params"], + [ + ( + {"nu": 0.4, "mu": 0.3, "energy": 0.9, "r": 7.5e14}, + { + "mu": 0.3120599529139568, + "r": 753060422542573.9, + "j": 8998701024436.969, + "nubar": 3598960894542.354, + }, + ), + ( + {"nu": 0.6, "mu": -0.5, "energy": 0.5, "r": 8.1e14}, + { + "mu": -0.4906548373534084, + "r": 805046582503149.2, + "j": 5001298975563.031, + "nubar": 3001558973156.1387, + }, + ), + ], +) +def test_move_r_packet( + packet_params, + expected_params, + packet, + model, + estimators, + ENABLE_FULL_RELATIVITY, +): + distance = 1.0e13 + packet.nu = packet_params["nu"] + packet.mu = packet_params["mu"] + packet.energy = packet_params["energy"] + packet.r = packet_params["r"] + + numba_config.ENABLE_FULL_RELATIVITY = ENABLE_FULL_RELATIVITY + r_packet.move_r_packet.recompile() # This must be done as move_r_packet was jitted with ENABLE_FULL_RELATIVITY + doppler_factor = r_packet.get_doppler_factor( + packet.r, packet.mu, model.time_explosion + ) + + r_packet.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"]) + + expected_j = expected_params["j"] + expected_nubar = expected_params["nubar"] + + if ENABLE_FULL_RELATIVITY: + expected_j *= doppler_factor + expected_nubar *= doppler_factor + + numba_config.ENABLE_FULL_RELATIVITY = False + assert_allclose( + estimators.j_estimator[packet.current_shell_id], expected_j, rtol=5e-7 + ) + assert_allclose( + estimators.nu_bar_estimator[packet.current_shell_id], + expected_nubar, + rtol=5e-7, + ) + + +@pytest.mark.xfail(reason="To be implemented") +def test_set_estimators(): + pass + + +@pytest.mark.xfail(reason="To be implemented") +def test_set_estimators_full_relativity(): + pass + + +@pytest.mark.xfail(reason="To be implemented") +def test_line_emission(): + pass + + +@pytest.mark.parametrize( + ["current_shell_id", "delta_shell", "no_of_shells"], + [(132, 11, 132), (132, 1, 133), (132, 2, 133)], +) +def test_move_packet_across_shell_boundary_emitted( + packet, current_shell_id, delta_shell, no_of_shells +): + packet.current_shell_id = current_shell_id + r_packet.move_packet_across_shell_boundary( + packet, delta_shell, no_of_shells + ) + assert packet.status == r_packet.PacketStatus.EMITTED + + +@pytest.mark.parametrize( + ["current_shell_id", "delta_shell", "no_of_shells"], + [(132, -133, 132), (132, -133, 133), (132, -1e9, 133)], +) +def test_move_packet_across_shell_boundary_reabsorbed( + packet, current_shell_id, delta_shell, no_of_shells +): + packet.current_shell_id = current_shell_id + r_packet.move_packet_across_shell_boundary( + packet, delta_shell, no_of_shells + ) + assert packet.status == r_packet.PacketStatus.REABSORBED + + +@pytest.mark.parametrize( + ["current_shell_id", "delta_shell", "no_of_shells"], + [(132, -1, 199), (132, 0, 132), (132, 20, 154)], +) +def test_move_packet_across_shell_boundary_increment( + packet, current_shell_id, delta_shell, no_of_shells +): + packet.current_shell_id = current_shell_id + r_packet.move_packet_across_shell_boundary( + packet, delta_shell, no_of_shells + ) + assert packet.current_shell_id == current_shell_id + delta_shell + + +@pytest.mark.parametrize( + ["line_id", "nu_line", "expected"], + [(5495, 1629252823683562.5, True), (3000, 0, False)], +) +def test_test_for_close_line( + packet, line_id, nu_line, verysimple_numba_plasma, expected +): + + r_packet.test_for_close_line( + packet, line_id, nu_line, verysimple_numba_plasma + ) + + assert packet.is_close_line == expected diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_single_packet_loop.py b/tardis/montecarlo/montecarlo_numba/tests/test_single_packet_loop.py new file mode 100644 index 00000000000..2cf7f70f4ee --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/tests/test_single_packet_loop.py @@ -0,0 +1,53 @@ +import pytest + +import numpy.testing as npt + +from tardis.montecarlo.montecarlo_numba import RPacket +from tardis.montecarlo.montecarlo_numba.single_packet_loop import ( + single_packet_loop, +) + +@pytest.mark.xfail(reason='Need to fix estimator differences across runs') +# TODO set RNG consistently +def test_verysimple_single_packet_loop( + verysimple_numba_model, + verysimple_numba_plasma, + verysimple_estimators, + verysimple_vpacket_collection, + verysimple_packet_collection, +): + packet_collection = verysimple_packet_collection + vpacket_collection = verysimple_vpacket_collection + numba_model = verysimple_numba_model + numba_plasma = verysimple_numba_plasma + numba_estimators = verysimple_estimators + + i = 0 + r_packet = RPacket( + verysimple_numba_model.r_inner[0], + packet_collection.packets_input_mu[i], + packet_collection.packets_input_nu[i], + packet_collection.packets_input_energy[i], + i, + ) + single_packet_loop( + r_packet, + numba_model, + numba_plasma, + numba_estimators, + vpacket_collection, + ) + + npt.assert_almost_equal(r_packet.nu, 1053057938883272.8) + npt.assert_almost_equal(r_packet.mu, 0.9611146425440562) + npt.assert_almost_equal(r_packet.energy, 0.10327717505563379) + + +@pytest.mark.xfail(reason="To be implemented") +def test_set_packet_props_partial_relativity(): + assert False + + +@pytest.mark.xfail(reason="To be implemented") +def test_set_packet_props_full_relativity(): + assert False diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py b/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py new file mode 100644 index 00000000000..899a1fdd60c --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/tests/test_vpacket.py @@ -0,0 +1,132 @@ +import os +import pytest +import numpy as np +import pandas as pd +import tardis.montecarlo.formal_integral as formal_integral +import tardis.montecarlo.montecarlo_numba.r_packet as r_packet +import tardis.montecarlo.montecarlo_numba.vpacket as vpacket +import tardis.montecarlo.montecarlo_configuration as mc +import tardis.montecarlo.montecarlo_numba.numba_interface as numba_interface +from tardis import constants as const +from tardis.montecarlo.montecarlo_numba.numba_interface import Estimators +from tardis.montecarlo.montecarlo_numba import macro_atom + +C_SPEED_OF_LIGHT = const.c.to("cm/s").value + +import numpy.testing as npt + + +@pytest.fixture(scope="function") +def v_packet(): + return vpacket.VPacket( + r=7.5e14, + nu=4e15, + mu=0.3, + energy=0.9, + current_shell_id=0, + next_line_id=0, + index=0, + is_close_line=0, + ) + + +def v_packet_initialize_line_id(v_packet, numba_plasma, numba_model): + inverse_line_list_nu = numba_plasma.line_list_nu[::-1] + doppler_factor = r_packet.get_doppler_factor( + v_packet.r, v_packet.mu, numba_model.time_explosion + ) + comov_nu = v_packet.nu * doppler_factor + next_line_id = len(numba_plasma.line_list_nu) - np.searchsorted( + inverse_line_list_nu, comov_nu + ) + v_packet.next_line_id = next_line_id + + +def test_trace_vpacket_within_shell( + v_packet, verysimple_numba_model, verysimple_numba_plasma +): + # Give the vpacket a reasonable line ID + v_packet_initialize_line_id( + v_packet, verysimple_numba_plasma, verysimple_numba_model + ) + + ( + tau_trace_combined, + distance_boundary, + delta_shell, + ) = vpacket.trace_vpacket_within_shell( + v_packet, verysimple_numba_model, verysimple_numba_plasma + ) + + npt.assert_almost_equal(tau_trace_combined, 8164850.891288479) + npt.assert_almost_equal(distance_boundary, 843684056256104.1) + assert delta_shell == 1 + + +def test_trace_vpacket( + v_packet, verysimple_numba_model, verysimple_numba_plasma +): + # Set seed because of RNG in trace_vpacket + np.random.seed(1) + + # Give the vpacket a reasonable line ID + v_packet_initialize_line_id( + v_packet, verysimple_numba_plasma, verysimple_numba_model + ) + + tau_trace_combined = vpacket.trace_vpacket( + v_packet, verysimple_numba_model, verysimple_numba_plasma + ) + + npt.assert_almost_equal(tau_trace_combined, 8164850.891288479) + npt.assert_almost_equal(v_packet.r, 1286064000000000.0) + npt.assert_almost_equal(v_packet.nu, 4.0e15) + npt.assert_almost_equal(v_packet.energy, 0.0) + npt.assert_almost_equal(v_packet.mu, 0.8309726858508629) + assert v_packet.next_line_id == 2773 + assert v_packet.is_close_line == False + assert v_packet.current_shell_id == 1 + + +# NEEDS TO TEST VPACKET COLLECTION OVERFLOW +@pytest.mark.xfail(reason="Needs to be implemented") +def test_trace_vpacket_volley( + packet, + verysimple_packet_collection, + verysimple_3vpacket_collection, + verysimple_numba_model, + verysimple_numba_plasma, +): + # Set seed because of RNG in trace_vpacket + np.random.seed(1) + + packet.initialize_line_id(verysimple_numba_plasma, verysimple_numba_model) + + vpacket.trace_vpacket_volley( + packet, + verysimple_3vpacket_collection, + verysimple_numba_model, + verysimple_numba_plasma, + ) + + +@pytest.fixture(scope="function") +def broken_packet(): + return vpacket.VPacket( + r=1286064000000000.0, + nu=1660428912896553.2, + mu=0.4916053094346575, + energy=2.474533071386993e-07, + index=3, + is_close_line=True, + current_shell_id=0, + next_line_id=5495, + ) + + +def test_trace_bad_vpacket( + broken_packet, verysimple_numba_model, verysimple_numba_plasma +): + vpacket.trace_vpacket( + broken_packet, verysimple_numba_model, verysimple_numba_plasma + ) diff --git a/tardis/montecarlo/montecarlo_numba/vpacket.py b/tardis/montecarlo/montecarlo_numba/vpacket.py new file mode 100644 index 00000000000..ef304dd2354 --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/vpacket.py @@ -0,0 +1,293 @@ +from numba import float64, int64, boolean +from numba import njit, gdb +from numba.experimental import jitclass + +from tardis.montecarlo.montecarlo_numba import njit_dict +from tardis.montecarlo import ( + montecarlo_configuration as montecarlo_configuration, +) +import math +import numpy as np + +from tardis.montecarlo.montecarlo_numba.r_packet import ( + calculate_distance_boundary, + get_doppler_factor, + calculate_distance_line, + calculate_tau_electron, + PacketStatus, + move_packet_across_shell_boundary, + angle_aberration_LF_to_CMF, + angle_aberration_CMF_to_LF, + test_for_close_line, +) + +vpacket_spec = [ + ("r", float64), + ("mu", float64), + ("nu", float64), + ("energy", float64), + ("next_line_id", int64), + ("current_shell_id", int64), + ("status", int64), + ("index", int64), + ("is_close_line", boolean), +] + + +@jitclass(vpacket_spec) +class VPacket(object): + def __init__( + self, + r, + mu, + nu, + energy, + current_shell_id, + next_line_id, + index=0, + is_close_line=0, + ): + self.r = r + self.mu = mu + self.nu = nu + self.energy = energy + self.current_shell_id = current_shell_id + self.next_line_id = next_line_id + self.status = PacketStatus.IN_PROCESS + self.index = index + self.is_close_line = is_close_line + + +@njit(**njit_dict) +def trace_vpacket_within_shell(v_packet, numba_model, numba_plasma): + """ + Trace VPacket within one shell (relatively simple operation) + """ + r_inner = numba_model.r_inner[v_packet.current_shell_id] + r_outer = numba_model.r_outer[v_packet.current_shell_id] + + distance_boundary, delta_shell = calculate_distance_boundary( + v_packet.r, v_packet.mu, r_inner, r_outer + ) + # defining start for line interaction + start_line_id = v_packet.next_line_id + + # e scattering initialization + + cur_electron_density = numba_plasma.electron_density[ + v_packet.current_shell_id + ] + tau_electron = calculate_tau_electron( + cur_electron_density, distance_boundary + ) + tau_trace_combined = tau_electron + + # Calculating doppler factor + doppler_factor = get_doppler_factor( + v_packet.r, v_packet.mu, numba_model.time_explosion + ) + comov_nu = v_packet.nu * doppler_factor + cur_line_id = start_line_id + + for cur_line_id in range(start_line_id, len(numba_plasma.line_list_nu)): + # if tau_trace_combined > 10: ### FIXME ????? + # break + + nu_line = numba_plasma.line_list_nu[cur_line_id] + # TODO: Check if this is what the C code does + + tau_trace_line = numba_plasma.tau_sobolev[ + cur_line_id, v_packet.current_shell_id + ] + + if cur_line_id == len(numba_plasma.line_list_nu) - 1: + is_last_line = True + else: + is_last_line = False + + distance_trace_line = calculate_distance_line( + v_packet, + comov_nu, + is_last_line, + nu_line, + numba_model.time_explosion, + ) + + if cur_line_id != (len(numba_plasma.line_list_nu) - 1): + test_for_close_line( + v_packet, + cur_line_id, + numba_plasma.line_list_nu[cur_line_id - 1], + numba_plasma, + ) + + if distance_boundary <= distance_trace_line: + break + + tau_trace_combined += tau_trace_line + + else: + if cur_line_id == (len(numba_plasma.line_list_nu) - 1): + cur_line_id += 1 + v_packet.next_line_id = cur_line_id + + return tau_trace_combined, distance_boundary, delta_shell + + +@njit(**njit_dict) +def trace_vpacket(v_packet, numba_model, numba_plasma): + """ + Trace single vpacket. + Parameters + ---------- + v_packet + numba_model + numba_plasma + + Returns + ------- + + """ + + tau_trace_combined = 0.0 + while True: + ( + tau_trace_combined_shell, + distance_boundary, + delta_shell, + ) = trace_vpacket_within_shell(v_packet, numba_model, numba_plasma) + tau_trace_combined += tau_trace_combined_shell + + move_packet_across_shell_boundary( + v_packet, delta_shell, len(numba_model.r_inner) + ) + + if tau_trace_combined > montecarlo_configuration.tau_russian: + event_random = np.random.random() + if event_random > montecarlo_configuration.survival_probability: + v_packet.energy = 0.0 + v_packet.status = PacketStatus.EMITTED + else: + v_packet.energy = ( + v_packet.energy + / montecarlo_configuration.survival_probability + * math.exp(-tau_trace_combined) + ) + tau_trace_combined = 0.0 + + # Moving the v_packet + new_r = math.sqrt( + v_packet.r * v_packet.r + + distance_boundary * distance_boundary + + 2.0 * v_packet.r * distance_boundary * v_packet.mu + ) + v_packet.mu = (v_packet.mu * v_packet.r + distance_boundary) / new_r + v_packet.r = new_r + + if v_packet.status == PacketStatus.EMITTED: + break + return tau_trace_combined + + +@njit(**njit_dict) +def trace_vpacket_volley( + r_packet, vpacket_collection, numba_model, numba_plasma +): + """ + Shoot a volley of vpackets (the vpacket collection specifies how many) + from the current position of the rpacket. + + Parameters + ---------- + r_packet : [type] + [description] + vpacket_collection : [type] + [description] + numba_model : [type] + [description] + numba_plasma : [type] + [description] + """ + + if (r_packet.nu < vpacket_collection.v_packet_spawn_start_frequency) or ( + r_packet.nu > vpacket_collection.v_packet_spawn_end_frequency + ): + + return + + no_of_vpackets = vpacket_collection.number_of_vpackets + if no_of_vpackets == 0: + return + + ### TODO theoretical check for r_packet nu within vpackets bins - is done somewhere else I think + if r_packet.r > numba_model.r_inner[0]: # not on inner_boundary + r_inner_over_r = numba_model.r_inner[0] / r_packet.r + mu_min = -math.sqrt(1 - r_inner_over_r * r_inner_over_r) + v_packet_on_inner_boundary = False + if montecarlo_configuration.full_relativity: + mu_min = angle_aberration_LF_to_CMF( + r_packet, numba_model.time_explosion, mu_min + ) + else: + v_packet_on_inner_boundary = True + mu_min = 0.0 + + mu_bin = (1.0 - mu_min) / no_of_vpackets + r_packet_doppler_factor = get_doppler_factor( + r_packet.r, r_packet.mu, numba_model.time_explosion + ) + for i in range(no_of_vpackets): + v_packet_mu = mu_min + i * mu_bin + np.random.random() * mu_bin + + if v_packet_on_inner_boundary: # The weights are described in K&S 2014 + weight = 2 * v_packet_mu / no_of_vpackets + else: + weight = (1 - mu_min) / (2 * no_of_vpackets) + + # C code: next line, angle_aberration_CMF_to_LF( & virt_packet, storage); + if montecarlo_configuration.full_relativity: + v_packet_mu = angle_aberration_CMF_to_LF( + r_packet, numba_model.time_explosion, v_packet_mu + ) + v_packet_doppler_factor = get_doppler_factor( + r_packet.r, v_packet_mu, numba_model.time_explosion + ) + + # transform between r_packet mu and v_packet_mu + + doppler_factor_ratio = r_packet_doppler_factor / v_packet_doppler_factor + + v_packet_nu = r_packet.nu * doppler_factor_ratio + v_packet_energy = r_packet.energy * weight * doppler_factor_ratio + + v_packet = VPacket( + r_packet.r, + v_packet_mu, + v_packet_nu, + v_packet_energy, + r_packet.current_shell_id, + r_packet.next_line_id, + i, + r_packet.is_close_line, + ) + + if r_packet.next_line_id <= (len(numba_plasma.line_list_nu) - 1): + test_for_close_line( + v_packet, + r_packet.next_line_id, + numba_plasma.line_list_nu[r_packet.next_line_id - 1], + numba_plasma, + ) + + tau_vpacket = trace_vpacket(v_packet, numba_model, numba_plasma) + + v_packet.energy *= math.exp(-tau_vpacket) + + vpacket_collection.set_properties( + v_packet.nu, + v_packet.energy, + r_packet.nu, + r_packet.last_interaction_type, + r_packet.next_line_id, + r_packet.last_line_interaction_out_id, + ) diff --git a/tardis/montecarlo/packet_source.py b/tardis/montecarlo/packet_source.py index eaed2ddd5d4..30cad9db005 100644 --- a/tardis/montecarlo/packet_source.py +++ b/tardis/montecarlo/packet_source.py @@ -3,44 +3,49 @@ import numpy as np import numexpr as ne from tardis import constants as const +from tardis.montecarlo import ( + montecarlo_configuration as montecarlo_configuration, +) class BasePacketSource(abc.ABC): - def __init__(self, seed): self.seed = seed np.random.seed(seed) - + @abc.abstractmethod def create_packets(self, seed=None, **kwargs): pass @staticmethod - def create_zero_limb_darkening_packet_mus(no_of_packets): + def create_zero_limb_darkening_packet_mus(no_of_packets, rng): """ Create zero-limb-darkening packet :math:`\mu` distributed according to :math:`\\mu=\\sqrt{z}, z \isin [0, 1]` - + Parameters ---------- no_of_packets : int number of packets to be created """ - return np.sqrt(np.random.random(no_of_packets)) + # For testing purposes + if montecarlo_configuration.LEGACY_MODE_ENABLED: + return np.sqrt(np.random.random(no_of_packets)) + else: + return np.sqrt(rng.random(no_of_packets)) @staticmethod - def create_uniform_packet_energies(no_of_packets): + def create_uniform_packet_energies(no_of_packets, rng): """ - 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. - Parameters ---------- no_of_packets : int number of packets - + Returns ------- : numpy.ndarray @@ -48,33 +53,26 @@ def create_uniform_packet_energies(no_of_packets): """ return np.ones(no_of_packets) / no_of_packets - @staticmethod - def create_blackbody_packet_nus(T, no_of_packets, l_samples=1000): + def create_blackbody_packet_nus(T, no_of_packets, rng, l_samples=1000): """ - - Create packet :math:`\\nu` distributed using the algorithm described in - Bjorkman & Wood 2001 (page 4) which references + Create packet :math:`\\nu` distributed using the algorithm described in + Bjorkman & Wood 2001 (page 4) which references Carter & Cashwell 1975: - First, generate a uniform random number, :math:`\\xi_0 \\in [0, 1]` and - determine the minimum value of + determine the minimum value of :math:`l, l_{\\rm min}`, that satisfies the condition - .. math:: \\sum_{i=1}^{l} i^{-4} \\ge {{\\pi^4}\\over{90}} m_0 \\; . - - Next obtain four additional uniform random numbers (in the range 0 - to 1) :math:`\\xi_1, \\xi_2, \\xi_3, {\\rm and } \\xi_4`. - + + Next obtain four additional uniform random numbers (in the range 0 + to 1) :math:`\\xi_1, \\xi_2, \\xi_3, {\\rm and } \\xi_4`. + Finally, the packet frequency is given by - + .. math:: x = -\\ln{(\\xi_1\\xi_2\\xi_3\\xi_4)}/l_{\\rm min}\\; . - - where :math:`x=h\\nu/kT` - Parameters ---------- T : float @@ -82,34 +80,37 @@ def create_blackbody_packet_nus(T, no_of_packets, l_samples=1000): no_of_packets: int l_samples: int number of l_samples needed in the algorithm - Returns ------- - : numpy.ndarray array of frequencies """ 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_array = np.cumsum(np.arange(1, l_samples, dtype=np.float64) ** -4) + l_coef = np.pi ** 4 / 90.0 + + # For testing purposes + if montecarlo_configuration.LEGACY_MODE_ENABLED: + xis = np.random.random((5, no_of_packets)) + else: + xis = rng.random((5, no_of_packets)) - xis = np.random.random((5, no_of_packets)) - l = l_array.searchsorted(xis[0]*l_coef) + 1. + l = l_array.searchsorted(xis[0] * l_coef) + 1.0 xis_prod = np.prod(xis[1:], 0) - x = ne.evaluate('-log(xis_prod)/l') + x = ne.evaluate("-log(xis_prod)/l") return x * (const.k_B.cgs.value * T) / const.h.cgs.value class BlackBodySimpleSource(BasePacketSource): """ - Simple packet source that generates Blackbody packets for the Montecarlo + Simple packet source that generates Blackbody packets for the Montecarlo part. """ - def create_packets(self, T, no_of_packets): - nus = self.create_blackbody_packet_nus(T, no_of_packets) - mus = self.create_zero_limb_darkening_packet_mus(no_of_packets) - energies = self.create_uniform_packet_energies(no_of_packets) + def create_packets(self, T, no_of_packets, rng): + nus = self.create_blackbody_packet_nus(T, no_of_packets, rng) + mus = self.create_zero_limb_darkening_packet_mus(no_of_packets, rng) + energies = self.create_uniform_packet_energies(no_of_packets, rng) - return nus, mus, energies \ No newline at end of file + return nus, mus, energies diff --git a/tardis/montecarlo/setup_package.py b/tardis/montecarlo/setup_package.py index 6f0a2e58bc2..c9eadb0d165 100644 --- a/tardis/montecarlo/setup_package.py +++ b/tardis/montecarlo/setup_package.py @@ -1,4 +1,4 @@ -#setting the right include +# setting the right include from setuptools import Extension import os from astropy_helpers.distutils_helpers import get_distutils_option @@ -7,41 +7,73 @@ from glob import glob -if get_distutils_option('with_openmp', ['build', 'install', 'develop']) is not None: - compile_args = ['-fopenmp', '-W', '-Wall', '-Wmissing-prototypes', '-std=c99'] - link_args = ['-fopenmp'] - define_macros = [('WITHOPENMP', None)] +if ( + get_distutils_option("with_openmp", ["build", "install", "develop"]) + is not None +): + compile_args = [ + "-fopenmp", + "-W", + "-Wall", + "-Wmissing-prototypes", + "-std=c99", + ] + link_args = ["-fopenmp"] + define_macros = [("WITHOPENMP", None)] else: - compile_args = ['-W', '-Wall', '-Wmissing-prototypes', '-std=c99'] + compile_args = ["-W", "-Wall", "-Wmissing-prototypes", "-std=c99"] link_args = [] define_macros = [] -if get_distutils_option('with_vpacket_logging', ['build', 'install', 'develop']) is not None: - define_macros.append(('WITH_VPACKET_LOGGING', None)) +if ( + get_distutils_option( + "with_vpacket_logging", ["build", "install", "develop"] + ) + is not None +): + define_macros.append(("WITH_VPACKET_LOGGING", None)) + def get_extensions(): - sources = ['tardis/montecarlo/montecarlo.pyx'] - sources += [os.path.relpath(fname) for fname in glob( - os.path.join(os.path.dirname(__file__), 'src', '*.c'))] - sources += [os.path.relpath(fname) for fname in glob( - os.path.join(os.path.dirname(__file__), 'src/randomkit', '*.c'))] - deps = [os.path.relpath(fname) for fname in glob( - os.path.join(os.path.dirname(__file__), 'src', '*.h'))] - deps += [os.path.relpath(fname) for fname in glob( - os.path.join(os.path.dirname(__file__), 'src/randomkit', '*.h'))] + sources = ["tardis/montecarlo/montecarlo.pyx"] + sources += [ + os.path.relpath(fname) + for fname in glob(os.path.join(os.path.dirname(__file__), "src", "*.c")) + ] + sources += [ + os.path.relpath(fname) + for fname in glob( + os.path.join(os.path.dirname(__file__), "src/randomkit", "*.c") + ) + ] + deps = [ + os.path.relpath(fname) + for fname in glob(os.path.join(os.path.dirname(__file__), "src", "*.h")) + ] + deps += [ + os.path.relpath(fname) + for fname in glob( + os.path.join(os.path.dirname(__file__), "src/randomkit", "*.h") + ) + ] return cythonize( - Extension('tardis.montecarlo.montecarlo', sources, - include_dirs=['tardis/montecarlo/src', - 'tardis/montecarlo/src/randomkit', - 'numpy'], - depends=deps, - extra_compile_args=compile_args, - extra_link_args=link_args, - define_macros=define_macros) - ) + Extension( + "tardis.montecarlo.montecarlo", + sources, + include_dirs=[ + "tardis/montecarlo/src", + "tardis/montecarlo/src/randomkit", + "numpy", + ], + depends=deps, + extra_compile_args=compile_args, + extra_link_args=link_args, + define_macros=define_macros, + ) + ) def get_package_data(): - return {'tardis.montecarlo.tests':['data/*.npy', 'data/*.hdf']} + return {"tardis.montecarlo.tests": ["data/*.npy", "data/*.hdf"]} diff --git a/tardis/montecarlo/spectrum.py b/tardis/montecarlo/spectrum.py index 0e8f62aa9d9..4f5ef9e4108 100644 --- a/tardis/montecarlo/spectrum.py +++ b/tardis/montecarlo/spectrum.py @@ -3,6 +3,7 @@ from astropy import units as u from tardis.io.util import HDFWriterMixin + class TARDISSpectrum(HDFWriterMixin): """ TARDISSpectrum(_frequency, luminosity) @@ -16,113 +17,122 @@ class TARDISSpectrum(HDFWriterMixin): After manually adding a distance attribute, the properties 'flux_nu' and 'flux_lambda' become available """ - hdf_properties = ['_frequency', 'luminosity'] + + hdf_properties = ["_frequency", "luminosity"] + def __init__(self, _frequency, luminosity): # Check for correct inputs if not _frequency.shape[0] == luminosity.shape[0] + 1: raise ValueError( - "shape of '_frequency' and 'luminosity' are not compatible" - ": '{}' and '{}'".format( - _frequency.shape[0], - luminosity.shape[0]) - ) - self._frequency = _frequency.to('Hz', u.spectral()) - self.luminosity = luminosity.to('erg / s') + "shape of '_frequency' and 'luminosity' are not compatible" + ": '{}' and '{}'".format( + _frequency.shape[0], luminosity.shape[0] + ) + ) + self._frequency = _frequency.to("Hz", u.spectral()) + self.luminosity = luminosity.to("erg / s") self.frequency = self._frequency[:-1] self.delta_frequency = self._frequency[1] - self._frequency[0] - self.wavelength = self.frequency.to('angstrom', u.spectral()) + self.wavelength = self.frequency.to("angstrom", u.spectral()) self.luminosity_density_nu = ( - self.luminosity / self.delta_frequency).to('erg / (s Hz)') + self.luminosity / self.delta_frequency + ).to("erg / (s Hz)") self.luminosity_density_lambda = self.f_nu_to_f_lambda( - self.luminosity_density_nu, - ) + self.luminosity_density_nu, + ) @property def flux_nu(self): - warnings.simplefilter('always', DeprecationWarning) + warnings.simplefilter("always", DeprecationWarning) warnings.warn( - "TARDISSpectrum.flux_nu is deprecated, " - "please use TARDISSpectrum.luminosity_to_flux() in the " - "future.", - category=DeprecationWarning, stacklevel=2) - warnings.simplefilter('default', DeprecationWarning) + "TARDISSpectrum.flux_nu is deprecated, " + "please use TARDISSpectrum.luminosity_to_flux() in the " + "future.", + category=DeprecationWarning, + stacklevel=2, + ) + warnings.simplefilter("default", DeprecationWarning) try: return self.luminosity_to_flux( - self.luminosity_density_nu, - self.distance) + self.luminosity_density_nu, self.distance + ) except AttributeError: - raise AttributeError( - 'distance is required as attribute of' - '{} to calculate "{}"'.format( - self.__class__.__name__, - 'flux_nu' - ) - ) + raise AttributeError( + "distance is required as attribute of" + '{} to calculate "{}"'.format( + self.__class__.__name__, "flux_nu" + ) + ) @property def flux_lambda(self): - warnings.simplefilter('always', DeprecationWarning) + warnings.simplefilter("always", DeprecationWarning) warnings.warn( - "TARDISSpectrum.flux_lambda is deprecated, " - "please use TARDISSpectrum.luminosity_to_flux() in the " - "future.", - category=DeprecationWarning, stacklevel=2) - warnings.simplefilter('default', DeprecationWarning) + "TARDISSpectrum.flux_lambda is deprecated, " + "please use TARDISSpectrum.luminosity_to_flux() in the " + "future.", + category=DeprecationWarning, + stacklevel=2, + ) + warnings.simplefilter("default", DeprecationWarning) try: return self.luminosity_to_flux( - self.luminosity_density_lambda, - self.distance - ) + self.luminosity_density_lambda, self.distance + ) except AttributeError: - raise AttributeError( - 'distance is required as attribute of' - '{} to calculate "{}"'.format( - self.__class__.__name__, - 'flux_lambda' - ) - ) + raise AttributeError( + "distance is required as attribute of" + '{} to calculate "{}"'.format( + self.__class__.__name__, "flux_lambda" + ) + ) @staticmethod def luminosity_to_flux(luminosity, distance): - return luminosity / (4 * np.pi * distance.to('cm')**2) + return luminosity / (4 * np.pi * distance.to("cm") ** 2) def f_nu_to_f_lambda(self, f_nu): return f_nu * self.frequency / self.wavelength - def plot(self, ax=None, mode='wavelength'): + def plot(self, ax=None, mode="wavelength"): if ax is None: from matplotlib.pyplot import gca + ax = gca() - if mode == 'wavelength': - ax.plot( - self.wavelength.value, - self.luminosity_density_lambda.value) - ax.set_xlabel('Wavelength [{}]'.format( - self.wavelength.unit._repr_latex_()) - ) + if mode == "wavelength": + ax.plot(self.wavelength.value, self.luminosity_density_lambda.value) + ax.set_xlabel( + "Wavelength [{}]".format(self.wavelength.unit._repr_latex_()) + ) ax.set_ylabel( - 'Flux [{:s}]'.format( - self.luminosity_density_lambda.unit._repr_latex_()) - ) + "Flux [{:s}]".format( + self.luminosity_density_lambda.unit._repr_latex_() + ) + ) else: warnings.warn( - 'Did not find plotting mode {}, doing nothing.'.format( - mode)) + "Did not find plotting mode {}, doing nothing.".format(mode) + ) - def to_ascii(self, fname, mode='luminosity_density'): - if mode == 'luminosity_density': + def to_ascii(self, fname, mode="luminosity_density"): + if mode == "luminosity_density": np.savetxt( - fname, list(zip( + fname, + list( + zip( self.wavelength.value, - self.luminosity_density_lambda.value))) - elif mode == 'flux': + self.luminosity_density_lambda.value, + ) + ), + ) + elif mode == "flux": np.savetxt( - fname, - zip(self.wavelength.value, self.flux_lambda.value)) + fname, zip(self.wavelength.value, self.flux_lambda.value) + ) else: raise NotImplementedError( - 'only mode "luminosity_density"' - 'and "flux" are implemented') + 'only mode "luminosity_density"' 'and "flux" are implemented' + ) diff --git a/tardis/montecarlo/struct.py b/tardis/montecarlo/struct.py index 80f4b95c015..fb6953cb4d7 100644 --- a/tardis/montecarlo/struct.py +++ b/tardis/montecarlo/struct.py @@ -8,130 +8,130 @@ class RPacket(Structure): _fields_ = [ - ('nu', c_double), - ('mu', c_double), - ('energy', c_double), - ('r', c_double), - ('tau_event', c_double), - ('nu_line', c_double), - ('current_shell_id', c_int64), - ('next_line_id', c_int64), - ('last_line', c_int64), - ('close_line', c_int64), - ('current_continuum_id', c_int64), - ('virtual_packet_flag', c_int64), - ('virtual_packet', c_int64), - ('d_line', c_double), - ('d_electron', c_double), - ('d_boundary', c_double), - ('d_cont', c_double), - ('next_shell_id', c_int64), - ('status', c_rpacket_status_t), - ('id', c_int64), - ('chi_th', c_double), - ('chi_cont', c_double), - ('chi_ff', c_double), - ('chi_bf', c_double), - ('chi_bf_tmp_partial', POINTER(c_double)), - ('macro_atom_activation_level', c_int64), - ('compute_chi_bf', c_bool), - ('vpacket_weight', c_double) + ("nu", c_double), + ("mu", c_double), + ("energy", c_double), + ("r", c_double), + ("tau_event", c_double), + ("nu_line", c_double), + ("current_shell_id", c_int64), + ("next_line_id", c_int64), + ("last_line", c_int64), + ("close_line", c_int64), + ("current_continuum_id", c_int64), + ("virtual_packet_flag", c_int64), + ("virtual_packet", c_int64), + ("d_line", c_double), + ("d_electron", c_double), + ("d_boundary", c_double), + ("d_cont", c_double), + ("next_shell_id", c_int64), + ("status", c_rpacket_status_t), + ("id", c_int64), + ("chi_th", c_double), + ("chi_cont", c_double), + ("chi_ff", c_double), + ("chi_bf", c_double), + ("chi_bf_tmp_partial", POINTER(c_double)), + ("macro_atom_activation_level", c_int64), + ("compute_chi_bf", c_bool), + ("vpacket_weight", c_double), ] class PhotoXsect1level(Structure): _fields_ = [ - ('nu', POINTER(c_double)), - ('x_sect', POINTER(c_double)), - ('no_of_points', c_int64) + ("nu", POINTER(c_double)), + ("x_sect", POINTER(c_double)), + ("no_of_points", c_int64), ] class StorageModel(Structure): _fields_ = [ - ('packet_nus', POINTER(c_double)), - ('packet_mus', POINTER(c_double)), - ('packet_energies', POINTER(c_double)), - ('output_nus', POINTER(c_double)), - ('output_energies', POINTER(c_double)), - ('last_interaction_in_nu', POINTER(c_double)), - ('last_line_interaction_in_id', POINTER(c_int64)), - ('last_line_interaction_out_id', POINTER(c_int64)), - ('last_line_interaction_shell_id', POINTER(c_int64)), - ('last_interaction_type', POINTER(c_int64)), - ('last_interaction_out_type', POINTER(c_int64)), - ('no_of_packets', c_int64), - ('no_of_shells', c_int64), - ('no_of_shells_i', c_int64), - ('r_inner', POINTER(c_double)), - ('r_outer', POINTER(c_double)), - ('r_inner_i', POINTER(c_double)), - ('r_outer_i', POINTER(c_double)), - ('v_inner', POINTER(c_double)), - ('time_explosion', c_double), - ('inverse_time_explosion', c_double), - ('electron_densities', POINTER(c_double)), - ('electron_densities_i', POINTER(c_double)), - ('inverse_electron_densities', POINTER(c_double)), - ('line_list_nu', POINTER(c_double)), - ('continuum_list_nu', POINTER(c_double)), - ('line_lists_tau_sobolevs', POINTER(c_double)), - ('line_lists_tau_sobolevs_i', POINTER(c_double)), - ('line_lists_tau_sobolevs_nd', c_int64), - ('line_lists_j_blues', POINTER(c_double)), - ('line_lists_j_blues_nd', c_int64), - ('line_lists_Edotlu', POINTER(c_double)), - ('no_of_lines', c_int64), - ('no_of_edges', c_int64), - ('line_interaction_id', c_int64), - ('transition_probabilities', POINTER(c_double)), - ('transition_probabilities_nd', c_int64), - ('line2macro_level_upper', POINTER(c_int64)), - ('macro_block_references', POINTER(c_int64)), - ('transition_type', POINTER(c_int64)), - ('destination_level_id', POINTER(c_int64)), - ('transition_line_id', POINTER(c_int64)), - ('js', POINTER(c_double)), - ('nubars', POINTER(c_double)), - ('spectrum_start_nu', c_double), - ('spectrum_delta_nu', c_double), - ('spectrum_end_nu', c_double), - ('spectrum_virt_start_nu', c_double), - ('spectrum_virt_end_nu', c_double), - ('spectrum_virt_nu', POINTER(c_double)), - ('sigma_thomson', c_double), - ('inverse_sigma_thomson', c_double), - ('inner_boundary_albedo', c_double), - ('reflective_inner_boundary', c_int64), - ('current_packet_id', c_int64), - ('photo_xsect', POINTER(POINTER(PhotoXsect1level))), - ('chi_ff_factor', POINTER(c_double)), - ('t_electrons', POINTER(c_double)), - ('l_pop', POINTER(c_double)), - ('l_pop_r', POINTER(c_double)), - ('cont_status', c_cont_status_t), - ('bf_treatment', c_int), - ('virt_packet_nus', POINTER(c_double)), - ('virt_packet_energies', POINTER(c_double)), - ('virt_packet_last_interaction_in_nu', POINTER(c_double)), - ('virt_packet_last_interaction_type', POINTER(c_int64)), - ('virt_packet_last_line_interaction_in_id', POINTER(c_int64)), - ('virt_packet_last_line_interaction_out_id', POINTER(c_int64)), - ('virt_packet_count', c_int64), - ('virt_array_size', c_int64), - ('kpacket2macro_level', c_int64), - ('cont_edge2macro_level', POINTER(c_int64)), - ('photo_ion_estimator', POINTER(c_double)), - ('stim_recomb_estimator', POINTER(c_double)), - ('photo_ion_estimator_statistics', POINTER(c_int64)), - ('bf_heating_estimator', POINTER(c_double)), - ('ff_heating_estimator', POINTER(c_double)), - ('stim_recomb_cooling_estimator', POINTER(c_double)), - ('full_relativity', c_int), - ('survival_probability',c_double), - ('tau_russian', c_double), - ('tau_bias', POINTER(c_double)), - ('enable_biasing', c_int) + ("packet_nus", POINTER(c_double)), + ("packet_mus", POINTER(c_double)), + ("packet_energies", POINTER(c_double)), + ("output_nus", POINTER(c_double)), + ("output_energies", POINTER(c_double)), + ("last_interaction_in_nu", POINTER(c_double)), + ("last_line_interaction_in_id", POINTER(c_int64)), + ("last_line_interaction_out_id", POINTER(c_int64)), + ("last_line_interaction_shell_id", POINTER(c_int64)), + ("last_interaction_type", POINTER(c_int64)), + ("last_interaction_out_type", POINTER(c_int64)), + ("no_of_packets", c_int64), + ("no_of_shells", c_int64), + ("no_of_shells_i", c_int64), + ("r_inner", POINTER(c_double)), + ("r_outer", POINTER(c_double)), + ("r_inner_i", POINTER(c_double)), + ("r_outer_i", POINTER(c_double)), + ("v_inner", POINTER(c_double)), + ("time_explosion", c_double), + ("inverse_time_explosion", c_double), + ("electron_densities", POINTER(c_double)), + ("electron_densities_i", POINTER(c_double)), + ("inverse_electron_densities", POINTER(c_double)), + ("line_list_nu", POINTER(c_double)), + ("continuum_list_nu", POINTER(c_double)), + ("line_lists_tau_sobolevs", POINTER(c_double)), + ("line_lists_tau_sobolevs_i", POINTER(c_double)), + ("line_lists_tau_sobolevs_nd", c_int64), + ("line_lists_j_blues", POINTER(c_double)), + ("line_lists_j_blues_nd", c_int64), + ("line_lists_Edotlu", POINTER(c_double)), + ("no_of_lines", c_int64), + ("no_of_edges", c_int64), + ("line_interaction_id", c_int64), + ("transition_probabilities", POINTER(c_double)), + ("transition_probabilities_nd", c_int64), + ("line2macro_level_upper", POINTER(c_int64)), + ("macro_block_references", POINTER(c_int64)), + ("transition_type", POINTER(c_int64)), + ("destination_level_id", POINTER(c_int64)), + ("transition_line_id", POINTER(c_int64)), + ("js", POINTER(c_double)), + ("nubars", POINTER(c_double)), + ("spectrum_start_nu", c_double), + ("spectrum_delta_nu", c_double), + ("spectrum_end_nu", c_double), + ("spectrum_virt_start_nu", c_double), + ("spectrum_virt_end_nu", c_double), + ("spectrum_virt_nu", POINTER(c_double)), + ("sigma_thomson", c_double), + ("inverse_sigma_thomson", c_double), + ("inner_boundary_albedo", c_double), + ("reflective_inner_boundary", c_int64), + ("current_packet_id", c_int64), + ("photo_xsect", POINTER(POINTER(PhotoXsect1level))), + ("chi_ff_factor", POINTER(c_double)), + ("t_electrons", POINTER(c_double)), + ("l_pop", POINTER(c_double)), + ("l_pop_r", POINTER(c_double)), + ("cont_status", c_cont_status_t), + ("bf_treatment", c_int), + ("virt_packet_nus", POINTER(c_double)), + ("virt_packet_energies", POINTER(c_double)), + ("virt_packet_last_interaction_in_nu", POINTER(c_double)), + ("virt_packet_last_interaction_type", POINTER(c_int64)), + ("virt_packet_last_line_interaction_in_id", POINTER(c_int64)), + ("virt_packet_last_line_interaction_out_id", POINTER(c_int64)), + ("virt_packet_count", c_int64), + ("virt_array_size", c_int64), + ("kpacket2macro_level", c_int64), + ("cont_edge2macro_level", POINTER(c_int64)), + ("photo_ion_estimator", POINTER(c_double)), + ("stim_recomb_estimator", POINTER(c_double)), + ("photo_ion_estimator_statistics", POINTER(c_int64)), + ("bf_heating_estimator", POINTER(c_double)), + ("ff_heating_estimator", POINTER(c_double)), + ("stim_recomb_cooling_estimator", POINTER(c_double)), + ("full_relativity", c_int), + ("survival_probability", c_double), + ("tau_russian", c_double), + ("tau_bias", POINTER(c_double)), + ("enable_biasing", c_int), ] @@ -159,6 +159,7 @@ class BoundFreeTreatment(Enum): LIN_INTERPOLATION = 0 HYDROGENIC = 1 + # Variables corresponding to macros defined in rpacket.h . MISS_DISTANCE = 1e99 C = 29979245800.0 @@ -172,8 +173,8 @@ class BoundFreeTreatment(Enum): class RKState(Structure): _fields_ = [ - ('key', c_ulong * RK_STATE_LEN), - ('pos', c_int), - ('has_gauss', c_int), - ('gauss', c_double) + ("key", c_ulong * RK_STATE_LEN), + ("pos", c_int), + ("has_gauss", c_int), + ("gauss", c_double), ] diff --git a/tardis/montecarlo/tests/conftest.py b/tardis/montecarlo/tests/conftest.py index 13cca00085e..77e1532c6f0 100644 --- a/tardis/montecarlo/tests/conftest.py +++ b/tardis/montecarlo/tests/conftest.py @@ -1,27 +1,25 @@ import os import pytest +from tardis.io import config_reader + from ctypes import ( - CDLL, - byref, - c_int64, - c_double, - c_ulong, - ) + CDLL, + byref, + c_int64, + c_double, + c_ulong, +) -from tardis.montecarlo import montecarlo from tardis.montecarlo.struct import ( - RPacket, StorageModel, RKState, + RPacket, + StorageModel, + RKState, TARDIS_PACKET_STATUS_IN_PROCESS, CONTINUUM_OFF, - BoundFreeTreatment + BoundFreeTreatment, ) - - -# Wrap the shared object containing C methods, which are tested here. -@pytest.fixture(scope='session') -def clib(): - return CDLL(os.path.join(montecarlo.__file__)) +from tardis.montecarlo.base import MontecarloRunner @pytest.fixture(scope="function") @@ -47,7 +45,7 @@ def packet(): chi_cont=6.652486e-16, chi_bf_tmp_partial=(c_double * 2)(), compute_chi_bf=True, - vpacket_weight=1.0 + vpacket_weight=1.0, ) @@ -61,71 +59,56 @@ def model(): last_line_interaction_shell_id=(c_int64 * 2)(*([0] * 2)), last_interaction_type=(c_int64 * 2)(*([2])), last_interaction_out_type=(c_int64 * 1)(*([0])), - no_of_shells=2, - r_inner=(c_double * 2)(*[6.912e14, 8.64e14]), r_outer=(c_double * 2)(*[8.64e14, 1.0368e15]), - time_explosion=5.2e7, inverse_time_explosion=1 / 5.2e7, - electron_densities=(c_double * 2)(*[1.0e9] * 2), inverse_electron_densities=(c_double * 2)(*[1.0e-9] * 2), - - line_list_nu=(c_double * 5)(*[ - 1.26318289e+16, - 1.26318289e+16, - 1.23357675e+16, - 1.23357675e+16, - 1.16961598e+16]), - - continuum_list_nu=(c_double * 2)(*([1.e13] * 2)), - - line_lists_tau_sobolevs=(c_double * 1000)(*([1.e-5] * 1000)), - line_lists_j_blues=(c_double * 2)(*([1.e-10] * 2)), + line_list_nu=(c_double * 5)( + *[ + 1.26318289e16, + 1.26318289e16, + 1.23357675e16, + 1.23357675e16, + 1.16961598e16, + ] + ), + continuum_list_nu=(c_double * 2)(*([1.0e13] * 2)), + line_lists_tau_sobolevs=(c_double * 1000)(*([1.0e-5] * 1000)), + line_lists_j_blues=(c_double * 2)(*([1.0e-10] * 2)), line_lists_j_blues_nd=0, - # Init to an explicit array line_lists_Edotlu=(c_double * 3)(*[0.0, 0.0, 1.0]), - no_of_lines=5, no_of_edges=2, - line_interaction_id=0, line2macro_level_upper=(c_int64 * 2)(*([0] * 2)), - js=(c_double * 2)(*([0.0] * 2)), nubars=(c_double * 2)(*([0.0] * 2)), - - spectrum_start_nu=1.e14, + spectrum_start_nu=1.0e14, spectrum_delta_nu=293796608840.0, - spectrum_end_nu=6.e15, - + spectrum_end_nu=6.0e15, spectrum_virt_start_nu=1e14, spectrum_virt_end_nu=6e15, spectrum_virt_nu=(c_double * 20000)(*([0.0] * 20000)), - sigma_thomson=6.652486e-25, inverse_sigma_thomson=1 / 6.652486e-25, - inner_boundary_albedo=0.0, reflective_inner_boundary=0, - chi_ff_factor=(c_double * 2)(*([1.0] * 2)), t_electrons=(c_double * 2)(*([1.0e4] * 2)), - l_pop=(c_double * 20000)(*([2.0] * 20000)), l_pop_r=(c_double * 20000)(*([3.0] * 20000)), cont_status=CONTINUUM_OFF, bf_treatment=BoundFreeTreatment.LIN_INTERPOLATION.value, ff_heating_estimator=(c_double * 2)(*([0.0] * 2)), cont_edge2macro_level=(c_int64 * 6)(*([1] * 6)), - survival_probability=0.0, tau_russian=10.0, tau_bias=(c_double * 3)(*([5.0, 0.5, 0.0])), - enable_biasing=0 + enable_biasing=0, ) @@ -133,10 +116,7 @@ def model(): def mt_state(): """Fixture to return `RKState` object with default params initialized.""" return RKState( - key=(c_ulong * 624)(*([0] * 624)), - pos=0, - has_gauss=0, - gauss=0.0 + key=(c_ulong * 624)(*([0] * 624)), pos=0, has_gauss=0, gauss=0.0 ) @@ -145,3 +125,11 @@ def mt_state_seeded(clib, mt_state): seed = 23111963 clib.rk_seed(seed, byref(mt_state)) return mt_state + + +@pytest.fixture(scope="function") +def runner(): + config_fname = "tardis/io/tests/data/tardis_configv1_verysimply.yml" + config = config_reader.Configuration.from_yaml(config_fname) + runner = MontecarloRunner.from_config(config) + return runner diff --git a/tardis/montecarlo/tests/test_base.py b/tardis/montecarlo/tests/test_base.py index c31b5a27209..283eb45bb6c 100644 --- a/tardis/montecarlo/tests/test_base.py +++ b/tardis/montecarlo/tests/test_base.py @@ -9,24 +9,32 @@ # Save and Load ### + @pytest.fixture(scope="module", autouse=True) def to_hdf_buffer(hdf_file_path, simulation_verysimple): - simulation_verysimple.runner.to_hdf(hdf_file_path, name='runner') + simulation_verysimple.runner.to_hdf(hdf_file_path, name="runner") + + +runner_properties = [ + "output_nu", + "output_energy", + "nu_bar_estimator", + "j_estimator", + "montecarlo_virtual_luminosity", + "last_interaction_in_nu", + "last_interaction_type", + "last_line_interaction_in_id", + "last_line_interaction_out_id", + "last_line_interaction_shell_id", + "packet_luminosity", +] -runner_properties = ['output_nu', 'output_energy', 'nu_bar_estimator', - 'j_estimator', 'montecarlo_virtual_luminosity', - 'last_interaction_in_nu', - 'last_interaction_type', - 'last_line_interaction_in_id', - 'last_line_interaction_out_id', - 'last_line_interaction_shell_id', - 'packet_luminosity'] @pytest.mark.parametrize("attr", runner_properties) def test_hdf_runner(hdf_file_path, simulation_verysimple, attr): actual = getattr(simulation_verysimple.runner, attr) - if hasattr(actual, 'cgs'): + if hasattr(actual, "cgs"): actual = actual.cgs.value - path = os.path.join('runner', attr) + path = os.path.join("runner", attr) expected = pd.read_hdf(hdf_file_path, path) assert_almost_equal(actual, expected.values) diff --git a/tardis/montecarlo/tests/test_cmontecarlo.py b/tardis/montecarlo/tests/test_cmontecarlo.py deleted file mode 100644 index ff7519e251a..00000000000 --- a/tardis/montecarlo/tests/test_cmontecarlo.py +++ /dev/null @@ -1,1058 +0,0 @@ -""" -Unit tests for methods in `tardis/montecarlo/src/cmontecarlo.c`. -* `ctypes` library is used to wrap C methods and expose them to python. - - -Probable Reasons for Failing Tests: ------------------------------------ - -1. Change made in C struct declarations: - - Reflect the changes done in C structs, into Python counterparts. - - Check **tardis/montecarlo/struct.py**. - -2. Return type of any method changed: - - Modify the `restype` parameter in the test method here. - - For example: - ``` - clib.rpacket_doppler_factor.restype = c_double - ``` - -3. Underlying logic modified: - - Check whether the changes made in C method are logically correct. - - If the changes made were correct and necessary, update the corresponding - test case. - - -General Test Design Procedure: ------------------------------- - -Please follow this design procedure while adding a new test: - -1. Parametrization as per desire of code coverage. - - C tests have different flows controlled by conditional statements. - Parameters checked in conditions can be provided in different testcases. - - Keep consistency with variable names as (in order): - - `packet_params` - - `model_params` - - `expected_params` (`expected` if only one value to be asserted.) - - Suggested variable names can be compromised if readability of the test - increases. - -2. Test Method body: - - Keep name as `test_` + `(name of C method)`. - - Refer to method `test_rpacket_doppler_factor` below for description. -""" - - -import os -import pytest -import numpy as np -import pandas as pd - - -from ctypes import ( - CDLL, - byref, - c_uint, - c_int64, - c_double, - c_ulong, - c_void_p, - cast, - POINTER, - pointer, - CFUNCTYPE - ) - -from numpy.testing import ( - assert_equal, - assert_almost_equal, - assert_array_equal, - assert_allclose - ) - -from tardis import __path__ as path -from tardis.montecarlo.struct import ( - RPacket, StorageModel, RKState, - PhotoXsect1level, - TARDIS_ERROR_OK, - TARDIS_ERROR_BOUNDS_ERROR, - TARDIS_ERROR_COMOV_NU_LESS_THAN_NU_LINE, - TARDIS_PACKET_STATUS_IN_PROCESS, - TARDIS_PACKET_STATUS_EMITTED, - TARDIS_PACKET_STATUS_REABSORBED, - CONTINUUM_OFF, - CONTINUUM_ON, - INVERSE_C, - BoundFreeTreatment -) - - -@pytest.fixture(scope='module') -def continuum_compare_data_fname(): - fname = 'continuum_compare_data.hdf' - return os.path.join(path[0], 'montecarlo', 'tests', 'data', fname) - - -@pytest.fixture(scope='module') -def continuum_compare_data(continuum_compare_data_fname, request): - compare_data = pd.HDFStore(continuum_compare_data_fname, mode='r') - - def fin(): - compare_data.close() - request.addfinalizer(fin) - - return compare_data - - -@pytest.fixture(scope="function") -def expected_ff_emissivity(continuum_compare_data): - emissivities = continuum_compare_data['ff_emissivity'] - - def ff_emissivity(t_electron): - emissivity = emissivities[t_electron] - nu_bins = emissivity['nu_bins'].values - emissivity_value = emissivity['emissivity'].dropna().values - - return nu_bins, emissivity_value - - return ff_emissivity - - -@pytest.fixture(scope='module') -def get_rkstate(continuum_compare_data): - data = continuum_compare_data['z2rkstate_key'] - pos_data = continuum_compare_data['z2rkstate_pos'] - - def z2rkstate(z_random): - key = (c_ulong * 624)(*data.loc[z_random].values) - pos = pos_data.loc[z_random] - return RKState( - key=key, - pos=pos, - has_gauss=0, - gauss=0.0 - ) - - return z2rkstate - - -@pytest.fixture(scope='function') -def model_w_edges(ion_edges, model): - photo_xsect = (POINTER(PhotoXsect1level) * len(ion_edges))() - - for i, edge in enumerate(ion_edges): - x_sect_1level = PhotoXsect1level() - for key, value in edge.items(): - if key in ['nu', 'x_sect']: - value = (c_double * len(value))(*value) - setattr(x_sect_1level, key, value) - photo_xsect[i] = pointer(x_sect_1level) - - no_of_edges = len(ion_edges) - continuum_list_nu = (c_double * no_of_edges)(*[edge['nu'][0] for edge in ion_edges]) - - model.photo_xsect = photo_xsect - model.continuum_list_nu = continuum_list_nu - model.no_of_edges = no_of_edges - - estimator_size = model.no_of_shells * no_of_edges - estims = ['photo_ion_estimator', 'stim_recomb_estimator', - 'bf_heating_estimator', 'stim_recomb_cooling_estimator' - ] - for estimator in estims: - setattr(model, estimator, (c_double * estimator_size)(*[0] * estimator_size)) - - model.photo_ion_estimator_statistics = (c_int64 * estimator_size)(*[0] * estimator_size) - return model - - -@pytest.fixture(scope='module') -def ion_edges(): - return [ - {'nu': [4.0e14, 4.1e14, 4.2e14, 4.3e14], - 'x_sect': [1.0, 0.9, 0.8, 0.7], 'no_of_points': 4}, - {'nu': [3.0e14, 3.1e14, 3.2e14, 3.3e14, 3.4e14], - 'x_sect': [1.0, 0.9, 0.8, 0.7, 0.6], 'no_of_points': 5}, - {'nu': [2.8e14, 3.0e14, 3.2e14, 3.4e14], - 'x_sect': [2.0, 1.8, 1.6, 1.4], 'no_of_points': 4} - ] - - -@pytest.fixture(scope='module') -def mock_sample_nu(): - SAMPLE_NUFUNC = CFUNCTYPE(c_double, POINTER(RPacket), - POINTER(StorageModel), POINTER(RKState)) - - def sample_nu_simple(packet, model, mt_state): - return packet.contents.nu - - return SAMPLE_NUFUNC(sample_nu_simple) - - -@pytest.fixture(scope='function') -def model_3lvlatom(model): - model.line2macro_level_upper = (c_int64 * 3)(*[2, 1, 2]) - model.macro_block_references = (c_int64 * 3)(*[0, 2, 5]) - - transition_probabilities = [ - 0.0, 0.0, 0.75, 0.25, 0.0, 0.25, 0.5, 0.25, 0.0, # shell_id = 0 - 0.0, 0.0, 1.00, 0.00, 0.0, 0.00, 0.0, 1.00, 0.0 # shell_id = 1 - ] - - nd = len(transition_probabilities)//2 - model.transition_type = (c_int64 * nd)(*[1, 1, -1, 1, 0, 0, -1, -1, 0]) - model.destination_level_id = (c_int64 * nd)(*[1, 2, 0, 2, 0, 1, 1, 0, 0]) - model.transition_line_id = (c_int64 * nd)(*[0, 1, 1, 2, 1, 2, 2, 0, 0]) - - model.transition_probabilities_nd = c_int64(nd) - model.transition_probabilities = (c_double * (nd * 2))(*transition_probabilities) - - model.last_line_interaction_out_id = (c_int64 * 1)(*[-5]) - - return model - - -def d_cont_setter(d_cont, model, packet): - model.inverse_electron_densities[packet.current_shell_id] = c_double(1.0) - model.inverse_sigma_thomson = c_double(1.0) - packet.tau_event = c_double(d_cont) - - -def d_line_setter(d_line, model, packet): - packet.mu = c_double(0.0) - scale = d_line * 1e1 - model.time_explosion = c_double(INVERSE_C * scale) - packet.nu = c_double(1.0) - nu_line = (1. - d_line/scale) - packet.nu_line = c_double(nu_line) - - -def d_boundary_setter(d_boundary, model, packet): - packet.mu = c_double(1e-16) - r_outer = 2. * d_boundary - model.r_outer[packet.current_shell_id] = r_outer - - r = np.sqrt(r_outer**2 - d_boundary**2) - packet.r = r - - - - -""" -Important Tests: ----------------- -The tests written further (till next block comment is encountered) have been -categorized as important tests, these tests correspond to methods which are -relatively old and stable code. -""" - - -@pytest.mark.parametrize( - ['x', 'x_insert', 'imin', 'imax', 'expected_params'], - [([5.0, 4.0, 3.0, 1.0], 2.0, 0, 3, - {'result': 2, 'ret_val': TARDIS_ERROR_OK}), - - ([5.0, 4.0, 3.0, 2.0], 0.0, 0, 3, - {'result': 0, 'ret_val': TARDIS_ERROR_BOUNDS_ERROR})] -) -def test_reverse_binary_search(clib, x, x_insert, imin, imax, expected_params): - x = (c_double * (imax - imin + 1))(*x) - x_insert = c_double(x_insert) - imin = c_int64(imin) - imax = c_int64(imax) - obtained_result = c_int64(0) - - clib.reverse_binary_search.restype = c_uint - obtained_tardis_error = clib.reverse_binary_search( - byref(x), x_insert, imin, imax, byref(obtained_result)) - - assert obtained_result.value == expected_params['result'] - assert obtained_tardis_error == expected_params['ret_val'] - - -@pytest.mark.parametrize( - ['nu', 'nu_insert', 'number_of_lines', 'expected_params'], - [([0.5, 0.4, 0.3, 0.1], 0.2, 4, - {'result': 3, 'ret_val': TARDIS_ERROR_OK}), - - ([0.5, 0.4, 0.3, 0.2], 0.1, 4, - {'result': 4, 'ret_val': TARDIS_ERROR_OK}), - - ([0.4, 0.3, 0.2, 0.1], 0.5, 4, - {'result': 0, 'ret_val': TARDIS_ERROR_OK})] -) -def test_line_search(clib, nu, nu_insert, number_of_lines, expected_params): - nu = (c_double * number_of_lines)(*nu) - nu_insert = c_double(nu_insert) - number_of_lines = c_int64(number_of_lines) - obtained_result = c_int64(0) - - clib.line_search.restype = c_uint - obtained_tardis_error = clib.line_search( - byref(nu), nu_insert, number_of_lines, byref(obtained_result)) - - assert obtained_result.value == expected_params['result'] - assert obtained_tardis_error == expected_params['ret_val'] - -@pytest.mark.parametrize( - ['x', 'x_insert', 'imin', 'imax', 'expected_params'], - [([2.0, 4.0, 6.0, 7.0], 5.0, 0, 3, - {'result': 2, 'ret_val': TARDIS_ERROR_OK}), - - ([2.0, 3.0, 5.0, 7.0], 8.0, 0, 3, - {'result': 0, 'ret_val': TARDIS_ERROR_BOUNDS_ERROR}), - - ([2.0, 4.0, 6.0, 7.0], 4.0, 0, 3, - {'result': 0, 'ret_val': TARDIS_ERROR_OK}) - ] -) -def test_binary_search(clib, x, x_insert, imin, imax, expected_params): - x = (c_double * (imax - imin + 1))(*x) - x_insert = c_double(x_insert) - imin = c_int64(imin) - imax = c_int64(imax) - obtained_result = c_int64(0) - - clib.binary_search.restype = c_uint - obtained_tardis_error = clib.binary_search( - byref(x), x_insert, imin, imax, byref(obtained_result)) - - assert obtained_result.value == expected_params['result'] - assert obtained_tardis_error == expected_params['ret_val'] - - -@pytest.mark.parametrize( - ['mu', 'r', 'inv_t_exp', 'expected'], - [(0.3, 7.5e14, 1 / 5.2e7, 0.9998556693818854), - (-.3, 8.1e14, 1 / 2.6e7, 1.0003117541351274)] -) -def test_rpacket_doppler_factor(clib, mu, r, inv_t_exp, expected, packet, model): - # Set the params from test cases here - packet.mu = mu - packet.r = r - model.inverse_time_explosion = inv_t_exp - - # Perform any other setups just before this, they can be additional calls - # to other methods or introduction of some temporary variables - - # Set `restype` attribute if returned quantity is used - clib.rpacket_doppler_factor.restype = c_double - # Call the C method (make sure to pass quantities as `ctypes` data types) - obtained = clib.rpacket_doppler_factor(byref(packet), byref(model)) - - # Perform required assertions - assert_almost_equal(obtained, expected) - - -@pytest.mark.parametrize( - ['packet_params', 'expected_params'], - [({'mu': 0.3, 'r': 7.5e14}, - {'d_boundary': 259376919351035.88}), - - ({'mu': -.3, 'r': 7.5e13}, - {'d_boundary': -664987228972291.5}), - - ({'mu': -.3, 'r': 7.5e14}, - {'d_boundary': 709376919351035.9})] -) -def test_compute_distance2boundary(clib, packet_params, expected_params, packet, model): - packet.mu = packet_params['mu'] - packet.r = packet_params['r'] - - clib.compute_distance2boundary(byref(packet), byref(model)) - - assert_almost_equal(packet.d_boundary, expected_params['d_boundary']) - - -# TODO: split this into two tests - one to assert errors and other for d_line -@pytest.mark.parametrize( - ['packet_params', 'expected_params'], - [({'nu_line': 0.1, 'next_line_id': 0, 'last_line': 1}, - {'tardis_error': TARDIS_ERROR_OK, 'd_line': 1e+99}), - - ({'nu_line': 0.2, 'next_line_id': 1, 'last_line': 0}, - {'tardis_error': TARDIS_ERROR_OK, 'd_line': 7.792353908000001e+17}), - - ({'nu_line': 0.5, 'next_line_id': 1, 'last_line': 0}, - {'tardis_error': TARDIS_ERROR_COMOV_NU_LESS_THAN_NU_LINE, 'd_line': 0.0}), - - ({'nu_line': 0.6, 'next_line_id': 0, 'last_line': 0}, - {'tardis_error': TARDIS_ERROR_COMOV_NU_LESS_THAN_NU_LINE, 'd_line': 0.0})] -) -def test_compute_distance2line(clib, packet_params, expected_params, packet, model): - packet.nu_line = packet_params['nu_line'] - packet.next_line_id = packet_params['next_line_id'] - packet.last_line = packet_params['last_line'] - - packet.d_line = 0.0 - clib.compute_distance2line.restype = c_uint - obtained_tardis_error = clib.compute_distance2line(byref(packet), byref(model)) - - assert_almost_equal(packet.d_line, expected_params['d_line']) - assert obtained_tardis_error == expected_params['tardis_error'] - - -@pytest.mark.parametrize( - ['packet_params', 'expected_params'], - [({'virtual_packet': 0}, - {'chi_cont': 6.652486e-16, 'd_cont': 4.359272608766106e+28}), - - ({'virtual_packet': 1}, - {'chi_cont': 6.652486e-16, 'd_cont': 1e+99})] -) -def test_compute_distance2continuum(clib, packet_params, expected_params, packet, model): - packet.virtual_packet = packet_params['virtual_packet'] - - clib.compute_distance2continuum(byref(packet), byref(model)) - - assert_almost_equal(packet.chi_cont, expected_params['chi_cont']) - assert_almost_equal(packet.d_cont, expected_params['d_cont']) - - -@pytest.mark.parametrize('full_relativity', [1, 0]) -@pytest.mark.parametrize( - ['packet_params', 'expected_params'], - [({'nu': 0.4, 'mu': 0.3, 'energy': 0.9, 'r': 7.5e14}, - {'mu': 0.3120599529139568, 'r': 753060422542573.9, - 'j': 8998701024436.969, 'nubar': 3598960894542.354}), - - ({'nu': 0.6, 'mu': -.5, 'energy': 0.5, 'r': 8.1e14}, - {'mu': -.4906548373534084, 'r': 805046582503149.2, - 'j': 5001298975563.031, 'nubar': 3001558973156.1387})] -) -def test_move_packet(clib, packet_params, expected_params, - packet, model, full_relativity): - packet.nu = packet_params['nu'] - packet.mu = packet_params['mu'] - packet.energy = packet_params['energy'] - packet.r = packet_params['r'] - model.full_relativity = full_relativity - - clib.rpacket_doppler_factor.restype = c_double - doppler_factor = clib.rpacket_doppler_factor(byref(packet), byref(model)) - clib.move_packet(byref(packet), byref(model), c_double(1.e13)) - - assert_almost_equal(packet.mu, expected_params['mu']) - assert_almost_equal(packet.r, expected_params['r']) - - expected_j = expected_params['j'] - expected_nubar = expected_params['nubar'] - if full_relativity: - expected_j *= doppler_factor - expected_nubar *= doppler_factor - - assert_allclose(model.js[packet.current_shell_id], - expected_j, rtol=5e-7) - assert_allclose(model.nubars[packet.current_shell_id], - expected_nubar, rtol=5e-7) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - ['packet_params', 'j_blue_idx', 'expected'], - [({'nu': 0.30, 'energy': 0.30}, 0, 1.0), - ({'nu': 0.20, 'energy': 1.e5}, 0, 5e5), - ({'nu': 2e15, 'energy': 0.50}, 1, 2.5e-16), - ({'nu': 0.40, 'energy': 1e-7}, 1, 2.5e-7)], -) -def test_increment_j_blue_estimator_full_relativity(clib, packet_params, - j_blue_idx, expected, - packet, model): - packet.nu = packet_params['nu'] - packet.energy = packet_params['energy'] - model.full_relativity = True - - clib.increment_j_blue_estimator(byref(packet), byref(model), - c_double(packet.d_line), - c_int64(j_blue_idx)) - - assert_almost_equal(model.line_lists_j_blues[j_blue_idx], expected) - - -@pytest.mark.parametrize( - ['packet_params', 'j_blue_idx', 'expected'], - [({'nu': 0.1, 'mu': 0.3, 'r': 7.5e14}, 0, 8.998643292289723), - ({'nu': 0.2, 'mu': -.3, 'r': 7.7e14}, 0, 4.499971133976377), - ({'nu': 0.5, 'mu': 0.5, 'r': 7.9e14}, 1, 0.719988453650551), - ({'nu': 0.6, 'mu': -.5, 'r': 8.1e14}, 1, 0.499990378058792)] -) -def test_increment_j_blue_estimator(clib, packet_params, j_blue_idx, expected, packet, model): - packet.nu = packet_params['nu'] - packet.mu = packet_params['mu'] - packet.r = packet_params['r'] - - clib.compute_distance2line(byref(packet), byref(model)) - clib.move_packet(byref(packet), byref(model), c_double(1.e13)) - clib.increment_j_blue_estimator(byref(packet), byref(model), - c_double(packet.d_line), c_int64(j_blue_idx)) - - assert_almost_equal(model.line_lists_j_blues[j_blue_idx], expected) - - -@pytest.mark.parametrize( - ['packet_params', 'expected_params'], - [({'virtual_packet': 0, 'current_shell_id': 0, 'next_shell_id': 1}, - {'status': TARDIS_PACKET_STATUS_IN_PROCESS, 'current_shell_id': 1}), - - ({'virtual_packet': 1, 'current_shell_id': 1, 'next_shell_id': 1}, - {'status': TARDIS_PACKET_STATUS_EMITTED, 'current_shell_id': 1, - 'tau_event': 29000000000000.008}), - - ({'virtual_packet': 1, 'current_shell_id': 0, 'next_shell_id': -1}, - {'status': TARDIS_PACKET_STATUS_REABSORBED, 'current_shell_id': 0, - 'tau_event': 29000000000000.008})] -) -def test_move_packet_across_shell_boundary(clib, packet_params, expected_params, - packet, model, mt_state): - packet.virtual_packet = packet_params['virtual_packet'] - packet.current_shell_id = packet_params['current_shell_id'] - packet.next_shell_id = packet_params['next_shell_id'] - - clib.move_packet_across_shell_boundary(byref(packet), byref(model), - c_double(1.e13), byref(mt_state)) - - if packet_params['virtual_packet'] == 1: - assert_almost_equal(packet.tau_event, expected_params['tau_event']) - assert packet.status == expected_params['status'] - assert packet.current_shell_id == expected_params['current_shell_id'] - - -@pytest.mark.parametrize( - ['packet_params', 'expected_params'], - [({'nu': 0.4, 'mu': 0.3, 'energy': 0.9, 'r': 7.5e14}, - {'nu': 0.39974659819356556, 'energy': 0.8994298459355226}), - - ({'nu': 0.6, 'mu': -.5, 'energy': 0.5, 'r': 8.1e14}, - {'nu': 0.5998422620533325, 'energy': 0.4998685517111104})] -) -def test_montecarlo_thomson_scatter(clib, packet_params, expected_params, packet, - model, mt_state): - packet.nu = packet_params['nu'] - packet.mu = packet_params['mu'] - packet.energy = packet_params['energy'] - packet.r = packet_params['r'] - - clib.montecarlo_thomson_scatter(byref(packet), byref(model), - c_double(1.e13), byref(mt_state)) - - assert_almost_equal(packet.nu, expected_params['nu']) - assert_almost_equal(packet.energy, expected_params['energy']) - - -@pytest.mark.parametrize( - ['packet_params', 'expected_params'], - # TODO: Add scientifically sound test cases. - [({'virtual_packet': 1, 'tau_event': 2.9e13, 'last_line': 0}, - {'tau_event': 2.9e13, 'next_line_id': 2}), - - ({'virtual_packet': 0, 'tau_event': 2.9e13, 'last_line': 0}, - {'tau_event': 2.9e13, 'next_line_id': 2}), - - ({'virtual_packet': 0, 'tau_event': 2.9e13, 'last_line': 0}, - {'tau_event': 2.9e13, 'next_line_id': 2}), - ] -) -def test_montecarlo_line_scatter(clib, packet_params, expected_params, packet, model, mt_state): - packet.virtual_packet = packet_params['virtual_packet'] - packet.tau_event = packet_params['tau_event'] - packet.last_line = packet_params['last_line'] - - clib.montecarlo_line_scatter(byref(packet), byref(model), - c_double(1.e13), byref(mt_state)) - - assert_almost_equal(packet.tau_event, expected_params['tau_event']) - assert_almost_equal(packet.next_line_id, expected_params['next_line_id']) - - -@pytest.mark.parametrize( - ['distances', 'expected'], - [({'boundary': 1.3e13, 'continuum': 1e14, 'line': 1e15}, - {'handler': 'move_packet_across_shell_boundary', 'distance': 1.3e13}), - - ({'boundary': 1.3e13, 'continuum': 1e14, 'line': 2.5e12}, - {'handler': 'montecarlo_line_scatter', 'distance': 2.5e12}), - - ({'boundary': 1.3e13, 'continuum': 1e11, 'line': 2.5e12}, - {'handler': 'montecarlo_thomson_scatter', 'distance': 1e11})] -) -def test_get_event_handler(clib, packet, model, mt_state, distances, expected): - d_cont_setter(distances['continuum'], model, packet) - d_line_setter(distances['line'], model, packet) - d_boundary_setter(distances['boundary'], model, packet) - obtained_distance = c_double() - - clib.get_event_handler.restype = c_void_p - obtained_handler = clib.get_event_handler(byref(packet), byref(model), - byref(obtained_distance), - byref(mt_state)) - - expected_handler = getattr(clib, expected['handler']) - expected_handler = cast(expected_handler, c_void_p).value - - assert_equal(obtained_handler, expected_handler) - assert_allclose(obtained_distance.value, expected['distance'], rtol=1e-10) - - -@pytest.mark.parametrize( - ['z_random', 'packet_params', 'expected'], - [(0.22443743797312765, - {'activation_level': 1, 'shell_id': 0}, 1), # Direct deactivation - - (0.78961460371187597, # next z_random = 0.818455414618 - {'activation_level': 1, 'shell_id': 0}, 0), # Upwards jump, then deactivation - - (0.22443743797312765, # next z_random = 0.545678896748 - {'activation_level': 2, 'shell_id': 0}, 1), # Downwards jump, then deactivation - - (0.765958602560605, # next z_random = 0.145914243888, 0.712382380384 - {'activation_level': 1, 'shell_id': 0}, 1), # Upwards jump, downwards jump, then deactivation - - (0.22443743797312765, - {'activation_level': 2, 'shell_id': 1}, 0)] # Direct deactivation -) -def test_macro_atom(clib, model_3lvlatom, packet, z_random, packet_params, get_rkstate, expected): - packet.macro_atom_activation_level = packet_params['activation_level'] - packet.current_shell_id = packet_params['shell_id'] - rkstate = get_rkstate(z_random) - - clib.macro_atom(byref(packet), byref(model_3lvlatom), byref(rkstate)) - obtained_line_id = model_3lvlatom.last_line_interaction_out_id[packet.id] - - assert_equal(obtained_line_id, expected) - - - -""" -Simple Tests: ----------------- -These test check very simple pices of code still work. -""" - -@pytest.mark.parametrize( - ['packet_params', 'line_idx', 'expected'], - [({'energy': 0.0}, 0, 0), - ({'energy': 1.0}, 1, 1), - ({'energy': 0.5}, 2, 1.5)] -) -@pytest.mark.skipif(True, reason="Needs rewrite to be relevant.") -def test_increment_Edotlu_estimator(clib, packet_params, line_idx, expected, packet, model): - packet.energy = packet_params['energy'] - - clib.increment_Edotlu_estimator(byref(packet), byref(model), c_int64(line_idx)) - - assert_almost_equal(model.line_lists_Edotlu[line_idx], expected) - - -""" -Difficult Tests: ----------------- -The tests written further are more complex than previous tests. They require -proper design procedure. They are not taken up yet and intended to be -completed together in future. -""" - - - -@pytest.mark.skipif(True, reason="Yet to be written.") -def test_montecarlo_one_packet(packet, model, mt_state): - pass - - -@pytest.mark.skipif(True, reason="Yet to be written.") -def test_montecarlo_one_packet_loop(packet, model, mt_state): - pass - - -@pytest.mark.skipif(True, reason="Yet to be written.") -def test_montecarlo_main_loop(packet, model, mt_state): - pass - - - -""" -Continuum Tests: ----------------- -The tests written further (till next block comment is encountered) are for the -methods related to continuum interactions. -""" - - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - 't_electron', [2500., 15000.] -) -def test_sample_nu_free_free(clib, t_electron, packet, model, mt_state_seeded, expected_ff_emissivity): - model.t_electrons[packet.current_shell_id] = t_electron - clib.sample_nu_free_free.restype = c_double - - nu_bins, expected_emissivity = expected_ff_emissivity(t_electron) - - nus = [] - for _ in range(int(1e5)): - nu = clib.sample_nu_free_free(byref(packet), byref(model), byref(mt_state_seeded)) - nus.append(nu) - - obtained_emissivity, _ = np.histogram(nus, density=True, bins=nu_bins) - - assert_allclose(obtained_emissivity, expected_emissivity, rtol=1e-10) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - ['packet_params', 't_electrons', 'chi_ff_factor', 'expected'], - [({'nu': 4.5e14, 'mu': 0.0, 'current_shell_id': 1}, 15000, 2.0, 1.6746639430359494e-44), - ({'nu': 3.0e15, 'mu': 0.0, 'current_shell_id': 0}, 5000, 3.0, 1.1111111111107644e-46), - ({'nu': 3.0e15, 'mu': 0.4, 'current_shell_id': 0}, 10000, 4.0, 1.5638286016098277e-46)] -) -def test_calculate_chi_ff(clib, packet, model, packet_params, t_electrons, chi_ff_factor, expected): - packet.mu = packet_params['mu'] - packet.nu = packet_params['nu'] - packet.current_shell_id = packet_params['current_shell_id'] - packet.r = 1.04e17 - - model.t_electrons[packet_params['current_shell_id']] = t_electrons - model.chi_ff_factor[packet_params['current_shell_id']] = chi_ff_factor - - clib.calculate_chi_ff(byref(packet), byref(model)) - obtained = packet.chi_ff - - assert_equal(obtained, expected) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - ['continuum_status', 'z_random', 'packet_params', 'expected'], - [(CONTINUUM_OFF, 0.94183547596539363, - {'chi_c': 1.0, 'chi_th': 0.4, 'chi_bf': 0.5}, - 'montecarlo_thomson_scatter'), - - (CONTINUUM_ON, 0.22443743797312765, - {'chi_c': 1.0, 'chi_th': 0.4, 'chi_bf': 0.5}, - 'montecarlo_thomson_scatter'), - - (CONTINUUM_ON, 0.54510721066252377, - {'chi_c': 1.0, 'chi_th': 0.4, 'chi_bf': 0.5}, - 'montecarlo_bound_free_scatter'), - - (CONTINUUM_ON, 0.94183547596539363, - {'chi_c': 1.0, 'chi_th': 0.4, 'chi_bf': 0.5}, - 'montecarlo_free_free_scatter'), - - (CONTINUUM_ON, 0.22443743797312765, - {'chi_c': 1e2, 'chi_th': 1e1, 'chi_bf': 2e1}, - 'montecarlo_bound_free_scatter')] -) -def test_montecarlo_continuum_event_handler(clib, continuum_status, expected, z_random, - packet_params, packet, model, get_rkstate): - packet.chi_cont = packet_params['chi_c'] - packet.chi_th = packet_params['chi_th'] - packet.chi_bf = packet_params['chi_bf'] - model.cont_status = continuum_status - - rkstate = get_rkstate(z_random) - - clib.montecarlo_continuum_event_handler.restype = c_void_p - obtained = clib.montecarlo_continuum_event_handler(byref(packet), - byref(model), byref(rkstate)) - expected = getattr(clib, expected) - expected = cast(expected, c_void_p).value - - assert_equal(obtained, expected) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - ['nu', 'continuum_id', 'expected', 'bf_treatment'], - [(4.40e14, 1, 0.00, BoundFreeTreatment.LIN_INTERPOLATION), - (3.25e14, 1, 0.75, BoundFreeTreatment.LIN_INTERPOLATION), - (4.03e14, 0, 0.97, BoundFreeTreatment.LIN_INTERPOLATION), - (4.10e14 + 1e-1, 0, 0.90, BoundFreeTreatment.LIN_INTERPOLATION), - pytest.param(4.1e14, 0, 0.90, BoundFreeTreatment.LIN_INTERPOLATION, - marks=pytest.mark.xfail), - (6.50e14, 0, 0.23304506144742834, BoundFreeTreatment.HYDROGENIC), - (3.40e14, 2, 1.1170364339507428, BoundFreeTreatment.HYDROGENIC)] -) -def test_bf_cross_section(clib, nu, continuum_id, model_w_edges, expected, bf_treatment): - model_w_edges.bf_treatment = bf_treatment.value - - clib.bf_cross_section.restype = c_double - obtained = clib.bf_cross_section(byref(model_w_edges), continuum_id, c_double(nu)) - - assert_almost_equal(obtained, expected) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - ['packet_params', 'expected'], - [({'nu': 4.13e14, 'mu': 0.0, 'current_shell_id': 1}, - [3.2882087455641473, 0.0, 0.0]), - - ({'nu': 3.27e14, 'mu': 0.0, 'current_shell_id': 0}, - [0.0, 1.3992114634681028, 5.702548202131454]), - - ({'nu': 3.27e14, 'mu': -0.4, 'current_shell_id': 0}, - [0.0, 1.2670858, 5.4446587])] -) -def test_calculate_chi_bf(clib, packet_params, expected, packet, model_w_edges): - model_w_edges.l_pop = (c_double * 6)(*range(1, 7)) - model_w_edges.l_pop_r = (c_double * 6)(*np.linspace(0.1, 0.6, 6)) - model_w_edges.t_electrons[packet_params['current_shell_id']] = 1e4 - - packet.mu = packet_params['mu'] - packet.nu = packet_params['nu'] - packet.r = 1.04e17 - packet.current_shell_id = packet_params['current_shell_id'] - packet.chi_bf_tmp_partial = (c_double * model_w_edges.no_of_edges)() - - clib.calculate_chi_bf(byref(packet), byref(model_w_edges)) - - obtained_chi_bf_tmp = np.ctypeslib.as_array(packet.chi_bf_tmp_partial, shape=(model_w_edges.no_of_edges,)) - expected_chi_bf_tmp = np.array(expected) - expected_chi_bf = expected_chi_bf_tmp[expected_chi_bf_tmp > 0][-1] - - assert_almost_equal(obtained_chi_bf_tmp, expected_chi_bf_tmp) - assert_almost_equal(packet.chi_bf, expected_chi_bf) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - ['comov_energy', 'distance', 'chi_ff', 'no_of_updates', 'expected'], - [(1.3, 1.3e14, 3e-12, 1, 507.), - (0.9, 0.7e15, 2e-12, 10, 1.260e4), - (0.8, 0.8e13, 1.5e-12, 35, 336.)] -) -def test_increment_continuum_estimators_ff_heating_estimator(clib, packet, model_w_edges, comov_energy, distance, - chi_ff, no_of_updates, expected): - packet.chi_ff = chi_ff - - for _ in range(no_of_updates): - clib.increment_continuum_estimators(byref(packet), byref(model_w_edges), c_double(distance), - c_double(0), c_double(comov_energy)) - obtained = model_w_edges.ff_heating_estimator[packet.current_shell_id] - - assert_almost_equal(obtained, expected) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - ['comov_nus', 'expected'], - [([4.05e14, 4.17e14, 3.3e14, 3.2e14, 2.9e14], [2, 2, 3]), - ([4.15e15, 3.25e14, 3.3e14, 2.85e14, 2.9e14], [0, 2, 4])] -) -def test_increment_continuum_estimators_photo_ion_estimator_statistics(clib, packet, model_w_edges, comov_nus, expected): - for comov_nu in comov_nus: - clib.increment_continuum_estimators(byref(packet), byref(model_w_edges), c_double(1e13), - c_double(comov_nu), c_double(1.0)) - - no_of_edges = model_w_edges.no_of_edges - no_of_shells = model_w_edges.no_of_shells - - obtained = np.ctypeslib.as_array(model_w_edges.photo_ion_estimator_statistics, - shape=(no_of_edges * no_of_shells,)) - obtained = np.reshape(obtained, newshape=(no_of_shells, no_of_edges), order='F') - obtained = obtained[packet.current_shell_id] - expected = np.array(expected) - - assert_array_equal(obtained, expected) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - ['comov_energy', 'distance', 'comov_nus', 'expected'], - [(1.3, 1.3e14, [4.05e14, 2.65e14], - {"photo_ion": [0.39641975308641975, 0., 0.], - "stim_recomb": [0.056757061269242064, 0., 0.], - "bf_heating": [1.9820987654321076e12, 0., 0.], - "stim_recomb_cooling": [283784812699.75476, 0., 0.]}), - - (0.9, 0.7e15, [3.25e14, 2.85e14], - {"photo_ion": [0., 1.4538461538461538, 7.315141700404858], - "stim_recomb": [0., 0.3055802, 1.7292954], - "bf_heating": [0., 36346153846153.82, 156760323886639.69], - "stim_recomb_cooling": [0., 7639505724285.9746, 33907776077426.875]})] -) -def test_increment_continuum_estimators_bf_estimators(clib, packet, model_w_edges, comov_energy, - distance, comov_nus, expected): - for comov_nu in comov_nus: - clib.increment_continuum_estimators(byref(packet), byref(model_w_edges), c_double(distance), - c_double(comov_nu), c_double(comov_energy)) - - no_of_edges = model_w_edges.no_of_edges - no_of_shells = model_w_edges.no_of_shells - - for estim_name, expected_value in expected.items(): - obtained = np.ctypeslib.as_array(getattr(model_w_edges, estim_name + "_estimator"), - shape=(no_of_edges * no_of_shells,)) - obtained = np.reshape(obtained, newshape=(no_of_shells, no_of_edges), order='F') - obtained = obtained[packet.current_shell_id] - - assert_almost_equal(obtained, np.array(expected_value)) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - ['packet_params', 'expected_params'], - [({'nu_comov': 1.1e16, 'mu': 0.0, 'r': 1.4e14}, - {'next_line_id': 5, 'last_line': True, 'nu': 1.1e16, 'type_id': 3}), - - ({'nu_comov': 1.3e16, 'mu': 0.3, 'r': 7.5e14}, - {'next_line_id': 0, 'last_line': False, 'nu': 1.30018766e+16, 'type_id': 3}), - - ({'nu_comov': 1.24e16, 'mu': -0.3, 'r': 7.5e14}, - {'next_line_id': 2, 'last_line': False, 'nu': 1.23982106e+16, 'type_id': 4})] -) -def test_continuum_emission(clib, packet, model, mock_sample_nu, packet_params, expected_params, mt_state): - packet.nu = packet_params['nu_comov'] # Is returned by mock function mock_sample_nu - packet.mu = packet_params['mu'] - packet.r = packet_params['r'] - expected_interaction_out_type = expected_params['type_id'] - - clib.continuum_emission(byref(packet), byref(model), byref(mt_state), - mock_sample_nu, expected_interaction_out_type) - - obtained_next_line_id = packet.next_line_id - obtained_last_interaction_out_type = model.last_interaction_out_type[0] - - assert_equal(obtained_next_line_id, expected_params['next_line_id']) - assert_equal(packet.last_line, expected_params['last_line']) - assert_equal(expected_interaction_out_type, obtained_last_interaction_out_type) - assert_allclose(packet.nu, expected_params['nu'], rtol=1e-7) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - ['packet_params', 'expected'], - [({'next_line_id': 3, 'last_line': 0}, 1), - ({'next_line_id': 5, 'last_line': 1}, 0), - ({'next_line_id': 2, 'last_line': 0}, 0), - ({'next_line_id': 1, 'last_line': 0}, 1)] -) -def test_test_for_close_line(clib, packet, model, packet_params, expected): - packet.nu_line = model.line_list_nu[packet_params['next_line_id'] - 1] - packet.next_line_id = packet_params['next_line_id'] - - clib.test_for_close_line(byref(packet), byref(model)) - - assert_equal(expected, packet.close_line) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - ['packet_params', 'z_random', 'expected'], - [({'current_continuum_id': 1, 'chi_bf_tmp_partial': [0.0, 0.23e13, 1.0e13]}, - 0.22443743797312765, 1), - - ({'current_continuum_id': 1, 'chi_bf_tmp_partial': [0.0, 0.23e10, 1.0e10]}, - 0.78961460371187597, 2), - - ({'current_continuum_id': 0, 'chi_bf_tmp_partial': [0.2e5, 0.5e5, 0.6e5, 1.0e5, 1.0e5]}, - 0.78961460371187597, 3)] -) -def test_montecarlo_bound_free_scatter_continuum_selection(clib, packet, model_3lvlatom, packet_params, - get_rkstate, z_random, expected): - rkstate = get_rkstate(z_random) - packet.current_continuum_id = packet_params['current_continuum_id'] - - chi_bf_tmp = packet_params['chi_bf_tmp_partial'] - packet.chi_bf_tmp_partial = (c_double * len(chi_bf_tmp))(*chi_bf_tmp) - packet.chi_bf = chi_bf_tmp[-1] - model_3lvlatom.no_of_edges = len(chi_bf_tmp) - - clib.montecarlo_bound_free_scatter(byref(packet), byref(model_3lvlatom), - c_double(1.e13), byref(rkstate)) - - assert_equal(packet.current_continuum_id, expected) - assert_equal(model_3lvlatom.last_line_interaction_in_id[packet.id], expected) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - ['mu', 'r', 'inv_t_exp', 'full_relativity'], - [(0.8, 7.5e14, 1 / 5.2e5, 1), - (-0.7, 7.5e14, 1 / 5.2e5, 1), - (0.3, 7.5e14, 1 / 2.2e5, 1), - (0.0, 7.5e14, 1 / 2.2e5, 1), - (-0.7, 7.5e14, 1 / 5.2e5, 0)] -) -def test_frame_transformations(clib, packet, model, mu, r, - inv_t_exp, full_relativity): - packet.r = r - packet.mu = mu - model.inverse_time_explosion = inv_t_exp - model.full_relativity = full_relativity - clib.rpacket_doppler_factor.restype = c_double - clib.rpacket_inverse_doppler_factor.restype = c_double - - inverse_doppler_factor = clib.rpacket_inverse_doppler_factor(byref(packet), byref(model)) - clib.angle_aberration_CMF_to_LF(byref(packet), byref(model)) - - doppler_factor = clib.rpacket_doppler_factor(byref(packet), byref(model)) - - assert_almost_equal(doppler_factor * inverse_doppler_factor, 1.0) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - ['mu', 'r', 'inv_t_exp'], - [(0.8, 7.5e14, 1 / 5.2e5), - (-0.7, 7.5e14, 1 / 5.2e5), - (0.3, 7.5e14, 1 / 2.2e5), - (0.0, 7.5e14, 1 / 2.2e5), - (-0.7, 7.5e14, 1 / 5.2e5)] -) -def test_angle_transformation_invariance(clib, packet, model, - mu, r, inv_t_exp): - packet.r = r - packet.mu = mu - model.inverse_time_explosion = inv_t_exp - model.full_relativity = 1 - clib.angle_aberration_LF_to_CMF.restype = c_double - - clib.angle_aberration_CMF_to_LF(byref(packet), byref(model)) - mu_obtained = clib.angle_aberration_LF_to_CMF( - byref(packet), byref(model), c_double(packet.mu)) - - assert_almost_equal(mu_obtained, mu) - - -@pytest.mark.continuumtest -@pytest.mark.parametrize( - 'full_relativity', - [1, 0] -) -@pytest.mark.parametrize( - ['mu', 'r', 't_exp', 'nu', 'nu_line'], - [(0.8, 7.5e14, 5.2e5, 1.0e15, 9.4e14), - (0.0, 6.3e14, 2.2e5, 6.0e12, 5.8e12), - (1.0, 9.0e14, 2.2e5, 4.0e8, 3.4e8), - (0.9, 9.0e14, 0.5e5, 1.0e15, 4.5e14), - (-0.7, 7.5e14, 5.2e5, 1.0e15, 9.8e14), - (-1.0, 6.3e14, 2.2e5, 6.0e12, 6.55e12)] -) -def test_compute_distance2line_relativistic(clib, mu, r, t_exp, nu, nu_line, - full_relativity, packet, model): - packet.r = r - packet.mu = mu - packet.nu = nu - packet.nu_line = nu_line - model.inverse_time_explosion = 1 / t_exp - model.time_explosion = t_exp - model.full_relativity = full_relativity - - clib.rpacket_doppler_factor.restype = c_double - - clib.compute_distance2line(byref(packet), byref(model)) - clib.move_packet(byref(packet), byref(model), c_double(packet.d_line)) - - doppler_factor = clib.rpacket_doppler_factor(byref(packet), byref(model)) - comov_nu = packet.nu * doppler_factor - - assert_allclose(comov_nu, nu_line, rtol=1e-14) - - -@pytest.mark.continuumtest -@pytest.mark.skipif(True, reason="Yet to be written.") -def test_montecarlo_free_free_scatter(packet, model, mt_state): - pass diff --git a/tardis/montecarlo/tests/test_formal_integral.py b/tardis/montecarlo/tests/test_formal_integral.py index 616de0d2d7d..ed9afed8600 100644 --- a/tardis/montecarlo/tests/test_formal_integral.py +++ b/tardis/montecarlo/tests/test_formal_integral.py @@ -2,75 +2,54 @@ import numpy as np from tardis import constants as c -import ctypes -from ctypes import ( - c_double, - c_int - ) - from numpy.ctypeslib import ( - as_array, - as_ctypes, - ndpointer - ) + as_array, + as_ctypes, +) import numpy.testing as ntest from tardis.util.base import intensity_black_body -from tardis.montecarlo.struct import StorageModel +import tardis.montecarlo.formal_integral as formal_integral + + +pytestmark = pytest.mark.skip(reason="Port from C to numba") @pytest.mark.parametrize( - ['nu', 'T'], - [ - (1e14, 1e4), - (0, 1), - (1, 1), - ] - ) -def test_intensity_black_body(clib, nu, T): - func = clib.intensity_black_body - func.restype = c_double - func.argtypes = [c_double, c_double] + ["nu", "T"], + [ + (1e14, 1e4), + (0, 1), + (1, 1), + ], +) +def test_intensity_black_body(nu, T): + func = formal_integral.intensity_black_body actual = func(nu, T) + print(actual, type(actual)) expected = intensity_black_body(nu, T) - ntest.assert_almost_equal( - actual, - expected - ) + ntest.assert_almost_equal(actual, expected) -@pytest.mark.parametrize( - 'N', - (1e2, 1e3, 1e4, 1e5) - ) -def test_trapezoid_integration(clib, N): - func = clib.trapezoid_integration - func.restype = c_double - h = 1. +@pytest.mark.parametrize("N", (1e2, 1e3, 1e4, 1e5)) +def test_trapezoid_integration(N): + func = formal_integral.trapezoid_integration + h = 1.0 N = int(N) - func.argtypes = [ - ndpointer(c_double), - c_double, - c_int - ] data = np.random.random(N) actual = func(data, h, int(N)) expected = np.trapz(data) - ntest.assert_almost_equal( - actual, - expected - ) + ntest.assert_almost_equal(actual, expected) @pytest.mark.skipif( - True, - reason='static inline functions are not inside the library' - ) -def test_calculate_z(clib): + True, reason="static inline functions are not inside the library" +) +def test_calculate_z(): pass @@ -79,48 +58,37 @@ def calculate_z(r, p): TESTDATA = [ - { - 'r': np.linspace(1, 2, 3, dtype=np.float64), - }, - { - 'r': np.linspace(0, 1, 3), - }, - { - 'r': np.linspace(1, 2, 10, dtype=np.float64) - }, - ] - - -@pytest.fixture( - scope='function', - params=TESTDATA) + { + "r": np.linspace(1, 2, 3, dtype=np.float64), + }, + { + "r": np.linspace(0, 1, 3), + }, + {"r": np.linspace(1, 2, 10, dtype=np.float64)}, +] + + +@pytest.fixture(scope="function", params=TESTDATA) def formal_integral_model(request, model): - r = request.param['r'] + r = request.param["r"] model.no_of_shells_i = r.shape[0] - 1 - model.inverse_time_explosion = c.c.cgs.value + model.time_explosion = 1 / c.c.cgs.value model.r_outer_i.contents = as_ctypes(r[1:]) model.r_inner_i.contents = as_ctypes(r[:-1]) return model -@pytest.mark.parametrize( - 'p', - [0, 0.5, 1] - ) -def test_populate_z_photosphere(clib, formal_integral_model, p): - ''' +@pytest.mark.xfail(reason="not implemented") +@pytest.mark.parametrize("p", [0, 0.5, 1]) +def test_populate_z_photosphere(formal_integral_model, p): + """ Test the case where p < r[0] That means we 'hit' all shells from inside to outside. - ''' - func = clib.populate_z - func.restype = ctypes.c_int64 - func.argtypes = [ - ctypes.POINTER(StorageModel), # storage - c_double, # p - ndpointer(dtype=np.float64), # oz - ndpointer(dtype=np.int64) # oshell_id - ] - + """ + integrator = formal_integral.FormalIntegrator( + formal_integral_model, None, None + ) + func = integrator.populate_z size = formal_integral_model.no_of_shells_i r_inner = as_array(formal_integral_model.r_inner_i, (size,)) r_outer = as_array(formal_integral_model.r_outer_i, (size,)) @@ -129,49 +97,30 @@ def test_populate_z_photosphere(clib, formal_integral_model, p): oz = np.zeros_like(r_inner) oshell_id = np.zeros_like(oz, dtype=np.int64) - N = func( - formal_integral_model, - p, - oz, - oshell_id - ) + N = func(p, oz, oshell_id) assert N == size - ntest.assert_allclose( - oshell_id, - np.arange(0, size, 1) - ) + ntest.assert_allclose(oshell_id, np.arange(0, size, 1)) - ntest.assert_allclose( - oz, - 1 - calculate_z(r_outer, p), - atol=1e-5 - ) + ntest.assert_allclose(oz, 1 - calculate_z(r_outer, p), atol=1e-5) -@pytest.mark.parametrize( - 'p', - [1e-5, 0.5, 0.99, 1] - ) -def test_populate_z_shells(clib, formal_integral_model, p): - ''' +@pytest.mark.parametrize("p", [1e-5, 0.5, 0.99, 1]) +def test_populate_z_shells(formal_integral_model, p): + """ Test the case where p > r[0] - ''' - func = clib.populate_z - func.restype = ctypes.c_int64 - func.argtypes = [ - ctypes.POINTER(StorageModel), # storage - c_double, # p - ndpointer(dtype=np.float64), # oz - ndpointer(dtype=np.int64) # oshell_id - ] + """ + integrator = formal_integral.FormalIntegrator( + formal_integral_model, None, None + ) + func = integrator.populate_z size = formal_integral_model.no_of_shells_i r_inner = as_array(formal_integral_model.r_inner_i, (size,)) r_outer = as_array(formal_integral_model.r_outer_i, (size,)) p = r_inner[0] + (r_outer[-1] - r_inner[0]) * p - idx = np.searchsorted(r_outer, p, side='right') + idx = np.searchsorted(r_outer, p, side="right") oz = np.zeros(size * 2) oshell_id = np.zeros_like(oz, dtype=np.int64) @@ -183,56 +132,40 @@ def test_populate_z_shells(clib, formal_integral_model, p): expected_oshell_id = np.zeros_like(oshell_id) # Calculated way to determine which shells get hit - expected_oshell_id[:expected_N] = np.abs( - np.arange(0.5, expected_N, 1) - offset) - 0.5 + idx + expected_oshell_id[:expected_N] = ( + np.abs(np.arange(0.5, expected_N, 1) - offset) - 0.5 + idx + ) expected_oz[0:offset] = 1 + calculate_z( - r_outer[np.arange(size, idx, -1) - 1], - p) + r_outer[np.arange(size, idx, -1) - 1], p + ) expected_oz[offset:expected_N] = 1 - calculate_z( - r_outer[np.arange(idx, size, 1)], - p) + r_outer[np.arange(idx, size, 1)], p + ) - N = func( - formal_integral_model, - p, - oz, - oshell_id - ) + N = func(p, oz, oshell_id) assert N == expected_N - ntest.assert_allclose( - oshell_id, - expected_oshell_id - ) + ntest.assert_allclose(oshell_id, expected_oshell_id) - ntest.assert_allclose( - oz, - expected_oz, - atol=1e-5 - ) + ntest.assert_allclose(oz, expected_oz, atol=1e-5) @pytest.mark.parametrize( - 'N', - [ - 100, - 1000, - 10000, - ]) -def test_calculate_p_values(clib, N): - r = 1. - func = clib.calculate_p_values - func.argtypes = [ - c_double, - c_int, - ndpointer(dtype=np.float64) - ] - expected = r/(N - 1) * np.arange(0, N, dtype=np.float64) + "N", + [ + 100, + 1000, + 10000, + ], +) +def test_calculate_p_values(N): + r = 1.0 + func = formal_integral.calculate_p_values + + expected = r / (N - 1) * np.arange(0, N, dtype=np.float64) actual = np.zeros_like(expected, dtype=np.float64) func(r, N, actual) - ntest.assert_allclose( - actual, - expected) + ntest.assert_allclose(actual, expected) diff --git a/tardis/montecarlo/tests/test_montecarlo.py b/tardis/montecarlo/tests/test_montecarlo.py new file mode 100644 index 00000000000..aa050fd3477 --- /dev/null +++ b/tardis/montecarlo/tests/test_montecarlo.py @@ -0,0 +1,1077 @@ +""" +Unit tests for methods in `tardis/montecarlo/src/cmontecarlo.c`. +* `ctypes` library is used to wrap C methods and expose them to python. + + +Probable Reasons for Failing Tests: +----------------------------------- + +1. Change made in C struct declarations: + - Reflect the changes done in C structs, into Python counterparts. + - Check **tardis/montecarlo/struct.py**. + +2. Return type of any method changed: + - Modify the `restype` parameter in the test method here. + - For example: + ``` + clib.rpacket_doppler_factor.restype = c_double + ``` + +3. Underlying logic modified: + - Check whether the changes made in C method are logically correct. + - If the changes made were correct and necessary, update the corresponding + test case. + + +General Test Design Procedure: +------------------------------ + +Please follow this design procedure while adding a new test: + +1. Parametrization as per desire of code coverage. + - C tests have different flows controlled by conditional statements. + Parameters checked in conditions can be provided in different testcases. + - Keep consistency with variable names as (in order): + - `packet_params` + - `model_params` + - `expected_params` (`expected` if only one value to be asserted.) + - Suggested variable names can be compromised if readability of the test + increases. + +2. Test Method body: + - Keep name as `test_` + `(name of C method)`. + - Refer to method `test_rpacket_doppler_factor` below for description. +""" + + +import os +import pytest +import numpy as np +import pandas as pd +import tardis.montecarlo.formal_integral as formal_integral +import tardis.montecarlo.montecarlo_numba.r_packet as r_packet +import tardis.montecarlo.montecarlo_configuration as mc +from tardis import constants as const +from tardis.montecarlo.montecarlo_numba.numba_interface import Estimators +from tardis.montecarlo.montecarlo_numba import macro_atom + +C_SPEED_OF_LIGHT = const.c.to("cm/s").value + +pytestmark = pytest.mark.skip(reason="Port from C to numba") + + +from ctypes import ( + CDLL, + byref, + c_uint, + c_int64, + c_double, + c_ulong, + c_void_p, + cast, + POINTER, + pointer, + CFUNCTYPE, +) + +from numpy.testing import ( + assert_equal, + assert_almost_equal, + assert_array_equal, + assert_allclose, +) + +from tardis import __path__ as path +from tardis.montecarlo.struct import ( + RPacket, + StorageModel, + RKState, + PhotoXsect1level, + TARDIS_ERROR_OK, + TARDIS_ERROR_BOUNDS_ERROR, + TARDIS_ERROR_COMOV_NU_LESS_THAN_NU_LINE, + TARDIS_PACKET_STATUS_IN_PROCESS, + TARDIS_PACKET_STATUS_EMITTED, + TARDIS_PACKET_STATUS_REABSORBED, + CONTINUUM_OFF, + CONTINUUM_ON, + INVERSE_C, + BoundFreeTreatment, +) + + +@pytest.fixture(scope="module") +def continuum_compare_data_fname(): + fname = "continuum_compare_data.hdf" + return os.path.join(path[0], "montecarlo", "tests", "data", fname) + + +@pytest.fixture(scope="module") +def continuum_compare_data(continuum_compare_data_fname, request): + compare_data = pd.HDFStore(continuum_compare_data_fname, mode="r") + + def fin(): + compare_data.close() + + request.addfinalizer(fin) + + return compare_data + + +@pytest.fixture(scope="function") +def expected_ff_emissivity(continuum_compare_data): + emissivities = continuum_compare_data["ff_emissivity"] + + def ff_emissivity(t_electron): + emissivity = emissivities[t_electron] + nu_bins = emissivity["nu_bins"].values + emissivity_value = emissivity["emissivity"].dropna().values + + return nu_bins, emissivity_value + + return ff_emissivity + + +@pytest.fixture(scope="module") +def get_rkstate(continuum_compare_data): + data = continuum_compare_data["z2rkstate_key"] + pos_data = continuum_compare_data["z2rkstate_pos"] + + def z2rkstate(z_random): + key = (c_ulong * 624)(*data.loc[z_random].values) + pos = pos_data.loc[z_random] + return RKState(key=key, pos=pos, has_gauss=0, gauss=0.0) + + return z2rkstate + + +@pytest.fixture(scope="function") +def model_w_edges(ion_edges, model): + photo_xsect = (POINTER(PhotoXsect1level) * len(ion_edges))() + + for i, edge in enumerate(ion_edges): + x_sect_1level = PhotoXsect1level() + for key, value in edge.items(): + if key in ["nu", "x_sect"]: + value = (c_double * len(value))(*value) + setattr(x_sect_1level, key, value) + photo_xsect[i] = pointer(x_sect_1level) + + no_of_edges = len(ion_edges) + continuum_list_nu = (c_double * no_of_edges)( + *[edge["nu"][0] for edge in ion_edges] + ) + + model.photo_xsect = photo_xsect + model.continuum_list_nu = continuum_list_nu + model.no_of_edges = no_of_edges + + estimator_size = model.no_of_shells * no_of_edges + estims = [ + "photo_ion_estimator", + "stim_recomb_estimator", + "bf_heating_estimator", + "stim_recomb_cooling_estimator", + ] + for estimator in estims: + setattr( + model, estimator, (c_double * estimator_size)(*[0] * estimator_size) + ) + + model.photo_ion_estimator_statistics = (c_int64 * estimator_size)( + *[0] * estimator_size + ) + return model + + +@pytest.fixture(scope="module") +def ion_edges(): + return [ + { + "nu": [4.0e14, 4.1e14, 4.2e14, 4.3e14], + "x_sect": [1.0, 0.9, 0.8, 0.7], + "no_of_points": 4, + }, + { + "nu": [3.0e14, 3.1e14, 3.2e14, 3.3e14, 3.4e14], + "x_sect": [1.0, 0.9, 0.8, 0.7, 0.6], + "no_of_points": 5, + }, + { + "nu": [2.8e14, 3.0e14, 3.2e14, 3.4e14], + "x_sect": [2.0, 1.8, 1.6, 1.4], + "no_of_points": 4, + }, + ] + + +@pytest.fixture(scope="module") +def mock_sample_nu(): + SAMPLE_NUFUNC = CFUNCTYPE( + c_double, POINTER(RPacket), POINTER(StorageModel), POINTER(RKState) + ) + + def sample_nu_simple(packet, model, mt_state): + return packet.contents.nu + + return SAMPLE_NUFUNC(sample_nu_simple) + + +@pytest.fixture(scope="function") +def model_3lvlatom(model): + model.line2macro_level_upper = (c_int64 * 3)(*[2, 1, 2]) + model.macro_block_references = (c_int64 * 3)(*[0, 2, 5]) + + transition_probabilities = [ + 0.0, + 0.0, + 0.75, + 0.25, + 0.0, + 0.25, + 0.5, + 0.25, + 0.0, # shell_id = 0 + 0.0, + 0.0, + 1.00, + 0.00, + 0.0, + 0.00, + 0.0, + 1.00, + 0.0, # shell_id = 1 + ] + + nd = len(transition_probabilities) // 2 + model.transition_type = (c_int64 * nd)(*[1, 1, -1, 1, 0, 0, -1, -1, 0]) + model.destination_level_id = (c_int64 * nd)(*[1, 2, 0, 2, 0, 1, 1, 0, 0]) + model.transition_line_id = (c_int64 * nd)(*[0, 1, 1, 2, 1, 2, 2, 0, 0]) + + model.transition_probabilities_nd = c_int64(nd) + model.transition_probabilities = (c_double * (nd * 2))( + *transition_probabilities + ) + + model.last_line_interaction_out_id = (c_int64 * 1)(*[-5]) + + return model + + +def d_cont_setter(d_cont, model, packet): + model.inverse_electron_densities[packet.current_shell_id] = c_double(1.0) + model.inverse_sigma_thomson = c_double(1.0) + packet.tau_event = c_double(d_cont) + + +def d_line_setter(d_line, model, packet): + packet.mu = c_double(0.0) + scale = d_line * 1e1 + model.time_explosion = c_double(INVERSE_C * scale) + packet.nu = c_double(1.0) + nu_line = 1.0 - d_line / scale + packet.nu_line = c_double(nu_line) + + +def d_boundary_setter(d_boundary, model, packet): + packet.mu = c_double(1e-16) + r_outer = 2.0 * d_boundary + model.r_outer[packet.current_shell_id] = r_outer + + r = np.sqrt(r_outer ** 2 - d_boundary ** 2) + packet.r = r + + +""" +Important Tests: +---------------- +The tests written further (till next block comment is encountered) have been +categorized as important tests, these tests correspond to methods which are +relatively old and stable code. +""" + + +@pytest.mark.parametrize( + ["x", "x_insert", "imin", "imax", "expected_params"], + [ + ( + [5.0, 4.0, 3.0, 1.0], + 2.0, + 0, + 3, + {"result": 2, "ret_val": TARDIS_ERROR_OK}, + ), + ( + [5.0, 4.0, 3.0, 2.0], + 0.0, + 0, + 3, + {"result": 0, "ret_val": TARDIS_ERROR_BOUNDS_ERROR}, + ), + ], +) +def test_reverse_binary_search(x, x_insert, imin, imax, expected_params): + # x = (c_double * (imax - imin + 1))(*x) + obtained_result = 0 + + obtained_tardis_error = formal_integral.reverse_binary_search( + x, x_insert, imin, imax, obtained_result + ) + + assert obtained_result.value == expected_params["result"] + assert obtained_tardis_error == expected_params["ret_val"] + + +@pytest.mark.parametrize( + ["nu", "nu_insert", "number_of_lines", "expected_params"], + [ + ( + [0.5, 0.4, 0.3, 0.1], + 0.2, + 4, + {"result": 3, "ret_val": TARDIS_ERROR_OK}, + ), + ( + [0.5, 0.4, 0.3, 0.2], + 0.1, + 4, + {"result": 4, "ret_val": TARDIS_ERROR_OK}, + ), + ( + [0.4, 0.3, 0.2, 0.1], + 0.5, + 4, + {"result": 0, "ret_val": TARDIS_ERROR_OK}, + ), + ], +) +def test_line_search(nu, nu_insert, number_of_lines, expected_params): + # nu = (c_double * number_of_lines)(*nu) + obtained_result = 0 + + obtained_tardis_error = formal_integral.line_search( + nu, nu_insert, number_of_lines, obtained_result + ) + + assert obtained_result.value == expected_params["result"] + assert obtained_tardis_error == expected_params["ret_val"] + + +@pytest.mark.parametrize( + ["x", "x_insert", "imin", "imax", "expected_params"], + [ + ( + [2.0, 4.0, 6.0, 7.0], + 5.0, + 0, + 3, + {"result": 2, "ret_val": TARDIS_ERROR_OK}, + ), + ( + [2.0, 3.0, 5.0, 7.0], + 8.0, + 0, + 3, + {"result": 0, "ret_val": TARDIS_ERROR_BOUNDS_ERROR}, + ), + ( + [2.0, 4.0, 6.0, 7.0], + 4.0, + 0, + 3, + {"result": 0, "ret_val": TARDIS_ERROR_OK}, + ), + ], +) +def test_binary_search(x, x_insert, imin, imax, expected_params): + obtained_result = 0 + + obtained_tardis_error = formal_integral.binary_search( + x, x_insert, imin, imax, obtained_result + ) + + assert obtained_result.value == expected_params["result"] + assert obtained_tardis_error == expected_params["ret_val"] + + +@pytest.mark.parametrize( + ["mu", "r", "inv_t_exp", "expected"], + [ + (0.3, 7.5e14, 1 / 5.2e7, 0.9998556693818854), + (-0.3, 0, 1 / 2.6e7, 1.0), + (0, 1, 1 / 2.6e7, 1.0), + ], +) +def test_get_doppler_factor(mu, r, inv_t_exp, expected, packet, model): + # Set the params from test cases here + # TODO: add relativity tests + time_explosion = 1 / inv_t_exp + + # Perform any other setups just before this, they can be additional calls + # to other methods or introduction of some temporary variables + + obtained = r_packet.get_doppler_factor(r, mu, time_explosion) + + # Perform required assertions + assert_almost_equal(obtained, expected) + + +@pytest.mark.parametrize(["mu", "r", "inv_t_exp"], [(-0.3, 5, 1e10)]) +def test_unphysical_doppler_factor(mu, r, inv_t_exp, packet, model): + # Set the params from test cases here + # TODO: add relativity tests + time_explosion = 1 / inv_t_exp + + # Perform any other setups just before this, they can be additional calls + # to other methods or introduction of some temporary variables + with pytest.raises(r_packet.SuperluminalError): + obtained = r_packet.get_doppler_factor(r, mu, time_explosion) + + +@pytest.mark.parametrize( + ["mu", "r", "inv_t_exp", "expected"], + [ + (0.3, 7.5e14, 1 / 5.2e7, 1 / 0.9998556693818854), + (-0.3, 0, 1 / 2.6e7, 1.0), + (0, 1, 1 / 2.6e7, 1.0), + ], +) +def test_get_inverse_doppler_factor(mu, r, inv_t_exp, expected, packet, model): + # Set the params from test cases here + # TODO: add relativity tests + time_explosion = 1 / inv_t_exp + + # Perform any other setups just before this, they can be additional calls + # to other methods or introduction of some temporary variables + + obtained = r_packet.get_inverse_doppler_factor(r, mu, time_explosion) + + # Perform required assertions + assert_almost_equal(obtained, expected) + + +@pytest.mark.parametrize(["mu", "r", "inv_t_exp"], [(-0.3, 5, 1e10)]) +def test_unphysical_inverse_doppler_factor(mu, r, inv_t_exp, packet, model): + # Set the params from test cases here + # TODO: add relativity tests + time_explosion = 1 / inv_t_exp + + # Perform any other setups just before this, they can be additional calls + # to other methods or introduction of some temporary variables + with pytest.raises(r_packet.SuperluminalError): + obtained = r_packet.get_inverse_doppler_factor(r, mu, time_explosion) + + +@pytest.mark.parametrize( + ["mu", "r", "inv_t_exp", "expected"], + [ + (-0.3, 10000000, 0.001, 1.0000001000692842), + (-0.3, 0, 1 / 2.6e7, 1.0), + (0, 1, 1 / 2.6e7, 1.0), + ], +) +def test_get_doppler_factor_full_relativity( + mu, r, inv_t_exp, expected, packet, model +): + # Set the params from test cases here + # TODO: add relativity tests + mc.full_relativity = True + time_explosion = 1 / inv_t_exp + + # Perform any other setups just before this, they can be additional calls + # to other methods or introduction of some temporary variables + + obtained = r_packet.get_doppler_factor(r, mu, time_explosion) + mc.full_relativity = False + # Perform required assertions + assert_almost_equal(obtained, expected) + + +@pytest.mark.parametrize( + ["mu", "r", "inv_t_exp", "expected"], + [ + (-0.3, 10000000, 0.001, 0.999999899930827), + (-0.3, 0, 1 / 2.6e7, 1.0), + (0, 1, 1 / 2.6e7, 1.0), + ], +) +def test_get_inverse_doppler_factor_full_relativity( + mu, r, inv_t_exp, expected, packet, model +): + # Set the params from test cases here + # TODO: add relativity tests + mc.full_relativity = True + time_explosion = 1 / inv_t_exp + + # Perform any other setups just before this, they can be additional calls + # to other methods or introduction of some temporary variables + + obtained = r_packet.get_inverse_doppler_factor(r, mu, time_explosion) + mc.full_relativity = False + # Perform required assertions + assert_almost_equal(obtained, expected) + + +def test_get_random_mu_different_output(): + """ + Ensure that different calls results + """ + output1 = r_packet.get_random_mu() + output2 = r_packet.get_random_mu() + assert output1 != output2 + + +def test_get_random_mu_different_output(): + """ + Ensure that different calls results + """ + output1 = r_packet.get_random_mu() + output2 = r_packet.get_random_mu() + assert output1 != output2 + + +@pytest.mark.parametrize( + ["mu", "r", "time_explosion"], [(1, C_SPEED_OF_LIGHT, 1)] +) +def test_angle_ab_LF_to_CMF_diverge(mu, r, time_explosion, packet): + """ + This equation should diverage as costheta --> 1 and beta --> 1 + """ + packet = r_packet.RPacket(packet.r, packet.mu, packet.nu, packet.energy) + packet.r = r + packet.mu = mu + with pytest.raises(ZeroDivisionError): + obtained = r_packet.angle_aberration_LF_to_CMF( + packet, time_explosion, mu + ) + + +@pytest.mark.parametrize( + ["mu", "r", "time_explosion"], [(0.3, 1e7, 1e10), (-0.3, 1e7, 1e11)] +) +def test_both_angle_aberrations(mu, r, time_explosion, packet): + """ + The angle aberration functions should be the functional inverse of one + another. + """ + packet = r_packet.RPacket(r, packet.mu, packet.nu, packet.energy) + packet.r = r + obtained_mu = r_packet.angle_aberration_LF_to_CMF( + packet, time_explosion, mu + ) + inverse_obtained_mu = r_packet.angle_aberration_CMF_to_LF( + packet, time_explosion, obtained_mu + ) + assert_almost_equal(inverse_obtained_mu, mu) + + +@pytest.mark.parametrize( + ["mu", "r", "time_explosion"], [(0.3, 7.5e14, 5.2e5), (-0.3, 7.5e14, 5.2e5)] +) +def test_both_angle_aberrations_inverse(mu, r, time_explosion, packet): + """ + The angle aberration functions should be the functional inverse of one + another. + """ + packet = r_packet.RPacket(r, packet.mu, packet.nu, packet.energy) + packet.r = r + obtained_mu = r_packet.angle_aberration_CMF_to_LF( + packet, time_explosion, mu + ) + inverse_obtained_mu = r_packet.angle_aberration_LF_to_CMF( + packet, time_explosion, obtained_mu + ) + assert_almost_equal(inverse_obtained_mu, mu) + + +@pytest.mark.parametrize( + ["current_shell_id", "delta_shell", "no_of_shells"], + [(132, 11, 132), (132, 1, 133), (132, 2, 133)], +) +def test_move_packet_across_shell_boundary_emitted( + packet, current_shell_id, delta_shell, no_of_shells +): + packet = r_packet.RPacket(packet.r, packet.mu, packet.nu, packet.energy) + packet.current_shell_id = current_shell_id + r_packet.move_packet_across_shell_boundary( + packet, delta_shell, no_of_shells + ) + assert packet.status == r_packet.PacketStatus.EMITTED + + +@pytest.mark.parametrize( + ["current_shell_id", "delta_shell", "no_of_shells"], + [(132, -133, 132), (132, -133, 133), (132, -1e9, 133)], +) +def test_move_packet_across_shell_boundary_reabsorbed( + packet, current_shell_id, delta_shell, no_of_shells +): + packet = r_packet.RPacket(packet.r, packet.mu, packet.nu, packet.energy) + packet.current_shell_id = current_shell_id + r_packet.move_packet_across_shell_boundary( + packet, delta_shell, no_of_shells + ) + assert packet.status == r_packet.PacketStatus.REABSORBED + + +@pytest.mark.parametrize( + ["current_shell_id", "delta_shell", "no_of_shells"], + [(132, -1, 199), (132, 0, 132), (132, 20, 154)], +) +def test_move_packet_across_shell_boundary_increment( + packet, current_shell_id, delta_shell, no_of_shells +): + packet = r_packet.RPacket(packet.r, packet.mu, packet.nu, packet.energy) + packet.current_shell_id = current_shell_id + r_packet.move_packet_across_shell_boundary( + packet, delta_shell, no_of_shells + ) + assert packet.current_shell_id == current_shell_id + delta_shell + + +@pytest.mark.parametrize( + ["distance_trace", "time_explosion", "mu", "r"], + [(0, 1, 0, 0), (0, 1, 1, 0), (0, 1, 0, 1)], +) +def test_packet_energy_limit_one(packet, distance_trace, time_explosion, mu, r): + initial_energy = packet.energy + packet = r_packet.RPacket(r, mu, packet.nu, packet.energy) + new_energy = r_packet.calc_packet_energy( + packet, distance_trace, time_explosion + ) + assert new_energy == initial_energy + + +@pytest.mark.parametrize( + ["packet_params", "expected_params"], + [ + ({"mu": 0.3, "r": 7.5e14}, {"d_boundary": 259376919351035.88}), + ({"mu": -0.3, "r": 7.5e13}, {"d_boundary": -664987228972291.5}), + ({"mu": -0.3, "r": 7.5e14}, {"d_boundary": 709376919351035.9}), + ], +) +def test_compute_distance2boundary( + packet_params, expected_params, packet, model +): + mu = packet_params["mu"] + r = packet_params["r"] + + d_boundary = r_packet.calculate_distance_boundary( + r, mu, model.r_inner[0], model.r_outer[0] + ) + + assert_almost_equal(d_boundary[0], expected_params["d_boundary"]) + + +# +# +# TODO: split this into two tests - one to assert errors and other for d_line +@pytest.mark.parametrize( + ["packet_params", "expected_params"], + [ + ( + {"nu_line": 0.1, "next_line_id": 0, "last_line": 1}, + {"tardis_error": None, "d_line": 1e99}, + ), + ( + {"nu_line": 0.2, "next_line_id": 1, "last_line": 0}, + {"tardis_error": None, "d_line": 7.792353908000001e17}, + ), + ( + {"nu_line": 0.5, "next_line_id": 1, "last_line": 0}, + {"tardis_error": r_packet.MonteCarloException, "d_line": 0.0}, + ), + ( + {"nu_line": 0.6, "next_line_id": 0, "last_line": 0}, + {"tardis_error": r_packet.MonteCarloException, "d_line": 0.0}, + ), + ], +) +def test_compute_distance2line(packet_params, expected_params, packet, model): + packet = r_packet.RPacket(packet.r, packet.mu, packet.nu, packet.energy) + nu_line = packet_params["nu_line"] + # packet.next_line_id = packet_params['next_line_id'] + # packet.last_line = packet_params['last_line'] + + time_explosion = model.time_explosion + + doppler_factor = r_packet.get_doppler_factor( + packet.r, packet.mu, time_explosion + ) + comov_nu = packet.nu * doppler_factor + + d_line = 0 + obtained_tardis_error = None + try: + d_line = r_packet.calculate_distance_line( + packet, comov_nu, nu_line, time_explosion + ) + except r_packet.MonteCarloException: + obtained_tardis_error = r_packet.MonteCarloException + + assert_almost_equal(d_line, expected_params["d_line"]) + assert obtained_tardis_error == expected_params["tardis_error"] + + +# @pytest.mark.parametrize( +# ['packet_params', 'expected_params'], +# [({'virtual_packet': 0}, +# {'chi_cont': 6.652486e-16, 'd_cont': 4.359272608766106e+28}), +# +# ({'virtual_packet': 1}, +# {'chi_cont': 6.652486e-16, 'd_cont': 1e+99})] +# ) +# def test_compute_distance2continuum(clib, packet_params, expected_params, packet, model): +# packet.virtual_packet = packet_params['virtual_packet'] +# +# clib.compute_distance2continuum(byref(packet), byref(model)) +# +# assert_almost_equal(packet.chi_cont, expected_params['chi_cont']) +# assert_almost_equal(packet.d_cont, expected_params['d_cont']) + + +@pytest.mark.parametrize("full_relativity", [1, 0]) +@pytest.mark.parametrize( + ["packet_params", "expected_params"], + [ + ( + {"nu": 0.4, "mu": 0.3, "energy": 0.9, "r": 7.5e14}, + { + "mu": 0.3120599529139568, + "r": 753060422542573.9, + "j": 8998701024436.969, + "nubar": 3598960894542.354, + }, + ), + ( + {"nu": 0.6, "mu": -0.5, "energy": 0.5, "r": 8.1e14}, + { + "mu": -0.4906548373534084, + "r": 805046582503149.2, + "j": 5001298975563.031, + "nubar": 3001558973156.1387, + }, + ), + ], +) +def test_move_packet( + packet_params, expected_params, packet, model, full_relativity +): + distance = 1e13 + packet.nu = packet_params["nu"] + packet.mu = packet_params["mu"] + packet.energy = packet_params["energy"] + packet.r = packet_params["r"] + # model.full_relativity = full_relativity + mc.full_relativity = full_relativity + + doppler_factor = r_packet.get_doppler_factor( + packet.r, packet.mu, model.time_explosion + ) + numba_estimator = Estimators( + packet_params["j"], packet_params["nu_bar"], 0, 0 + ) + r_packet.move_r_packet( + packet, distance, model.time_explosion, numba_estimator + ) + + assert_almost_equal(packet.mu, expected_params["mu"]) + assert_almost_equal(packet.r, expected_params["r"]) + + expected_j = expected_params["j"] + expected_nubar = expected_params["nubar"] + if full_relativity: + expected_j *= doppler_factor + expected_nubar *= doppler_factor + + mc.full_relativity = False + + assert_allclose( + numba_estimator.j_estimator[packet.current_shell_id], + expected_j, + rtol=5e-7, + ) + assert_allclose( + numba_estimator.nu_bar_estimator[packet.current_shell_id], + expected_nubar, + rtol=5e-7, + ) + + +# @pytest.mark.continuumtest +# @pytest.mark.parametrize( +# ['packet_params', 'j_blue_idx', 'expected'], +# [({'nu': 0.30, 'energy': 0.30}, 0, 1.0), +# ({'nu': 0.20, 'energy': 1.e5}, 0, 5e5), +# ({'nu': 2e15, 'energy': 0.50}, 1, 2.5e-16), +# ({'nu': 0.40, 'energy': 1e-7}, 1, 2.5e-7)], +# ) +# def test_increment_j_blue_estimator_full_relativity(packet_params, +# j_blue_idx, expected, +# packet, model): +# packet.nu = packet_params['nu'] +# packet.energy = packet_params['energy'] +# model.full_relativity = True +# +# r_packet.increment_j_blue_estimator(byref(packet), byref(model), +# c_double(packet.d_line), +# c_int64(j_blue_idx)) +# +# assert_almost_equal(model.line_lists_j_blues[j_blue_idx], expected) +# +# +# @pytest.mark.parametrize( +# ['packet_params', 'cur_line_id', 'expected'], +# [({'nu': 0.1, 'mu': 0.3, 'r': 7.5e14}, 0, 8.998643292289723), +# ({'nu': 0.2, 'mu': -.3, 'r': 7.7e14}, 0, 4.499971133976377), +# ({'nu': 0.5, 'mu': 0.5, 'r': 7.9e14}, 1, 0.719988453650551), +# ({'nu': 0.6, 'mu': -.5, 'r': 8.1e14}, 1, 0.499990378058792)] +# ) +# def test_increment_j_blue_estimator(packet_params, cur_line_id, expected, packet): +# +# numba_interface +# packet = r_packet.RPacket(packet_params['r'], +# packet_params['mu'], +# packet_params['nu'], +# packet.energy) +# +# r_packet.compute_distance2line(byref(packet), byref(model)) +# r_packet.move_r_packet(packet, +# distance, +# model.time_explosion, +# numba_estimator) +# r_packet.move_packet(byref(packet), byref(model), c_double(1.e13)) +# r_packet.update_line_estimators(estimators, r_packet, +# cur_line_id, +# model.distance_trace, +# model.time_explosion) +# +# assert_almost_equal(model.line_lists_j_blues[j_blue_idx], expected) + + +@pytest.mark.parametrize( + ["packet_params", "expected_params"], + [ + ( + {"nu": 0.4, "mu": 0.3, "energy": 0.9, "r": 7.5e14}, + {"nu": 0.39974659819356556, "energy": 0.8994298459355226}, + ), + ( + {"nu": 0.6, "mu": -0.5, "energy": 0.5, "r": 8.1e14}, + {"nu": 0.5998422620533325, "energy": 0.4998685517111104}, + ), + ], +) +def test_montecarlo_thomson_scatter( + clib, packet_params, expected_params, packet, model, mt_state +): + packet.nu = packet_params["nu"] + packet.mu = packet_params["mu"] + packet.energy = packet_params["energy"] + packet.r = packet_params["r"] + + clib.montecarlo_thomson_scatter( + byref(packet), byref(model), c_double(1.0e13), byref(mt_state) + ) + + assert_almost_equal(packet.nu, expected_params["nu"]) + assert_almost_equal(packet.energy, expected_params["energy"]) + + +@pytest.mark.parametrize( + ["packet_params", "expected_params"], + # TODO: Add scientifically sound test cases. + [ + ( + {"virtual_packet": 1, "tau_event": 2.9e13, "last_line": 0}, + {"tau_event": 2.9e13, "next_line_id": 2}, + ), + ( + {"virtual_packet": 0, "tau_event": 2.9e13, "last_line": 0}, + {"tau_event": 2.9e13, "next_line_id": 2}, + ), + ( + {"virtual_packet": 0, "tau_event": 2.9e13, "last_line": 0}, + {"tau_event": 2.9e13, "next_line_id": 2}, + ), + ], +) +def test_montecarlo_line_scatter( + clib, packet_params, expected_params, packet, model, mt_state +): + packet.virtual_packet = packet_params["virtual_packet"] + packet.tau_event = packet_params["tau_event"] + packet.last_line = packet_params["last_line"] + + clib.montecarlo_line_scatter( + byref(packet), byref(model), c_double(1.0e13), byref(mt_state) + ) + + assert_almost_equal(packet.tau_event, expected_params["tau_event"]) + assert_almost_equal(packet.next_line_id, expected_params["next_line_id"]) + + +@pytest.mark.parametrize( + ["z_random", "packet_params", "expected"], + [ + ( + 0.22443743797312765, + {"activation_level": 1, "shell_id": 0}, + 1, + ), # Direct deactivation + ( + 0.78961460371187597, # next z_random = 0.818455414618 + {"activation_level": 1, "shell_id": 0}, + 0, + ), # Upwards jump, then deactivation + ( + 0.22443743797312765, # next z_random = 0.545678896748 + {"activation_level": 2, "shell_id": 0}, + 1, + ), # Downwards jump, then deactivation + ( + 0.765958602560605, # next z_random = 0.145914243888, 0.712382380384 + {"activation_level": 1, "shell_id": 0}, + 1, + ), # Upwards jump, downwards jump, then deactivation + (0.22443743797312765, {"activation_level": 2, "shell_id": 1}, 0), + ], # Direct deactivation +) +def test_macro_atom(packet, z_random, packet_params, get_rkstate, expected): + packet.macro_atom_activation_level = packet_params["activation_level"] + packet = r_packet.RPacket( + r=packet.r, mu=packet.mu, nu=packet.nu, energy=packet.energy + ) + packet.current_shell_id = packet_params["shell_id"] + rkstate = get_rkstate(z_random) + + macro_atom.macro_atom(packet, numba_plasma) + obtained_line_id = model_3lvlatom.last_line_interaction_out_id[packet.id] + + assert_equal(obtained_line_id, expected) + + +""" +Simple Tests: +---------------- +These test check very simple pieces of code still work. +""" + + +@pytest.mark.parametrize( + ["packet_params", "line_idx", "expected"], + [ + ({"energy": 0.0}, 0, 0), + ({"energy": 1.0}, 1, 1), + ({"energy": 0.5}, 2, 1.5), + ], +) +@pytest.mark.skipif(True, reason="Needs rewrite to be relevant.") +def test_increment_Edotlu_estimator( + clib, packet_params, line_idx, expected, packet, model +): + packet.energy = packet_params["energy"] + + clib.increment_Edotlu_estimator( + byref(packet), byref(model), c_int64(line_idx) + ) + + assert_almost_equal(model.line_lists_Edotlu[line_idx], expected) + + +""" +Continuum Tests: +---------------- +The tests written further (till next block comment is encountered) are for the +methods related to continuum interactions. +""" + + +@pytest.mark.continuumtest +@pytest.mark.parametrize( + ["mu", "r", "inv_t_exp", "full_relativity"], + [ + (0.8, 7.5e14, 1 / 5.2e5, 1), + (-0.7, 7.5e14, 1 / 5.2e5, 1), + (0.3, 7.5e14, 1 / 2.2e5, 1), + (0.0, 7.5e14, 1 / 2.2e5, 1), + (-0.7, 7.5e14, 1 / 5.2e5, 0), + ], +) +def test_frame_transformations(packet, mu, r, inv_t_exp, full_relativity): + packet = r_packet.RPacket(r=r, mu=mu, energy=packet.energy, nu=packet.nu) + mc.full_relativity = bool(full_relativity) + mc.full_relativity = full_relativity + + inverse_doppler_factor = r_packet.get_inverse_doppler_factor( + r, mu, 1 / inv_t_exp + ) + r_packet.angle_aberration_CMF_to_LF(packet, 1 / inv_t_exp, packet.mu) + + doppler_factor = r_packet.get_doppler_factor(r, mu, 1 / inv_t_exp) + mc.full_relativity = False + + assert_almost_equal(doppler_factor * inverse_doppler_factor, 1.0) + + +@pytest.mark.continuumtest +@pytest.mark.parametrize( + ["mu", "r", "inv_t_exp"], + [ + (0.8, 7.5e14, 1 / 5.2e5), + (-0.7, 7.5e14, 1 / 5.2e5), + (0.3, 7.5e14, 1 / 2.2e5), + (0.0, 7.5e14, 1 / 2.2e5), + (-0.7, 7.5e14, 1 / 5.2e5), + ], +) +def test_angle_transformation_invariance(packet, model, mu, r, inv_t_exp): + packet = r_packet.RPacket(r, mu, packet.nu, packet.energy) + model.full_relativity = 1 + + mu1 = r_packet.angle_aberration_CMF_to_LF(packet, 1 / inv_t_exp, mu) + mu_obtained = r_packet.angle_aberration_LF_to_CMF( + packet, 1 / inv_t_exp, mu1 + ) + + assert_almost_equal(mu_obtained, mu) + + +@pytest.mark.continuumtest +@pytest.mark.parametrize("full_relativity", [1, 0]) +@pytest.mark.parametrize( + ["mu", "r", "t_exp", "nu", "nu_line"], + [ + (0.8, 7.5e14, 5.2e5, 1.0e15, 9.4e14), + (0.0, 6.3e14, 2.2e5, 6.0e12, 5.8e12), + (1.0, 9.0e14, 2.2e5, 4.0e8, 3.4e8), + (0.9, 9.0e14, 0.5e5, 1.0e15, 4.5e14), + (-0.7, 7.5e14, 5.2e5, 1.0e15, 9.8e14), + (-1.0, 6.3e14, 2.2e5, 6.0e12, 6.55e12), + ], +) +def test_compute_distance2line_relativistic( + mu, r, t_exp, nu, nu_line, full_relativity, packet, runner +): + packet = r_packet.RPacket(r=r, nu=nu, mu=mu, energy=packet.energy) + # packet.nu_line = nu_line + numba_estimator = Estimators( + runner.j_estimator, + runner.nu_bar_estimator, + runner.j_blue_estimator, + runner.Edotlu_estimator, + ) + mc.full_relativity = bool(full_relativity) + + doppler_factor = r_packet.get_doppler_factor(r, mu, t_exp) + comov_nu = packet.nu * doppler_factor + distance = r_packet.calculate_distance_line( + packet, comov_nu, nu_line, t_exp + ) + r_packet.move_r_packet(packet, distance, t_exp, numba_estimator) + + doppler_factor = r_packet.get_doppler_factor(r, mu, t_exp) + comov_nu = packet.nu * doppler_factor + mc.full_relativity = False + + assert_allclose(comov_nu, nu_line, rtol=1e-14) diff --git a/tardis/montecarlo/tests/test_packet_source.py b/tardis/montecarlo/tests/test_packet_source.py index 010348c8d53..6c560fdc1a2 100644 --- a/tardis/montecarlo/tests/test_packet_source.py +++ b/tardis/montecarlo/tests/test_packet_source.py @@ -20,18 +20,19 @@ def packet_unit_test_fpath(tardis_ref_path): def test_bb_packet_sampling(request, tardis_ref_data, packet_unit_test_fpath): bb = BlackBodySimpleSource(2508) + rng = np.random.default_rng(seed=1963) # ref_df = pd.read_hdf('test_bb_sampling.h5') if request.config.getoption("--generate-reference"): ref_bb = pd.read_hdf(packet_unit_test_fpath, key="/blackbody") ref_bb.to_hdf( tardis_ref_data, key="/packet_unittest/blackbody", mode="a" ) + pytest.skip("Reference data was generated during this run.") ref_df = tardis_ref_data["/packet_unittest/blackbody"] - nus = bb.create_blackbody_packet_nus(10000, 100) - mus = bb.create_zero_limb_darkening_packet_mus(100) - unif_energies = bb.create_uniform_packet_energies(100) + nus = bb.create_blackbody_packet_nus(10000, 100, rng) + mus = bb.create_zero_limb_darkening_packet_mus(100, rng) + unif_energies = bb.create_uniform_packet_energies(100, rng) assert np.all(np.isclose(nus, ref_df["nus"])) assert np.all(np.isclose(mus, ref_df["mus"])) assert np.all(np.isclose(unif_energies, ref_df["energies"])) - diff --git a/tardis/montecarlo/tests/test_spectrum.py b/tardis/montecarlo/tests/test_spectrum.py index 2ff064cb122..a16ea419666 100644 --- a/tardis/montecarlo/tests/test_spectrum.py +++ b/tardis/montecarlo/tests/test_spectrum.py @@ -2,47 +2,41 @@ import numpy as np import pandas as pd import os -from astropy import ( - units as u, - constants as c - ) +from astropy import units as u, constants as c import astropy.tests.helper as test_helper from numpy.testing import assert_almost_equal from tardis.montecarlo.spectrum import ( - TARDISSpectrum, - ) + TARDISSpectrum, +) BIN = 5 TESTDATA = [ - { - 'nu': u.Quantity(np.linspace(1, 50, BIN + 1), 'Hz'), - 'lum': u.Quantity(np.linspace(1e27, 1e28, BIN), 'erg / s'), - 'distance': None - }, - { - 'nu': u.Quantity(np.linspace(1, 50, BIN + 1), 'Hz'), - 'lum': u.Quantity(np.linspace(1e27, 1e28, BIN), 'erg / s'), - 'distance': u.Quantity(1, 'Mpc') - }, - { - 'nu': u.Quantity([1, 2, 3, 4], 'Hz'), - 'lum': u.Quantity([1, 2, 3], 'erg / s') * np.pi, - 'distance': u.Quantity(.5, 'cm'), - } - ] - - -@pytest.fixture( - scope='module', - params=TESTDATA - ) + { + "nu": u.Quantity(np.linspace(1, 50, BIN + 1), "Hz"), + "lum": u.Quantity(np.linspace(1e27, 1e28, BIN), "erg / s"), + "distance": None, + }, + { + "nu": u.Quantity(np.linspace(1, 50, BIN + 1), "Hz"), + "lum": u.Quantity(np.linspace(1e27, 1e28, BIN), "erg / s"), + "distance": u.Quantity(1, "Mpc"), + }, + { + "nu": u.Quantity([1, 2, 3, 4], "Hz"), + "lum": u.Quantity([1, 2, 3], "erg / s") * np.pi, + "distance": u.Quantity(0.5, "cm"), + }, +] + + +@pytest.fixture(scope="module", params=TESTDATA) def spectrum(request): data = TARDISSpectrum( - request.param['nu'], - request.param['lum'], - ) - distance = request.param['distance'] + request.param["nu"], + request.param["lum"], + ) + distance = request.param["distance"] if distance: data.distance = distance return data @@ -52,64 +46,60 @@ def spectrum(request): # Testing properties ### def test_frequency(spectrum): - assert np.all( - spectrum.frequency == spectrum._frequency[:-1] - ) + assert np.all(spectrum.frequency == spectrum._frequency[:-1]) def test_luminosity_density_nu(spectrum): expected = spectrum.luminosity / np.diff(spectrum._frequency) test_helper.assert_quantity_allclose( - spectrum.luminosity_density_nu, - expected - ) + spectrum.luminosity_density_nu, expected + ) def test_luminosity_density_lambda(spectrum): expected = spectrum.f_nu_to_f_lambda( - spectrum.luminosity / np.diff(spectrum._frequency) - ) + spectrum.luminosity / np.diff(spectrum._frequency) + ) test_helper.assert_quantity_allclose( - spectrum.luminosity_density_lambda, - expected - ) + spectrum.luminosity_density_lambda, expected + ) def test_flux_nu(spectrum): - if getattr(spectrum, 'distance', None): + if getattr(spectrum, "distance", None): - if pytest.__version__ < '2.8': + if pytest.__version__ < "2.8": pytest.xfail( - reason='requires pytest.warns (introduced in pytest v2.8)', - ) + reason="requires pytest.warns (introduced in pytest v2.8)", + ) with pytest.warns(DeprecationWarning): test_helper.assert_quantity_allclose( - spectrum.flux_nu, - spectrum.luminosity_to_flux( - spectrum.luminosity_density_nu, - spectrum.distance) - ) + spectrum.flux_nu, + spectrum.luminosity_to_flux( + spectrum.luminosity_density_nu, spectrum.distance + ), + ) else: with pytest.raises(AttributeError): spectrum.flux_nu def test_flux_lambda(spectrum): - if getattr(spectrum, 'distance', None): + if getattr(spectrum, "distance", None): - if pytest.__version__ < '2.8': + if pytest.__version__ < "2.8": pytest.xfail( - reason='requires pytest.warns (introduced in pytest v2.8)', - ) + reason="requires pytest.warns (introduced in pytest v2.8)", + ) with pytest.warns(DeprecationWarning): test_helper.assert_quantity_allclose( - spectrum.flux_lambda, - spectrum.luminosity_to_flux( - spectrum.luminosity_density_lambda, - spectrum.distance) - ) + spectrum.flux_lambda, + spectrum.luminosity_to_flux( + spectrum.luminosity_density_lambda, spectrum.distance + ), + ) else: with pytest.raises(AttributeError): spectrum.flux_nu @@ -119,68 +109,63 @@ def test_flux_lambda(spectrum): # Testing methods ### + def test_luminosity_to_flux(): - lum = u.Quantity(np.arange(1, 4, 1) * np.pi, 'erg / s') - distance = u.Quantity(.5, 'cm') + lum = u.Quantity(np.arange(1, 4, 1) * np.pi, "erg / s") + distance = u.Quantity(0.5, "cm") flux = TARDISSpectrum.luminosity_to_flux(lum, distance) test_helper.assert_quantity_allclose( - flux, - u.Quantity(lum.value / np.pi, 'erg s^-1 cm^-2') - ) + flux, u.Quantity(lum.value / np.pi, "erg s^-1 cm^-2") + ) def test_f_nu_to_f_lambda(spectrum): expected = ( - spectrum.luminosity_density_nu * - spectrum.frequency**2 / c.c.to('angstrom/s') - ).to('erg / (s angstrom)') + spectrum.luminosity_density_nu + * spectrum.frequency ** 2 + / c.c.to("angstrom/s") + ).to("erg / (s angstrom)") print(expected) test_helper.assert_quantity_allclose( - spectrum.luminosity_density_lambda, - expected - ) + spectrum.luminosity_density_lambda, expected + ) expected = ( - spectrum.luminosity_density_nu * - spectrum.frequency**2 / c.c.to('angstrom/s') - ).to('erg / (s angstrom)') + spectrum.luminosity_density_nu + * spectrum.frequency ** 2 + / c.c.to("angstrom/s") + ).to("erg / (s angstrom)") np.testing.assert_allclose( - spectrum.luminosity_density_lambda.value, - expected.value - ) + spectrum.luminosity_density_lambda.value, expected.value + ) ### # Save and Load ### + def compare_spectra(actual, desired): - test_helper.assert_quantity_allclose( - actual.frequency, - desired.frequency - ) - test_helper.assert_quantity_allclose( - actual.luminosity, - desired.luminosity - ) - if getattr(actual, 'distance', None): - test_helper.assert_quantity_allclose( - actual.distance, - desired.distance - ) - + test_helper.assert_quantity_allclose(actual.frequency, desired.frequency) + test_helper.assert_quantity_allclose(actual.luminosity, desired.luminosity) + if getattr(actual, "distance", None): + test_helper.assert_quantity_allclose(actual.distance, desired.distance) + + @pytest.fixture(autouse=True) -def to_hdf_buffer(hdf_file_path,spectrum): - spectrum.to_hdf(hdf_file_path, name='spectrum') +def to_hdf_buffer(hdf_file_path, spectrum): + spectrum.to_hdf(hdf_file_path, name="spectrum") + @pytest.mark.parametrize("attr", TARDISSpectrum.hdf_properties) def test_hdf_spectrum(hdf_file_path, spectrum, attr): actual = getattr(spectrum, attr) - if hasattr(actual, 'cgs'): + if hasattr(actual, "cgs"): actual = actual.cgs.value - path = os.path.join('spectrum', attr) + path = os.path.join("spectrum", attr) expected = pd.read_hdf(hdf_file_path, path) assert_almost_equal(actual, expected.values) + ### # Test creation from nonstandard units ### @@ -188,17 +173,15 @@ def test_hdf_spectrum(hdf_file_path, spectrum, attr): def test_creat_from_wl(spectrum): actual = TARDISSpectrum( - spectrum._frequency.to('angstrom', u.spectral()), - spectrum.luminosity - ) + spectrum._frequency.to("angstrom", u.spectral()), spectrum.luminosity + ) compare_spectra(actual, spectrum) def test_creat_from_J(spectrum): actual = TARDISSpectrum( - spectrum._frequency, - spectrum.luminosity.to('J / s') - ) + spectrum._frequency, spectrum.luminosity.to("J / s") + ) compare_spectra(actual, spectrum) diff --git a/tardis/plasma/properties/atomic.py b/tardis/plasma/properties/atomic.py index b491518fccf..9e4f0e32293 100644 --- a/tardis/plasma/properties/atomic.py +++ b/tardis/plasma/properties/atomic.py @@ -279,9 +279,9 @@ def _set_index(self, zeta_data): class NLTEData(ProcessingPlasmaProperty): """ - Attributes: - nlte_data : -#Finish later (need atomic dataset with NLTE data). + Attributes: + nlte_data : + #Finish later (need atomic dataset with NLTE data). """ outputs = ("nlte_data",) diff --git a/tardis/plasma/properties/j_blues.py b/tardis/plasma/properties/j_blues.py index bb1d575a3f3..9ffbe3c1b5c 100644 --- a/tardis/plasma/properties/j_blues.py +++ b/tardis/plasma/properties/j_blues.py @@ -65,9 +65,10 @@ def calculate( for i in range(len(t_rad)): zero_j_blues = j_blues[i] == 0.0 - j_blues[i][zero_j_blues] = ( - self.w_epsilon - * intensity_black_body(nu[zero_j_blues].values, t_rad[i]) + j_blues[i][ + zero_j_blues + ] = self.w_epsilon * intensity_black_body( + nu[zero_j_blues].values, t_rad[i] ) return j_blues diff --git a/tardis/plasma/properties/nlte.py b/tardis/plasma/properties/nlte.py index ef7e58c3cc2..7dcddb53dd0 100644 --- a/tardis/plasma/properties/nlte.py +++ b/tardis/plasma/properties/nlte.py @@ -30,7 +30,10 @@ class PreviousElectronDensities(PreviousIterationProperty): outputs = ("previous_electron_densities",) def set_initial_value(self, kwargs): - initial_value = pd.Series(1000000.0, index=kwargs["abundance"].columns,) + initial_value = pd.Series( + 1000000.0, + index=kwargs["abundance"].columns, + ) self._set_initial_value(initial_value) diff --git a/tardis/plasma/properties/radiative_properties.py b/tardis/plasma/properties/radiative_properties.py index 43aaed04164..0deb7282358 100644 --- a/tardis/plasma/properties/radiative_properties.py +++ b/tardis/plasma/properties/radiative_properties.py @@ -249,8 +249,8 @@ def calculate( self.initialize_macro_atom_transition_type_filters( atomic_data, macro_atom_data ) - self.transition_probability_coef = self._get_transition_probability_coefs( - macro_atom_data + self.transition_probability_coef = ( + self._get_transition_probability_coefs(macro_atom_data) ) self.initialize = False transition_probabilities = self._calculate_transition_probability( diff --git a/tardis/plasma/properties/util/macro_atom.py b/tardis/plasma/properties/util/macro_atom.py new file mode 100644 index 00000000000..296c28e043a --- /dev/null +++ b/tardis/plasma/properties/util/macro_atom.py @@ -0,0 +1,59 @@ +from numba import njit +from tardis.montecarlo.montecarlo_numba import njit_dict +import numpy as np +from tardis import constants as const + +h_cgs = const.h.cgs.value +c = const.c.to("cm/s").value +kb = const.k_B.cgs.value +inv_c2 = 1 / (c ** 2) + + +@njit(**njit_dict) +def calculate_transition_probabilities( + transition_probability_coef, + beta_sobolev, + j_blues, + stimulated_emission_factor, + transition_type, + lines_idx, + block_references, + transition_probabilities, +): + """ + Calculates transition probabilities for macro_atom interactions + + transition_probability_coef must be a 1D array + transition_type, lines_idx, and block_references must be int-type arrays + beta_sobolev, j_blues,stimulated_emission_factor, and transition_probabilities must be 2D array + """ + + norm_factor = np.zeros(transition_probabilities.shape[1]) + + for i in range(transition_probabilities.shape[0]): + line_idx = lines_idx[i] + for j in range(transition_probabilities.shape[1]): + transition_probabilities[i, j] = ( + transition_probability_coef[i] * beta_sobolev[line_idx, j] + ) + if transition_type[i] == 1: + for j in range(transition_probabilities.shape[1]): + transition_probabilities[i, j] *= ( + stimulated_emission_factor[line_idx, j] + * j_blues[line_idx, j] + ) + + for i in range(block_references.shape[0] - 1): + for k in range(transition_probabilities.shape[1]): + norm_factor[k] = 0.0 + for j in range(block_references[i], block_references[i + 1]): + for k in range(transition_probabilities.shape[1]): + norm_factor[k] += transition_probabilities[j, k] + for k in range(transition_probabilities.shape[1]): + if norm_factor[k] != 0.0: + norm_factor[k] = 1 / norm_factor[k] + else: + norm_factor[k] = 1.0 + for j in range(block_references[i], block_references[i + 1]): + for k in range(0, transition_probabilities.shape[1]): + transition_probabilities[j, k] *= norm_factor[k] diff --git a/tardis/plasma/properties/util/macro_atom.pyx b/tardis/plasma/properties/util/macro_atom.pyx deleted file mode 100644 index e1c6fcceaf2..00000000000 --- a/tardis/plasma/properties/util/macro_atom.pyx +++ /dev/null @@ -1,61 +0,0 @@ -# module for fast macro_atom calculations - -# cython: profile=False -# cython: boundscheck=False -# cython: cdivision=True -# cython: wraparound=False -# cython: language_level=3 - -import numpy as np - -cimport numpy as np - -ctypedef np.int64_t int_type_t - -from astropy import constants - - -cdef extern from "math.h": - double exp(double) - - -cdef double h_cgs = constants.h.cgs.value -cdef double c = constants.c.cgs.value -cdef double kb = constants.k_B.cgs.value -cdef double inv_c2 = 1 / (c ** 2) - - -def calculate_transition_probabilities( - double [:] transition_probability_coef, - double [:, ::1] beta_sobolev, double [:, ::1] j_blues, - double [:, ::1] stimulated_emission_factor, - int_type_t [:] transition_type, - int_type_t [:] lines_idx, - int_type_t [:] block_references, - double [:, ::1] transition_probabilities): - - cdef int i, j, k, line_idx - cdef np.ndarray[double, ndim=1] norm_factor = np.zeros(transition_probabilities.shape[1]) - - for i in range(transition_probabilities.shape[0]): - line_idx = lines_idx[i] - for j in range(transition_probabilities.shape[1]): - transition_probabilities[i, j] = (transition_probability_coef[i] * beta_sobolev[line_idx, j]) - if transition_type[i] == 1: - for j in range(transition_probabilities.shape[1]): - transition_probabilities[i, j] *= stimulated_emission_factor[line_idx, j] * j_blues[line_idx, j] - - for i in range(block_references.shape[0] - 1): - for k in range(transition_probabilities.shape[1]): - norm_factor[k] = 0.0 - for j in range(block_references[i], block_references[i + 1]): - for k in range(transition_probabilities.shape[1]): - norm_factor[k] += transition_probabilities[j, k] - for k in range(transition_probabilities.shape[1]): - if norm_factor[k] != 0.0: - norm_factor[k] = 1 / norm_factor[k] - else: - norm_factor[k] = 1.0 - for j in range(block_references[i], block_references[i + 1]): - for k in range(0, transition_probabilities.shape[1]): - transition_probabilities[j, k] *= norm_factor[k] diff --git a/tardis/plasma/setup_package.py b/tardis/plasma/setup_package.py index 0a33ad7fbb9..7a68ae7bddf 100644 --- a/tardis/plasma/setup_package.py +++ b/tardis/plasma/setup_package.py @@ -37,13 +37,5 @@ def get_package_data(): def get_extensions(): - sources = ["tardis/plasma/properties/util/macro_atom.pyx"] - return [ - Extension( - "tardis.plasma.properties.util.macro_atom", - sources, - include_dirs=["numpy"], - extra_compile_args=compile_args, - extra_link_args=link_args, - ) - ] + sources = [] + return [] diff --git a/tardis/scripts/debug/run_numba_single.py b/tardis/scripts/debug/run_numba_single.py new file mode 100644 index 00000000000..1ab234b282c --- /dev/null +++ b/tardis/scripts/debug/run_numba_single.py @@ -0,0 +1,22 @@ +from tardis import run_tardis +import numpy as np +from tardis.montecarlo.montecarlo_numba.base import montecarlo_main_loop +import os +import numba +import sys +import yaml + + +SEED = eval(sys.argv[1].split("=")[1])[0] + +yaml_file, params = "tardis_example_single.yml", None + +with open(yaml_file) as f: + params = yaml.safe_load(f) + +params["montecarlo"]["single_packet_seed"] = SEED + +with open(yaml_file, "w") as f: + yaml.safe_dump(params, f) + +mdl = run_tardis(yaml_file) diff --git a/tardis/scripts/debug/run_numba_single.run.xml b/tardis/scripts/debug/run_numba_single.run.xml new file mode 100644 index 00000000000..bec9d6a5e38 --- /dev/null +++ b/tardis/scripts/debug/run_numba_single.run.xml @@ -0,0 +1,25 @@ + + + + + diff --git a/tardis/scripts/debug/tardis_example_single.yml b/tardis/scripts/debug/tardis_example_single.yml new file mode 100644 index 00000000000..35c17071b18 --- /dev/null +++ b/tardis/scripts/debug/tardis_example_single.yml @@ -0,0 +1,45 @@ +atom_data: kurucz_cd23_chianti_H_He.h5 +model: + abundances: + Si: 1.0 + type: uniform + structure: + density: + type: branch85_w7 + type: specific + velocity: + num: 20 + start: 1.1e4 km/s + stop: 20000 km/s +montecarlo: + convergence_strategy: + damping_constant: 1.0 + fraction: 0.8 + hold_iterations: 3 + t_inner: + damping_constant: 1.0 + threshold: 0.05 + type: damped + debug_packets: true + iterations: 2 + last_no_of_packets: 1000000.0 + logger_buffer: 1 + no_of_packets: 40000.0 + no_of_virtual_packets: 0 + nthreads: 6 + seed: 23111963 + single_packet_seed: 46 +plasma: + disable_electron_scattering: false + excitation: lte + ionization: lte + line_interaction_type: macroatom + radiative_rates_type: dilute-blackbody +spectrum: + num: 10000 + start: 500 angstrom + stop: 20000 angstrom +supernova: + luminosity_requested: 9.44 log_lsun + time_explosion: 10 day +tardis_config_version: v1.0 diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 623980acf2c..23ddfda50b2 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -10,6 +10,7 @@ from tardis.plasma.standard_plasmas import assemble_plasma from tardis.io.util import HDFWriterMixin from tardis.io.config_reader import ConfigurationError +from tardis.montecarlo import montecarlo_configuration as mc_config_module # Adding logging support logger = logging.getLogger(__name__) @@ -332,6 +333,7 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0, last_run=False): no_of_virtual_packets=no_of_virtual_packets, nthreads=self.nthreads, last_run=last_run, + iteration=self.iterations_executed, ) output_energy = self.runner.output_energy if np.sum(output_energy < 0) == len(output_energy): @@ -370,7 +372,9 @@ def run(self): self.plasma.electron_densities, self.model.t_inner, ) - self.iterate(self.last_no_of_packets, self.no_of_virtual_packets, True) + self.iterate( + self.last_no_of_packets, self.no_of_virtual_packets, last_run=True + ) self.reshape_plasma_state_store(self.iterations_executed) diff --git a/tardis/stats/base.py b/tardis/stats/base.py index e0ea57f684e..16519285e7a 100644 --- a/tardis/stats/base.py +++ b/tardis/stats/base.py @@ -2,9 +2,7 @@ def get_trivial_poisson_uncertainty(model): - """ - - """ + """""" emitted_nu = model.montecarlo_nu[model.montecarlo_luminosity >= 0] emitted_luminosity = model.montecarlo_luminosity[ model.montecarlo_luminosity >= 0 diff --git a/tardis/tests/integration_tests/test_integration.py b/tardis/tests/integration_tests/test_integration.py index c2719558416..d3bf223901d 100644 --- a/tardis/tests/integration_tests/test_integration.py +++ b/tardis/tests/integration_tests/test_integration.py @@ -12,6 +12,8 @@ from tardis.simulation import Simulation from tardis.io.config_reader import Configuration +pytestmarker = [pytest.mark.no_cover, pytest.mark.integration] + quantity_comparison = [ ( "/simulation/runner/last_line_interaction_in_id", @@ -71,7 +73,7 @@ def model_quantities(request): @pytest.mark.integration class TestIntegration(object): """Slow integration test for various setups present in subdirectories of - ``tardis/tests/integration_tests``. + ``tardis/tests/integration_tests``. """ @classmethod diff --git a/tardis/tests/test_tardis_full.py b/tardis/tests/test_tardis_full.py index 825909e165c..066b1a6bc00 100644 --- a/tardis/tests/test_tardis_full.py +++ b/tardis/tests/test_tardis_full.py @@ -62,12 +62,12 @@ def test_virtual_spectrum(self, runner, refdata): assert_quantity_allclose(runner.spectrum_virtual.luminosity, luminosity) def test_runner_properties(self, runner): - """Tests whether a number of runner attributes exist and also verifies + """ + Tests whether a number of runner attributes exist and also verifies their types Currently, runner attributes needed to call the model routine to_hdf5 are checked. - """ virt_type = np.ndarray @@ -87,7 +87,8 @@ def test_runner_properties(self, runner): for prop, prop_type in required_props.items(): actual = getattr(runner, prop) - assert type(actual) == prop_type, ( - "wrong type of attribute '{}':" - "expected {}, found {}".format(prop, prop_type, type(actual)) + assert ( + type(actual) == prop_type + ), "wrong type of attribute '{}':" "expected {}, found {}".format( + prop, prop_type, type(actual) ) diff --git a/tardis/tests/test_tardis_full_formal_integral.py b/tardis/tests/test_tardis_full_formal_integral.py index 53a7343e7ea..dcf40a7b898 100644 --- a/tardis/tests/test_tardis_full_formal_integral.py +++ b/tardis/tests/test_tardis_full_formal_integral.py @@ -1,13 +1,13 @@ import os import pytest -import numpy as np import numpy.testing as npt from astropy import units as u from astropy.tests.helper import assert_quantity_allclose from tardis.simulation.base import Simulation from tardis.io.config_reader import Configuration -import astropy + +pytestmark = pytest.mark.skip(reason="memory problem") config_line_modes = ["downbranch", "macroatom"] interpolate_shells = [-1, 30] diff --git a/tardis/tests/test_util.py b/tardis/tests/test_util.py index 820bd6ef5dd..9fa47d19c40 100644 --- a/tardis/tests/test_util.py +++ b/tardis/tests/test_util.py @@ -149,7 +149,11 @@ def test_species_tuple_to_string(species_tuple, roman_numerals, species_string): @pytest.mark.parametrize( ["species_string", "species_tuple"], - [("si ii", (14, 1)), ("si 2", (14, 1)), ("si ix", (14, 8)),], + [ + ("si ii", (14, 1)), + ("si 2", (14, 1)), + ("si ix", (14, 8)), + ], ) def test_species_string_to_tuple(species_string, species_tuple): assert species_string_to_tuple(species_string) == species_tuple @@ -198,7 +202,13 @@ def test_atomic_number2element_symbol(): @pytest.mark.parametrize( ["unformatted_element_string", "formatted_element_string"], - [("si", "Si"), ("sI", "Si"), ("Si", "Si"), ("c", "C"), ("C", "C"),], + [ + ("si", "Si"), + ("sI", "Si"), + ("Si", "Si"), + ("c", "C"), + ("C", "C"), + ], ) def test_reformat_element_symbol( unformatted_element_string, formatted_element_string diff --git a/tardis/widgets/shell_info.py b/tardis/widgets/shell_info.py index 94a11ad5539..c006786dfdf 100644 --- a/tardis/widgets/shell_info.py +++ b/tardis/widgets/shell_info.py @@ -13,8 +13,7 @@ class BaseShellInfo: - """The simulation information that is used by shell info widget - """ + """The simulation information that is used by shell info widget""" def __init__( self, diff --git a/tardis/widgets/tests/test_line_info.py b/tardis/widgets/tests/test_line_info.py index 5d9d055b269..a3be0208408 100644 --- a/tardis/widgets/tests/test_line_info.py +++ b/tardis/widgets/tests/test_line_info.py @@ -24,6 +24,7 @@ def line_info_widget(simulation_verysimple): class TestLineInfoWidgetData: """Tests for methods that handles data in LineInfoWidget.""" + @pytest.mark.xfail(reason="Add real packet information to runner") def test_get_species_interactions( self, line_info_widget, wavelength_range, filter_mode ): @@ -170,6 +171,8 @@ def liw_with_selection(self, simulation_verysimple, request): return liw, selection_range + + @pytest.mark.xfail(reason="Add real packet information to runner") def test_selection_on_plot(self, liw_with_selection): """ Test if selection on spectrum plot, updates correct data in both @@ -217,6 +220,7 @@ def test_selection_on_plot(self, liw_with_selection): line_info_widget.total_packets_label.widget.children[1].value ) + @pytest.mark.xfail(reason="Add real packet information to runner") @pytest.mark.parametrize("selected_filter_mode_idx", [0, 1]) def test_filter_mode_toggle( self, diff --git a/tardis/widgets/tests/test_shell_info.py b/tardis/widgets/tests/test_shell_info.py index 0daac6969ed..e739d693c97 100644 --- a/tardis/widgets/tests/test_shell_info.py +++ b/tardis/widgets/tests/test_shell_info.py @@ -67,14 +67,16 @@ def test_ion_count_data( self, base_shell_info, simulation_verysimple, atomic_num, shell_num ): ion_count_data = base_shell_info.ion_count(atomic_num, shell_num) - sim_ion_number_density = simulation_verysimple.plasma.ion_number_density[ - shell_num - 1 - ].loc[ - atomic_num - ] - sim_element_number_density = simulation_verysimple.plasma.number_density.loc[ - atomic_num, shell_num - 1 - ] + sim_ion_number_density = ( + simulation_verysimple.plasma.ion_number_density[shell_num - 1].loc[ + atomic_num + ] + ) + sim_element_number_density = ( + simulation_verysimple.plasma.number_density.loc[ + atomic_num, shell_num - 1 + ] + ) assert ion_count_data.shape == (len(sim_ion_number_density), 2) assert np.allclose( ion_count_data.iloc[:, -1].map(np.float), @@ -95,16 +97,16 @@ def test_level_count_data( level_count_data = base_shell_info.level_count( ion_num, atomic_num, shell_num ) - sim_level_number_density = simulation_verysimple.plasma.level_number_density[ - shell_num - 1 - ].loc[ - atomic_num, ion_num - ] - sim_ion_number_density = simulation_verysimple.plasma.ion_number_density[ - shell_num - 1 - ].loc[ - atomic_num, ion_num - ] + sim_level_number_density = ( + simulation_verysimple.plasma.level_number_density[ + shell_num - 1 + ].loc[atomic_num, ion_num] + ) + sim_ion_number_density = ( + simulation_verysimple.plasma.ion_number_density[shell_num - 1].loc[ + atomic_num, ion_num + ] + ) assert level_count_data.shape == (len(sim_level_number_density), 1) assert np.allclose( level_count_data.iloc[:, 0].map(np.float), diff --git a/tardis/widgets/util.py b/tardis/widgets/util.py index cef0dedbc0f..398bafabfa2 100644 --- a/tardis/widgets/util.py +++ b/tardis/widgets/util.py @@ -24,7 +24,7 @@ def create_table_widget( table_options : dict, optional A dictionary to specify options to use when creating interactive table widget (same as :grid_options: of qgrid, specified in Notes section of - their `API documentation _`). + their `API documentation _`). Invalid keys will have no effect and valid keys will override or add to the options used by this function. changeable_col : dict, optional @@ -105,8 +105,8 @@ def create_table_widget( class TableSummaryLabel: """ Label like widget to show summary of a qgrid table widget. - - Also handles aligning the label with the table columns exactly like a + + Also handles aligning the label with the table columns exactly like a summary row. """ @@ -145,10 +145,10 @@ def __init__(self, target_table, table_col_widths, label_key, label_value): def update_and_resize(self, value): """ Update the label value and resize it as per the size of target table. - + Resizing is done in such a way so as to match the width of components - of label with columns of target table, making it look like another row. - This method should be called whenever there is any update in data or + of label with columns of target table, making it look like another row. + This method should be called whenever there is any update in data or layout of target table. Parameters @@ -158,8 +158,8 @@ def update_and_resize(self, value): Note ---- - The width resizing operation is highly dependent on qgrid tables' - layout. So it may not remain precise if there happens any CSS change + The width resizing operation is highly dependent on qgrid tables' + layout. So it may not remain precise if there happens any CSS change in upcoming versions of qgrid. """ self.widget.children[1].value = str(value) @@ -216,7 +216,13 @@ def _create(self, key, value): f"
{key}:
", layout=component_layout_options, ), - ipw.HTML(str(value), layout=component_layout_options,), + ipw.HTML( + str(value), + layout=component_layout_options, + ), ], - layout=dict(display="flex", justify_content="flex-start",), + layout=dict( + display="flex", + justify_content="flex-start", + ), )