allbandsimcubev2.py

# !/usr/bin/env python
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from astropy.io import fits
from multiprocessing import Pool, RLock
from copy import deepcopy
from jwst.stpipe import Pipeline
from jwst import datamodels
from glob import glob
import json, os, time
from jwst.associations import Association

try:
    pool.close()
except Exception:
    pass


def runsim(kw):
    global lock

    scenefile=kw.get('scenefile','scene.ini')
    local = kw.get('local', None)
    channel=kw.get('channel',1)
    band=kw.get('band','short')
    frames=kw.get('frames',10)
    dither=kw.get('dither',True)
    ditherndx=kw.get('ditherndx',1)
    ndither=kw.get('ndither',2)
    detector=kw.get('detector','BOTH')
    noise=kw.get('noise', False)
    simulatorfile=kw.get('simulatorfile','simulator.ini')
    #band & channel for POP/dither center only
    # detector -> BOTH SW LW
    band=band.upper()
    subband=str(channel)+band
    
    cfgpath=kw.get('configpath',subband.upper())
    
    args={'name':"Default Simulation", 'rel_obsdate':0.0, 'scene':scenefile,
          'POP':'MRS', 'ConfigPath':'MRS_'+cfgpath, 'Dither':str(dither), 'StartInd':ditherndx,
          'NDither':ndither, 'DitherPat':"mrs_recommended_dither.dat",
          'disperser':band.upper(), 'detector':detector, 'mrs_mode':'FAST',
          'mrs_exposures':1, 'mrs_integrations':1, 'mrs_frames':frames,
          'filter':"F1130W", 'readDetect':'FULL','ima_mode':'FAST',
          'ima_exposures':1, 'ima_integrations':1, 'ima_frames':10}
    
    from mirisim.config_parser import SimConfig, SimulatorConfig, SceneConfig
    from mirisim import MiriSimulation

    with lock:
        time.sleep(1)
        scenecfg = SceneConfig(scenefile)
        simcfg = SimConfig.makeSim(**args)
        simucfg = SimulatorConfig(simulatorfile)
        noisechar = 'T' if noise else 'F'
        simucfg['SCASim']['include_readnoise'] = noisechar
        simucfg['SCASim']['include_poisson'] = noisechar
        wd = kw.get('wd','')
        if wd:
            wdo=os.getcwd()
            if not os.path.exists(wd):
                os.mkdir(wd)
            os.chdir(wd)
        ms = MiriSimulation(sim_config=simcfg, scene_config=scenecfg, local=local,
                            simulator_config=simucfg, loglevel='DEBUG')
        time.sleep(1)
    ms.run()
    if wd:
        os.chdir(wdo)
    return ms.path_out

def init(l):
    global lock
    lock = l


class MRSMetaPipeline(Pipeline):
    from jwst.dq_init import dq_init_step
    from jwst.saturation import saturation_step
    from jwst.linearity import linearity_step
    from jwst.rscd import rscd_step
    from jwst.lastframe import lastframe_step
    #from jwst.dark_current import dark_current_step
    from jwst.refpix import refpix_step
    from jwst.persistence import persistence_step
    from jwst.jump import jump_step
    from jwst.ramp_fitting import ramp_fit_step

    from jwst.assign_wcs import assign_wcs_step
    from jwst.flatfield import flat_field_step
    #from jwst.straylight import straylight_step
    from jwst.fringe import fringe_step
    from jwst.photom import photom_step
    from jwst.cube_build import cube_build_step

    step_defs = {'dq_init': dq_init_step.DQInitStep,
                 'saturation': saturation_step.SaturationStep,
                 'linearity': linearity_step.LinearityStep,
                 #'rscd': rscd_step.RSCD_Step,  #drifts & latency
                 'lastframe': lastframe_step.LastFrameStep,
                 #'dark_current': dark_current_step.DarkCurrentStep,
                 'refpix': refpix_step.RefPixStep,
                 'persistence': persistence_step.PersistenceStep,
                 'jump': jump_step.JumpStep,
                 'ramp_fit': ramp_fit_step.RampFitStep,
                 'assign_wcs': assign_wcs_step.AssignWcsStep,
                 'flat_field': flat_field_step.FlatFieldStep, # include_flat
                 #'straylight': straylight_step.StraylightStep,
                 'fringe': fringe_step.FringeStep,
                 'photom': photom_step.PhotomStep,
                 'cube_build': cube_build_step.CubeBuildStep}

    # start the actual processing
    def process(self, input):
        # open the input
        input = datamodels.open(input)

        # apply steps
        input = self.dq_init(input)
        input = self.saturation(input)
        input = self.linearity(input)
        # input = self.rscd(input)  #features off in mirisim
        input = self.lastframe(input)
        # input = self.dark_current(input)  #Not compatible with sim
        input = self.refpix(input)
        input = self.persistence(input)
        input = self.jump(input)
        input = self.ramp_fit(input)
        input = self.assign_wcs(input[0] if isinstance(input, tuple) else input)  #rc3 compatibility test
        input = self.flat_field(input)
        # input = self.straylight(input)  #Not simulated by mirisim
        input = self.fringe(input)
        input = self.photom(input)
        input = self.cube_build(input)
        return input


class MRSCubePipeline(Pipeline):
    from jwst.cube_build import cube_build_step
    step_defs = {'cube_build': cube_build_step.CubeBuildStep}

    def process(self, input):
        input = datamodels.open(input)
        return self.cube_build(input)


def run_pipe(f):
    pipe = MRSMetaPipeline(steps={'photom': {'output_dir': os.path.dirname(f),
                                             'output_file': f[:-4] + 'photom.fits'},
                                  'cube_build': {'output_dir': os.path.dirname(f),
                                                 'output_file': f[:-4] + 'cube_radec.fits'}})
    pipe(f)


def cubearg(f, channel='', band='', ofile=None, odir=None, radec=True, dither4=False):
    SROID2 = {sb:v for v,sb in zip([0.15000001,  0.15000001,  0.15000001,  0.15000001,  0.22500001,
                                    0.22500001,  0.22500001,  0.22500001,  0.22500001,  0.60000002,
                                    0.45000002,  0.45000002],
                                   ['1SHORT', '1MEDIUM', '1LONG', '2SHORT', '2MEDIUM', '2LONG',
                                    '3SHORT', '3MEDIUM', '3LONG', '4SHORT', '4MEDIUM', '4LONG'])}
    SROID2['']=1.0
    odir = os.path.dirname(f) if odir is None else odir
    if not ofile:
        ofile = os.path.basename(f[:-4])+'cube{}{}{}.fits'.format(channel, band,
                                               'radec' if radec else 'ab')
    sb=str(channel)+band.upper()
    
    args={'output_dir': odir, 'output_file': ofile,
          'coord_system': 'alpha-beta', 'rois':SROID2[sb],
          'channel': str(channel), 'band':band.upper()}
    if dither4 or not sb:
        args.pop('rois')
    if radec:
        args.pop('coord_system')
    if not (channel and band):
        args.pop('channel')
        args.pop('band')
    return (f,args)

def runCubePipe((f, args)):
    MRSCubePipeline(steps={'cube_build': args}).run(f)





run=False
if run:

    #Do Simulations
    #arglist=[{'ditherndx':dind, 'ndither':1, 'band':band}
    #         for dind in range(1,5) for band in ('short','medium','long')]
    os.chdir('psfnormed')
    import time
    tic=time.time()
    arglist=[{'configpath':'1SHORT','ndither':2, 'band':band, 'frames':10, 'noise':False,
              'local':True}
             for band in ('short','medium','long')]
    lock=RLock()
    pool = Pool(3, initializer=init, initargs=(lock,))
    pool.map(runsim, arglist)
    pool.close()


    #Run the pipeline
    files = [f for f in glob('./*mirisim/det_images/*exp1.fits')]
    run_pipe(files[0])
    pool = Pool(6)
    pool.map(run_pipe, files[1:])
    pool.close()

    #Make Photometry associations
    id='dither2'
    from jwst.associations import asn_from_list, load_asn
    files = [f for f in glob('./*mirisim/det_images/*exp1.photom.fits')]
    asn_from_list.Main('-o '+id+'_asn.json --product-name dither '+
                       ' '.join([f for f in files if 'seq3' not in f and 'seq4' not in f]))
#    asn_from_list.Main('-o dither4_asn.json --product-name dither '+ ' '.join(files))

    #Make band cubes
    arglist = [cubearg(id+'_asn.json',channel=c,band=b,
                   odir=os.path.abspath(os.path.curdir))
           for c in range(1,5) for b in ('SHORT', 'MEDIUM', 'LONG')]
    arglist.append(cubearg(id+'_asn.json', odir=os.path.abspath(os.path.curdir)))
    pool = Pool(4)
    pool.map(runCubePipe, arglist[1:])
    pool.close()
    
    toc=time.time()

    #Make 4dither sim res
    arglist = [cubearg('dither4_asn.json',channel=c,band=b,dither4=True,
                   odir=os.path.abspath(os.path.curdir))
           for c in range(1,5) for b in ('SHORT', 'MEDIUM', 'LONG')]
    arglist.append(cubearg('dither4_asn.json', dither4=True, odir=os.path.abspath(os.path.curdir)))
    pool = Pool(4)
    pool.map(runCubePipe, arglist)
    pool.close()


    # scenes=['scene0.ini','scene0nol.ini','scene10.ini','scene10nol.ini','scene100.ini','scene100nol.ini']
    #
    # arglist=[{'configpath':'1SHORT','ndither':2, 'band':'medium', 'frames':f,'noise':n,
    #           'scenefile':sf,'wd':'{}_f{}_n{}'.format(sf[:-4],f,n)} for sf in scenes for f in (10,65) for n in (False,True)]
    # wds=['{}_f{}_n{}'.format(sf[:-4],f,n) for sf in scenes for f in (10,65) for n in (False,True)]
    #
    # tic=time.time()
    # lock=RLock()
    # pool = Pool(6, initializer=init, initargs=(lock,))
    # pool.map(runsim, arglist)
    # pool.close()
    # toc=time.time()
    # print toc-tic
    #
    #
    # files = [f for f in glob('./*/*mirisim/det_images/*exp1.fits')]
    # run_pipe(files[0])
    # pool = Pool(6)
    # pool.map(run_pipe, files[1:])
    # pool.close()
    #
    #
    # id='dither2'
    # from jwst.associations import asn_from_list, load_asn
    # for wd in wds:
    #     files = [f for f in glob('./'+wd+'/*mirisim/det_images/*exp1.photom.fits')]
    #     asn_from_list.Main('-o '+os.path.join('.',wd+'_'+id+'_asn.json')+' --product-name dither '+
    #                        ' '.join([f for f in files if 'seq3' not in f and 'seq4' not in f]))
    #
    # bands=('MEDIUM',)
    # chans=[3]
    # arglist = [cubearg(os.path.join('./',wd+'_'+id+'_asn.json'),channel=c,band=b,
    #                    odir=os.path.abspath(wd)) for wd in wds for c in chans for b in bands]
    # #arglist.append(cubearg(id+'_asn.json', odir=os.path.abspath(os.path.curdir)))
    # pool = Pool(4)
    # pool.map(runCubePipe, arglist)
    # pool.close()
    #