setup for fitting artis models

artis_fitness_function

from dalek import assemble_tardis_model
from dalek.base.model import SimpleTARDISUncertaintyModel, NormRedSpectrum, LogLikelihood
from dalek.base.base import BaseLikelihoodModel
from dalek.base.transform import uniform_transform
from astropy import units as u, constants as const
import os

base_dir = '/home/hpc/pr94se/di73qom/projects/tardis/dalek/test_artis_run/artis_fit'

artis_filename = os.path.join(base_dir, 'aaflux_057.dat')
log_fname = os.path.join(base_dir, 'dalek_log.txt')
spec_dir = os.path.join(base_dir, 'spectral_store')

class ArtisDELogLikelihood(BaseLikelihoodModel):
    def transform(self, si, s, ca, fe, co, ni, mg, ti, cr, c, lum, v_inner):
        v_inner = uniform_transform(v_inner, 7000, 15000)
        lum = uniform_transform(lum, 4e42, 3e43)

        s = uniform_transform(s, 0, 0.5)
        fe = uniform_transform(fe, 0, 0.2)
        co = uniform_transform(co, 0, 0.2)
        ni = uniform_transform(ni, 0, 0.2)

        c = uniform_transform(c, 0, 0.1)
        mg = uniform_transform(mg, 0, 0.1)
        ca = uniform_transform(ca, 0, 0.1)

        ti = uniform_transform(ti, 0, 0.02)
        cr = uniform_transform(cr, 0, 0.02)

        o = 1 - np.sum([si, s, ca, fe, co, ni, mg, ti, cr, c])

        return o, si, s, ca, fe, co, ni, mg, ti, cr, c, lum, v_inner

    def calculate_priors(self, o, si, s, ca, fe, co, ni, mg, ti, cr, c, lum, v_inner):
        priors = 0

        if o <= 0.:
            priors += np.inf

        if s > si:
            priors += np.inf

        return priors



parameters = ['model.abundances.o',
    'model.abundances.si',
    'model.abundances.s',
    'model.abundances.ca',
    'model.abundances.fe',
    'model.abundances.co',
    'model.abundances.ni',
    'model.abundances.mg',
    'model.abundances.ti',
    'model.abundances.cr',
    'model.abundances.c',
    'supernova.luminosity_requested',
    'model.structure.velocity.item0']

artisw, artisf = loadtxt(artis_filename, unpack=1)
artisw = artisw[198:]
artisf = artisf[198:]
artisf *= ((1 * u.Mpc)**2).to(u.cm**2).value * 4 * np.pi
artisu = artisf * 0.01

tardis_model = assemble_tardis_model('tardis_artis_fits.yml', parameters)
full_model = (tardis_model | SimpleTARDISUncertaintyModel(artisw)
| NormRedSpectrum() | LogLikelihood(artisw, artisf, artisu, wavelength_end=9000.))
deartis = ArtisDELogLikelihood(full_model, log_fname, spec_dir=spec_dir)