<commit_before>import os
import subprocess
import warnings

import h5py
import numpy as np
import matplotlib.pyplot as plt

import hyperion
from hyperion.util.functions import delete_file, random_id
from hyperion.grid import CartesianGrid, SphericalPolarGrid, CylindricalPolarGrid, OcTreeGrid, AMRGrid
from hyperion.sources import PointSource, SphericalSource, ExternalSphericalSource, ExternalBoxSource, MapSource, PlaneParallelSource
from hyperion.conf import RunConf, PeeledImageConf, BinnedImageConf, OutputConf
from hyperion.util.constants import c, pi
from hyperion.util.functions import FreezableClass
from hyperion.dust import SphericalDust

STOKESD = {}
STOKESD['I'] = 0
STOKESD['Q'] = 1
STOKESD['U'] = 2
STOKESD['V'] = 3

LABEL = {}
LABEL['I'] = '$\lambda\, F_\lambda$ (ergs/cm$^2$/s)'
LABEL['Q'] = '$\lambda\, F_\lambda$ [stokes=Q] (ergs/cm$^2$/s)'
LABEL['U'] = '$\lambda\, F_\lambda$ [stokes=U] (ergs/cm$^2$/s)'
LABEL['V'] = '$\lambda\, F_\lambda$ [stokes=V] (ergs/cm$^2$/s)'
LABEL['linpol'] = "Total linear polarization fraction"
LABEL['circpol'] = "Total circular polarization fraction"


def mc_linear_polarization(I, sigma_I, Q, sigma_Q, U, sigma_U):
    n = 1000
    if I > 0.:
        if sigma_I == 0.:
            Is = np.repeat(I, n)
        else:
            Is = np.random.normal(loc=I, scale=sigma_I, size=n)
        if sigma_Q == 0.:
            Qs = np.repeat(Q, n)
        else:
            Qs = np.random.normal(loc=Q, scale=sigma_Q, size=n)
        if sigma_U == 0.:
            Us = np.repeat(U, n)
        else:
            Us = np.random.normal(loc=U, scale=sigma_U, size=n)
        Ps = np.sqrt((Qs ** 2 + Us ** 2) / Is ** 2)
        return np.mean(Ps), np.std(Ps)
    else:
        return 0., 0.


def mc_circular_polarization(I, sigma_I, V, sigma_V):
    n = 1000
    if I > 0.:
        if sigma_I == 0.:
            Is = np.repeat(I, n)
        else:
            Is = np.random.normal(loc=I, scale=sigma_I, size=n)
        if sigma_V == 0.:
            Vs = np.repeat(V, n)
        else:
            Vs = np.random.normal(loc=V, scale=sigma_V, size=n)
        Ps = np.abs(Vs) / Is
        return np.mean(Ps), np.std(Ps)
    else:
        return 0., 0.


def bool2str(value):
    if value:
        return "yes"
    else:
        return "no"


class File(file):

    def __init__(self, filename, *args, **kwargs):
        self.abspath = os.path.abspath(filename)
        file.__init__(self, filename, *args, **kwargs)


class Configuration(FreezableClass):

    def __init__(self):
        self.output = OutputConf()
        self.run = RunConf()
        self._freeze()


class Model(FreezableClass):

    def __init__(self, name):
        '''
        Create a Model instance

        Parameters
        ----------
        name : str
            The name of the model. This is used to name the .rtin file
            output with the write() method.
        '''

        if name[-6:] == '.rtout':
            name = name[:-6]

        self.name = name
        self.reset_density()
        self.reset_sources()
        self.reset_images()
        self.set_monochromatic(False)

        self.conf = Configuration()

        self.grid = None

        # Import methods for convenience
        self.set_n_temperature_iterations = self.conf.run.set_n_temperature_iterations
        self.set_n_photons = self.conf.run.set_n_photons
        self.set_raytracing = self.conf.run.set_raytracing
        self.set_max_interactions = self.conf.run.set_max_interactions
        self.set_pda = self.conf.run.set_pda
        self.set_mrw = self.conf.run.set_mrw
        self.set_convergence = self.conf.run.set_convergence
        self.set_kill_on_absorb = self.conf.run.set_kill_on_absorb
        self.set_forced_first_scattering = self.conf.run.set_forced_first_scattering
        self.set_output_bytes = self.conf.run.set_output_bytes
        self.set_sample_sources_evenly = self.conf.run.set_sample_sources_evenly

        self._freeze()

    def reset_density(self):
        self.density = []
        self.temperature = []
        self.dust = []

    def reset_sources(self):
        self.sources = []

    def reset_images(self):
        self.binned_output = None
        self.peeled_output = []

    def set_monochromatic(self, monochromatic, frequencies=None, wavelengths=None):
        '''
        Set whether to do the radiation transfer at specific
        frequencies/wavelengths.

        Parameters
        ----------
        monochromatic : bool
            Whether to carry out radiation transfer at specific frequencies
            or wavelengths
        frequencies : iterable of floats, optional
            The frequencies to compute the radiation transfer for, in Hz
        wavelengths : iterable of floats, optional
            The wavelengths to compute the radiation transfer for, in microns

        If `monochromatic` is True, then one of `frequencies` or
        `wavelengths` is required
        '''

        self._monochromatic = monochromatic

        if self._monochromatic:
            if wavelengths is not None and frequencies is not None:
                raise Exception("Cannot specify both frequencies and wavelengths")
            elif wavelengths is not None:
                frequencies = c / (np.array(wavelengths) * 1.e-4)
            elif frequencies is not None:
                frequencies = np.array(frequencies)
            else:
                raise Exception("Need to specify frequencies or wavelengths")
        else:
            if wavelengths is not None or frequencies is not None:
                raise Exception("Cannot specify monochromatic frequencies or wavelengths if monochromatic=False")

        self._frequencies = frequencies

        for images in self.peeled_output:
            images._monochromatic = True
            if type(images.wav_min) != int or type(images.wav_max) != int:
                images.set_wavelength_range(len(frequencies), 1, len(frequencies))
        if self.binned_output is not None:
            raise Exception("Binned images cannot be computed in monochromatic mode")

    def _write_monochromatic(self, group, compression=True, dtype=np.float64):
        group.attrs['monochromatic'] = bool2str(self._monochromatic)
        if self._monochromatic:
            group.create_dataset('Frequencies', data=np.array(zip(self._frequencies), dtype=[('nu', dtype)]), compression=compression)

    def write(self, compression=True, copy_dust=False, absolute_paths=False, geo_dtype=float, wall_dtype=float, physics_dtype=float):
        '''
        Write the model to an HDF5 file

        Optional Keyword Arguments:

            *compression*: [ True | False ]
                Whether to compress the datasets inside the HDF5 file

            *copy_dust*: [ True | False ]
                Whether to copy the dust properties when a filename is specified (True), or to just link to the file (False)

            *paths*: [ True | False ]
                Whether to use absolute paths (True) or relative paths (False) when linking to files
        '''

        if not self.name:
            raise Exception("Model name is not defined")

        if not os.path.dirname(self.name) == "" and not os.path.exists(os.path.dirname(self.name)):
            raise Exception("Directory %s does not exist" % os.path.dirname(self.name))

        # Create output file
        delete_file('%s.rtin' % self.name)
        root = h5py.File('%s.rtin' % self.name, 'w')
        root.attrs['python_version'] = hyperion.__version__

        g_dust = root.create_group('Dust')
        g_grid = root.create_group('Grid')
        g_geometry = g_grid.create_group('Geometry')
        g_physics = g_grid.create_group('Physics')
        g_sources = root.create_group('Sources')
        g_output = root.create_group('Output')
        g_peeled = g_output.create_group('Peeled')
        g_binned = g_output.create_group('Binned')

        if self.grid is None:
            raise Exception("No grid has been set up")

        # Generate random geometry ID
        self.grid.geometry_id = random_id()

        # Output sources
        for i, source in enumerate(self.sources):
            if isinstance(source, MapSource):
                source.write(g_sources, 'Source %i' % i, self.grid, compression=compression, map_dtype=physics_dtype)
            else:
                source.write(g_sources, 'Source %i' % i)

        # Output configuration files
        for i, peel in enumerate(self.peeled_output):
            if not self._frequencies is None:
                if not peel._monochromatic:
                    raise Exception("Peeled images need to be set to monochromatic mode")
            peel.write(g_peeled.create_group('Group %i' % (i + 1)))
        if self.binned_output is not None:
            self.binned_output.write(g_binned.create_group('Group 1'))

        self._write_monochromatic(root, compression=compression)
        self.conf.run.write(root)
        self.conf.output.write(g_output)

        if len(self.density) > 0:

            self.grid.write_physical_array(g_physics, self.density, "Density", dust=True, compression=compression, physics_dtype=physics_dtype)

            if self.temperature:
                self.grid.write_physical_array(g_physics, self.temperature, "Temperature", dust=True, compression=compression, physics_dtype=physics_dtype)

            # Output dust file
            present = {}
            for i, dust in enumerate(self.dust):

                short_name = 'dust_%03i' % (i + 1)

                if copy_dust:

                    if type(dust) == str:
                        dust = SphericalDust(dust)

                    if id(dust) in present:
                        g_dust[short_name] = h5py.SoftLink(present[id(dust)])
                    else:
                        dust.write(g_dust.create_group(short_name))
                        present[id(dust)] = short_name

                else:

                    if type(dust) != str:
                        dust = dust.filename

                    if absolute_paths:
                        g_dust[short_name] = h5py.ExternalLink(os.path.abspath(dust), '/')
                    else:
                        g_dust[short_name] = h5py.ExternalLink(os.path.relpath(dust), '/')

        # Output geometry
        self.grid.write_geometry(g_geometry, \
                                 compression=compression, geo_dtype=geo_dtype, wall_dtype=wall_dtype)

        root.close()

    def add_point_source(self, *args, **kwargs):
        source = PointSource(*args, **kwargs)
        self.add_source(source)
        return source

    def add_spherical_source(self, *args, **kwargs):
        source = SphericalSource(*args, **kwargs)
        self.add_source(source)
        return source

    def add_external_spherical_source(self, *args, **kwargs):
        source = ExternalSphericalSource(*args, **kwargs)
        self.add_source(source)
        return source

    def add_external_box_source(self, *args, **kwargs):
        source = ExternalBoxSource(*args, **kwargs)
        self.add_source(source)
        return source

    def add_map_source(self, *args, **kwargs):
        source = MapSource(*args, **kwargs)
        self.add_source(source)
        return source

    def add_plane_parallel_source(self, *args, **kwargs):
        source = PlaneParallelSource(*args, **kwargs)
        self.add_source(source)
        return source

    def add_source(self, source):
        self.sources.append(source)

    def add_density_grid(self, density, dust, temperature=None):
        if not self.grid:
            raise Exception("Grid not defined")
        if not isinstance(self.grid, AMRGrid):
            if not density.shape == self.grid.shape:
                raise Exception("Density shape does not match that of grid")
            if temperature is not None:
                if not temperature.shape == self.grid.shape:
                    raise Exception("Temperature shape does not match that of grid")
        if np.all(density == 0.):
            warnings.warn("All density values are zero - ignoring density grid")
            return
        self.density.append(density)
        self.dust.append(dust)
        if temperature is not None:
            self.temperature.append(temperature)<commit_msg>Can now import temperature from rtout file to a new model by filename reference<commit_after>import os
import subprocess
import warnings

import h5py
import numpy as np
import matplotlib.pyplot as plt

import hyperion
from hyperion.util.functions import delete_file, random_id
from hyperion.grid import CartesianGrid, SphericalPolarGrid, CylindricalPolarGrid, OcTreeGrid, AMRGrid
from hyperion.sources import PointSource, SphericalSource, ExternalSphericalSource, ExternalBoxSource, MapSource, PlaneParallelSource
from hyperion.conf import RunConf, PeeledImageConf, BinnedImageConf, OutputConf
from hyperion.util.constants import c, pi
from hyperion.util.functions import FreezableClass
from hyperion.dust import SphericalDust

STOKESD = {}
STOKESD['I'] = 0
STOKESD['Q'] = 1
STOKESD['U'] = 2
STOKESD['V'] = 3

LABEL = {}
LABEL['I'] = '$\lambda\, F_\lambda$ (ergs/cm$^2$/s)'
LABEL['Q'] = '$\lambda\, F_\lambda$ [stokes=Q] (ergs/cm$^2$/s)'
LABEL['U'] = '$\lambda\, F_\lambda$ [stokes=U] (ergs/cm$^2$/s)'
LABEL['V'] = '$\lambda\, F_\lambda$ [stokes=V] (ergs/cm$^2$/s)'
LABEL['linpol'] = "Total linear polarization fraction"
LABEL['circpol'] = "Total circular polarization fraction"


def mc_linear_polarization(I, sigma_I, Q, sigma_Q, U, sigma_U):
    n = 1000
    if I > 0.:
        if sigma_I == 0.:
            Is = np.repeat(I, n)
        else:
            Is = np.random.normal(loc=I, scale=sigma_I, size=n)
        if sigma_Q == 0.:
            Qs = np.repeat(Q, n)
        else:
            Qs = np.random.normal(loc=Q, scale=sigma_Q, size=n)
        if sigma_U == 0.:
            Us = np.repeat(U, n)
        else:
            Us = np.random.normal(loc=U, scale=sigma_U, size=n)
        Ps = np.sqrt((Qs ** 2 + Us ** 2) / Is ** 2)
        return np.mean(Ps), np.std(Ps)
    else:
        return 0., 0.


def mc_circular_polarization(I, sigma_I, V, sigma_V):
    n = 1000
    if I > 0.:
        if sigma_I == 0.:
            Is = np.repeat(I, n)
        else:
            Is = np.random.normal(loc=I, scale=sigma_I, size=n)
        if sigma_V == 0.:
            Vs = np.repeat(V, n)
        else:
            Vs = np.random.normal(loc=V, scale=sigma_V, size=n)
        Ps = np.abs(Vs) / Is
        return np.mean(Ps), np.std(Ps)
    else:
        return 0., 0.


def bool2str(value):
    if value:
        return "yes"
    else:
        return "no"


class File(file):

    def __init__(self, filename, *args, **kwargs):
        self.abspath = os.path.abspath(filename)
        file.__init__(self, filename, *args, **kwargs)


class Configuration(FreezableClass):

    def __init__(self):
        self.output = OutputConf()
        self.run = RunConf()
        self._freeze()


class Model(FreezableClass):

    def __init__(self, name):
        '''
        Create a Model instance

        Parameters
        ----------
        name : str
            The name of the model. This is used to name the .rtin file
            output with the write() method.
        '''

        if name[-6:] == '.rtout':
            name = name[:-6]

        self.name = name
        self.reset_density()
        self.reset_sources()
        self.reset_images()
        self.temperature = None
        self.set_monochromatic(False)

        self.conf = Configuration()

        self.grid = None

        # Import methods for convenience
        self.set_n_temperature_iterations = self.conf.run.set_n_temperature_iterations
        self.set_n_photons = self.conf.run.set_n_photons
        self.set_raytracing = self.conf.run.set_raytracing
        self.set_max_interactions = self.conf.run.set_max_interactions
        self.set_pda = self.conf.run.set_pda
        self.set_mrw = self.conf.run.set_mrw
        self.set_convergence = self.conf.run.set_convergence
        self.set_kill_on_absorb = self.conf.run.set_kill_on_absorb
        self.set_forced_first_scattering = self.conf.run.set_forced_first_scattering
        self.set_output_bytes = self.conf.run.set_output_bytes
        self.set_sample_sources_evenly = self.conf.run.set_sample_sources_evenly

        self._freeze()

    def reset_density(self):
        self.density = []
        self.dust = []

    def reset_sources(self):
        self.sources = []

    def reset_images(self):
        self.binned_output = None
        self.peeled_output = []

    def set_monochromatic(self, monochromatic, frequencies=None, wavelengths=None):
        '''
        Set whether to do the radiation transfer at specific
        frequencies/wavelengths.

        Parameters
        ----------
        monochromatic : bool
            Whether to carry out radiation transfer at specific frequencies
            or wavelengths
        frequencies : iterable of floats, optional
            The frequencies to compute the radiation transfer for, in Hz
        wavelengths : iterable of floats, optional
            The wavelengths to compute the radiation transfer for, in microns

        If `monochromatic` is True, then one of `frequencies` or
        `wavelengths` is required
        '''

        self._monochromatic = monochromatic

        if self._monochromatic:
            if wavelengths is not None and frequencies is not None:
                raise Exception("Cannot specify both frequencies and wavelengths")
            elif wavelengths is not None:
                frequencies = c / (np.array(wavelengths) * 1.e-4)
            elif frequencies is not None:
                frequencies = np.array(frequencies)
            else:
                raise Exception("Need to specify frequencies or wavelengths")
        else:
            if wavelengths is not None or frequencies is not None:
                raise Exception("Cannot specify monochromatic frequencies or wavelengths if monochromatic=False")

        self._frequencies = frequencies

        for images in self.peeled_output:
            images._monochromatic = True
            if type(images.wav_min) != int or type(images.wav_max) != int:
                images.set_wavelength_range(len(frequencies), 1, len(frequencies))
        if self.binned_output is not None:
            raise Exception("Binned images cannot be computed in monochromatic mode")

    def _write_monochromatic(self, group, compression=True, dtype=np.float64):
        group.attrs['monochromatic'] = bool2str(self._monochromatic)
        if self._monochromatic:
            group.create_dataset('Frequencies', data=np.array(zip(self._frequencies), dtype=[('nu', dtype)]), compression=compression)

    def write(self, compression=True, copy_dust=False, absolute_paths=False, geo_dtype=float, wall_dtype=float, physics_dtype=float):
        '''
        Write the model to an HDF5 file

        Optional Keyword Arguments:

            *compression*: [ True | False ]
                Whether to compress the datasets inside the HDF5 file

            *copy_dust*: [ True | False ]
                Whether to copy the dust properties when a filename is specified (True), or to just link to the file (False)

            *paths*: [ True | False ]
                Whether to use absolute paths (True) or relative paths (False) when linking to files
        '''

        if not self.name:
            raise Exception("Model name is not defined")

        if not os.path.dirname(self.name) == "" and not os.path.exists(os.path.dirname(self.name)):
            raise Exception("Directory %s does not exist" % os.path.dirname(self.name))

        # Create output file
        delete_file('%s.rtin' % self.name)
        root = h5py.File('%s.rtin' % self.name, 'w')
        root.attrs['python_version'] = hyperion.__version__

        g_dust = root.create_group('Dust')
        g_grid = root.create_group('Grid')
        g_geometry = g_grid.create_group('Geometry')
        g_physics = g_grid.create_group('Physics')
        g_sources = root.create_group('Sources')
        g_output = root.create_group('Output')
        g_peeled = g_output.create_group('Peeled')
        g_binned = g_output.create_group('Binned')

        if self.grid is None:
            raise Exception("No grid has been set up")

        # Generate random geometry ID
        self.grid.geometry_id = random_id()

        # Output sources
        for i, source in enumerate(self.sources):
            if isinstance(source, MapSource):
                source.write(g_sources, 'Source %i' % i, self.grid, compression=compression, map_dtype=physics_dtype)
            else:
                source.write(g_sources, 'Source %i' % i)

        # Output configuration files
        for i, peel in enumerate(self.peeled_output):
            if not self._frequencies is None:
                if not peel._monochromatic:
                    raise Exception("Peeled images need to be set to monochromatic mode")
            peel.write(g_peeled.create_group('Group %i' % (i + 1)))
        if self.binned_output is not None:
            self.binned_output.write(g_binned.create_group('Group 1'))

        self._write_monochromatic(root, compression=compression)
        self.conf.run.write(root)
        self.conf.output.write(g_output)

        if len(self.density) > 0:

            self.grid.write_physical_array(g_physics, self.density, "Density", dust=True, compression=compression, physics_dtype=physics_dtype)

            if self.temperature is not None:
                if type(self.temperature) is list:
                    self.grid.write_physical_array(g_physics, self.temperature, "Temperature", dust=True, compression=compression, physics_dtype=physics_dtype)
                elif type(self.temperature) is str:
                    f = h5py.File(self.temperature, 'r')
                    max_iteration = 0
                    for group_name in f:
                        if "Iteration" in group_name:
                            iteration = int(group_name.split()[1])
                            max_iteration = max(iteration, max_iteration)
                    print "Retrieving temperature from iteration %i of %s" % (max_iteration, self.temperature)
                    g_physics.create_dataset('Temperature', data=f['Iteration %05i' % max_iteration]['temperature'])
                else:
                    raise Exception("Unknown type %s for Model.temperature" % str(type(self.temperature)))

            # Output dust file
            present = {}
            for i, dust in enumerate(self.dust):

                short_name = 'dust_%03i' % (i + 1)

                if copy_dust:

                    if type(dust) == str:
                        dust = SphericalDust(dust)

                    if id(dust) in present:
                        g_dust[short_name] = h5py.SoftLink(present[id(dust)])
                    else:
                        dust.write(g_dust.create_group(short_name))
                        present[id(dust)] = short_name

                else:

                    if type(dust) != str:
                        dust = dust.filename

                    if absolute_paths:
                        g_dust[short_name] = h5py.ExternalLink(os.path.abspath(dust), '/')
                    else:
                        g_dust[short_name] = h5py.ExternalLink(os.path.relpath(dust), '/')

        # Output geometry
        self.grid.write_geometry(g_geometry, \
                                 compression=compression, geo_dtype=geo_dtype, wall_dtype=wall_dtype)

        root.close()

    def add_point_source(self, *args, **kwargs):
        source = PointSource(*args, **kwargs)
        self.add_source(source)
        return source

    def add_spherical_source(self, *args, **kwargs):
        source = SphericalSource(*args, **kwargs)
        self.add_source(source)
        return source

    def add_external_spherical_source(self, *args, **kwargs):
        source = ExternalSphericalSource(*args, **kwargs)
        self.add_source(source)
        return source

    def add_external_box_source(self, *args, **kwargs):
        source = ExternalBoxSource(*args, **kwargs)
        self.add_source(source)
        return source

    def add_map_source(self, *args, **kwargs):
        source = MapSource(*args, **kwargs)
        self.add_source(source)
        return source

    def add_plane_parallel_source(self, *args, **kwargs):
        source = PlaneParallelSource(*args, **kwargs)
        self.add_source(source)
        return source

    def add_source(self, source):
        self.sources.append(source)

    def add_density_grid(self, density, dust, temperature=None):
        if not self.grid:
            raise Exception("Grid not defined")
        if not isinstance(self.grid, AMRGrid):
            if not density.shape == self.grid.shape:
                raise Exception("Density shape does not match that of grid")
            if temperature is not None:
                if not temperature.shape == self.grid.shape:
                    raise Exception("Temperature shape does not match that of grid")
        if np.all(density == 0.):
            warnings.warn("All density values are zero - ignoring density grid")
            return
        self.density.append(density)
        self.dust.append(dust)
        if temperature is not None:
            if self.temperature is None:
                self.temperature = []
            self.temperature.append(temperature)
{'fraction_non_alphanumeric': 0.10180109631949882, 'fraction_numerical': 0.004346123727486296, 'mean_word_length': 2.2127044025157234, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6851572', 'n_tokens_mistral': 7743, 'n_tokens_neox': 7280, 'n_words': 1859}