Exemplo n.º 1
0
def simulate_frame(night, expid, camera, ccdshape=None, **kwargs):
    """
    Simulate a single frame, including I/O

    Args:
        night: YEARMMDD string
        expid: integer exposure ID
        camera: b0, r1, .. z9

    Options:
        ccdshape = (npix_y, npix_x) primarily used to limit memory while testing

    Additional keyword args are passed to pixsim.simulate()

    Reads:
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/simspec-{expid}.fits

    Writes:
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/simpix-{camera}-{expid}.fits
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/desi-{expid}.fits
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/pix-{camera}-{expid}.fits

    For a lower-level pixel simulation interface that doesn't perform I/O,
    see pixsim.simulate()
    """
    #- night, expid, camera -> input file names
    simspecfile = io.findfile('simspec', night=night, expid=expid)

    #- Read inputs
    psf = desimodel.io.load_psf(camera[0])
    simspec = io.read_simspec(simspecfile)

    #- Trim effective CCD size; mainly to limit memory for testing
    if ccdshape is not None:
        psf.npix_y, psf.npix_x = ccdshape

    if 'cosmics' in kwargs:
        shape = (psf.npix_y, psf.npix_x)
        kwargs['cosmics'] = io.read_cosmics(kwargs['cosmics'],
                                            expid,
                                            shape=shape)

    image, rawpix, truepix = simulate(camera, simspec, psf, **kwargs)

    #- Outputs; force "real" data files into simulation directory
    simpixfile = io.findfile('simpix', night=night, expid=expid, camera=camera)
    io.write_simpix(simpixfile, truepix, camera=camera, meta=image.meta)

    simdir = io.simdir(night=night)
    rawfile = desispec.io.findfile('desi', night=night, expid=expid)
    rawfile = os.path.join(simdir, os.path.basename(rawfile))
    desispec.io.write_raw(rawfile, rawpix, image.meta, camera=camera)

    pixfile = desispec.io.findfile('pix',
                                   night=night,
                                   expid=expid,
                                   camera=camera)
    pixfile = os.path.join(simdir, os.path.basename(pixfile))
    desispec.io.write_image(pixfile, image)
Exemplo n.º 2
0
def simulate_exposure(simspecfile, rawfile, cameras=None,
        ccdshape=None, simpixfile=None, addcosmics=None, comm=None,
        **kwargs):
    """
    Simulate frames from an exposure, including I/O

    Args:
        simspecfile: input simspec format file with spectra
        rawfile: output raw data file to write

    Options:
        cameras: str or list of str, e.g. b0, r1, .. z9
        ccdshape: (npix_y, npix_x) primarily used to limit memory while testing
        simpixfile: output file for noiseless truth pixels
        addcosmics: if True (must be specified via command input), add cosmics from real data
        comm: MPI communicator object

    Additional keyword args are passed to pixsim.simulate()

    For a lower-level pixel simulation interface that doesn't perform I/O,
    see pixsim.simulate()

    Note: call desi_preproc or desispec.preproc.preproc to pre-process the
    output desi*.fits file for overscan subtraction, noise estimation, etc.
    """
    #- Split communicator by nodes; each node processes N frames
    #- Assumes / requires equal number of ranks per node
    if comm is not None:
        rank, size = comm.rank, comm.size
        num_nodes = mpi_count_nodes(comm)
        comm_node, node_index, num_nodes = mpi_split_by_node(comm, 1)
        node_rank = comm_node.rank
        node_size = comm_node.size
    else:
        log.debug('Not using MPI')
        rank, size = 0, 1
        comm_node = None
        node_index = 0
        num_nodes = 1
        node_rank = 0
        node_size = 1

    if rank == 0:
        log.debug('Starting simulate_exposure at {}'.format(asctime()))

    if cameras is None:
        if rank == 0:
            from astropy.io import fits
            fibermap = fits.getdata(simspecfile, 'FIBERMAP')
            cameras = io.fibers2cameras(fibermap['FIBER'])
            log.debug('Found cameras {} in input simspec file'.format(cameras))
            if len(cameras) % num_nodes != 0:
                raise ValueError('Number of cameras {} should be evenly divisible by number of nodes {}'.format(
                    len(cameras), num_nodes))

    if comm is not None:
        cameras = comm.bcast(cameras, root=0)

    #- Fail early if camera alreaady in output file
    if rank == 0 and os.path.exists(rawfile):
        from astropy.io import fits
        err = False
        fx = fits.open(rawfile)
        for camera in cameras:
            if camera in fx:
                log.error('Camera {} already in {}'.format(camera, rawfile))
                err = True
        if err:
            raise ValueError('Some cameras already in output file')

    #- Read simspec input; I/O layer handles MPI broadcasting
    if rank == 0:
        log.debug('Reading simspec at {}'.format(asctime()))

    mycameras = cameras[node_index::num_nodes]
    if node_rank == 0:
        log.info("Assigning cameras {} to comm_exp node {}".format(mycameras, node_index))

    simspec = io.read_simspec(simspecfile, cameras=mycameras,
        readflux=False, comm=comm)
    night = simspec.header['NIGHT']
    expid = simspec.header['EXPID']

    if rank == 0:
        log.debug('Reading PSFs at {}'.format(asctime()))

    psfs = dict()
    #need to initialize previous channel
    previous_channel = 'a'
    for camera in mycameras:
        #- Note: current PSF object can't be pickled and thus every
        #- rank must read it instead of rank 0 read + bcast
        channel = camera[0]
        if channel not in psfs:
            log.info('Reading {} PSF at {}'.format(channel, asctime()))
            psfs[channel] = desimodel.io.load_psf(channel)

            #- Trim effective CCD size; mainly to limit memory for testing
            if ccdshape is not None:
                psfs[channel].npix_y, psfs[channel].npix_x = ccdshape

        psf = psfs[channel]

        cosmics=None
        #avoid re-broadcasting cosmics if we can
        if previous_channel != channel:
            if (addcosmics is True) and (node_rank == 0):
                cosmics_file = io.find_cosmics(camera, simspec.header['EXPTIME'])
                log.info('Reading cosmics templates {} at {}'.format(
                    cosmics_file, asctime()))
                shape = (psf.npix_y, psf.npix_x)
                cosmics = io.read_cosmics(cosmics_file, expid, shape=shape)
            if (addcosmics is True) and (comm_node is not None):
                if node_rank == 0:
                    log.info('Broadcasting cosmics at {}'.format(asctime()))
                cosmics = comm_node.bcast(cosmics, root=0)
            else:
                log.debug("Cosmics not requested")

        if node_rank == 0:
            log.info("Starting simulate for camera {} on node {}".format(camera,node_index))
        image, rawpix, truepix = simulate(camera, simspec, psf, comm=comm_node, preproc=False, cosmics=cosmics, **kwargs)

        #- Use input communicator as barrier since multiple sub-communicators
        #- will write to the same output file
        if rank == 0:
            log.debug('Writing outputs at {}'.format(asctime()))

        tmprawfile = rawfile + '.tmp'
        if comm is not None:
            for i in range(comm.size):
                if (i == comm.rank) and (comm_node.rank == 0):
                    desispec.io.write_raw(tmprawfile, rawpix, image.meta,
                                          camera=camera)
                    if simpixfile is not None:
                        io.write_simpix(simpixfile, truepix, camera=camera,
                                        meta=image.meta)
                comm.barrier()
        else:
            desispec.io.write_raw(tmprawfile, rawpix, image.meta, camera=camera)
            if simpixfile is not None:
                io.write_simpix(simpixfile, truepix, camera=camera,
                                meta=image.meta)

        if rank == 0:
            log.info('Wrote {}'.format(rawfile))
            log.debug('done at {}'.format(asctime()))

        previous_channel = channel

    #- All done; rename temporary raw file to final location
    if comm is None or comm.rank == 0:
        os.rename(tmprawfile, rawfile)
Exemplo n.º 3
0
def simulate_exposure(simspecfile, rawfile, cameras=None,
        ccdshape=None, simpixfile=None, addcosmics=None, comm=None,
        **kwargs):
    """
    Simulate frames from an exposure, including I/O

    Args:
        simspecfile: input simspec format file with spectra
        rawfile: output raw data file to write

    Options:
        cameras: str or list of str, e.g. b0, r1, .. z9
        ccdshape: (npix_y, npix_x) primarily used to limit memory while testing
        simpixfile: output file for noiseless truth pixels
        addcosmics: if True (must be specified via command input), add cosmics from real data
        comm: MPI communicator object

    Additional keyword args are passed to pixsim.simulate()

    For a lower-level pixel simulation interface that doesn't perform I/O,
    see pixsim.simulate()

    Note: call desi_preproc or desispec.preproc.preproc to pre-process the
    output desi*.fits file for overscan subtraction, noise estimation, etc.
    """
    #- Split communicator by nodes; each node processes N frames
    #- Assumes / requires equal number of ranks per node
    if comm is not None:
        rank, size = comm.rank, comm.size
        num_nodes = mpi_count_nodes(comm)
        comm_node, node_index, num_nodes = mpi_split_by_node(comm, 1)
        node_rank = comm_node.rank
        node_size = comm_node.size
    else:
        log.debug('Not using MPI')
        rank, size = 0, 1
        comm_node = None
        node_index = 0
        num_nodes = 1
        node_rank = 0
        node_size = 1
    
    if rank == 0:
        log.debug('Starting simulate_exposure at {}'.format(asctime()))

    if cameras is None:
        if rank == 0:
            from astropy.io import fits
            fibermap = fits.getdata(simspecfile, 'FIBERMAP')
            cameras = io.fibers2cameras(fibermap['FIBER'])
            log.debug('Found cameras {} in input simspec file'.format(cameras))
            if len(cameras) % num_nodes != 0:
                raise ValueError('Number of cameras {} should be evenly divisible by number of nodes {}'.format(
                    len(cameras), num_nodes))

    if comm is not None:
        cameras = comm.bcast(cameras, root=0)

    #- Fail early if camera alreaady in output file
    if rank == 0 and os.path.exists(rawfile):
        from astropy.io import fits
        err = False
        fx = fits.open(rawfile)
        for camera in cameras:
            if camera in fx:
                log.error('Camera {} already in {}'.format(camera, rawfile))
                err = True
        if err:
            raise ValueError('Some cameras already in output file')

    #- Read simspec input; I/O layer handles MPI broadcasting
    if rank == 0:
        log.debug('Reading simspec at {}'.format(asctime()))

    mycameras = cameras[node_index::num_nodes]
    if node_rank == 0:
        log.info("Assigning cameras {} to comm_exp node {}".format(mycameras, node_index))

    simspec = io.read_simspec(simspecfile, cameras=mycameras,
        readflux=False, comm=comm)
    night = simspec.header['NIGHT']
    expid = simspec.header['EXPID']

    if rank == 0:
        log.debug('Reading PSFs at {}'.format(asctime()))

    psfs = dict()
    #need to initialize previous channel
    previous_channel = 'a'
    for camera in mycameras:
        #- Note: current PSF object can't be pickled and thus every
        #- rank must read it instead of rank 0 read + bcast
        channel = camera[0]
        if channel not in psfs:
            psfs[channel] = desimodel.io.load_psf(channel)

            #- Trim effective CCD size; mainly to limit memory for testing
            if ccdshape is not None:
                psfs[channel].npix_y, psfs[channel].npix_x = ccdshape

        psf = psfs[channel]

        cosmics=None
        #avoid re-broadcasting cosmics if we can
        if previous_channel != channel:
            if (addcosmics is True) and (node_rank == 0):
                cosmics_file = io.find_cosmics(camera, simspec.header['EXPTIME'])
                log.info('cosmics templates {}'.format(cosmics_file))
                shape = (psf.npix_y, psf.npix_x)
                cosmics = io.read_cosmics(cosmics_file, expid, shape=shape)
            if (addcosmics is True) and (comm_node is not None):
                cosmics = comm_node.bcast(cosmics, root=0)
            else:
                log.debug("Cosmics not requested")

        if node_rank == 0: 
            log.info("Starting simulate for camera {} on node {}".format(camera,node_index)) 
        image, rawpix, truepix = simulate(camera, simspec, psf, comm=comm_node, preproc=False, cosmics=cosmics, **kwargs)

        #- Use input communicator as barrier since multiple sub-communicators
        #- will write to the same output file
        if rank == 0:
            log.debug('Writing outputs at {}'.format(asctime()))

        tmprawfile = rawfile + '.tmp'
        if comm is not None:
            for i in range(comm.size):
                if (i == comm.rank) and (comm_node.rank == 0):
                    desispec.io.write_raw(tmprawfile, rawpix, image.meta,
                                          camera=camera)
                    if simpixfile is not None:
                        io.write_simpix(simpixfile, truepix, camera=camera,
                                        meta=image.meta)
                comm.barrier()
        else:
            desispec.io.write_raw(tmprawfile, rawpix, image.meta, camera=camera)
            if simpixfile is not None:
                io.write_simpix(simpixfile, truepix, camera=camera,
                                meta=image.meta)

        if rank == 0:
            log.info('Wrote {}'.format(rawfile))
            log.debug('done at {}'.format(asctime()))

        previous_channel = channel

    #- All done; rename temporary raw file to final location
    if comm is None or comm.rank == 0:
        os.rename(tmprawfile, rawfile)
Exemplo n.º 4
0
def main(args=None):
    log.info('Starting pixsim at {}'.format(asctime()))
    if isinstance(args, (list, tuple, type(None))):
        args = parse(args)

    if args.verbose:
        import logging
        log.setLevel(logging.DEBUG)

    if args.mpi:
        from mpi4py import MPI
        mpicomm = MPI.COMM_WORLD
        log.debug('Using mpi4py MPI communicator')
    else:
        from desisim.util import _FakeMPIComm
        log.debug('Using fake MPI communicator')
        mpicomm = _FakeMPIComm()

    log.info('Starting pixsim rank {} at {}'.format(mpicomm.rank, asctime()))
    log.debug('MPI rank {} size {} / {} {}'.format(mpicomm.rank, mpicomm.size,
                                                   mpicomm.Get_rank(),
                                                   mpicomm.Get_size()))

    #- Pre-flight check that these cameras haven't been done yet
    if (mpicomm.rank == 0) and (not args.overwrite) and os.path.exists(
            args.rawfile):
        log.debug('Checking if cameras are already in output file')
        from astropy.io import fits
        fx = fits.open(args.rawfile)
        oops = False
        for camera in args.cameras:
            if camera.upper() in fx:
                log.error('Camera {} already in {}'.format(
                    camera, args.rawfile))
                oops = True

        fx.close()
        if oops:
            log.fatal('Exiting due to repeat cameras already in output file')
            mpicomm.Abort()

    ncameras = len(args.cameras)
    if ncameras % mpicomm.size != 0:
        log.fatal('Processing cameras {}'.format(args.cameras))
        log.fatal(
            'Number of cameras {} must be evenly divisible by MPI size {}'.
            format(ncameras, mpicomm.size))
        mpicomm.Abort()

    #- Use original seed to generate different random seeds for each MPI rank
    np.random.seed(args.seed)
    while True:
        seeds = np.random.randint(0, 2**32 - 1, size=mpicomm.size)
        if np.unique(seeds).size == mpicomm.size:
            random.seed(seeds[mpicomm.rank])
            np.random.seed(seeds[mpicomm.rank])
            break

    if args.psf is not None:
        from specter.psf import load_psf
        psf = load_psf(args.psf)

    simspec = io.read_simspec(args.simspec)

    if args.fibermap:
        fibermap = desispec.io.read_fibermap(args.fibermap)
        fibers = fibermap['FIBER']
        if args.nspec is not None:
            fibers = fibers[0:args.nspec]
    else:
        fibers = None

    if args.overwrite and os.path.exists(args.rawfile):
        log.debug('removing {}'.format(args.rawfile))
        os.remove(args.rawfile)

    if args.overwrite and os.path.exists(args.simpixfile):
        log.debug('removing {}'.format(args.simpixfile))
        os.remove(args.simpixfile)

    for i in range(mpicomm.rank, ncameras, mpicomm.size):
        camera = args.cameras[i]
        log.debug('Rank {} processing camera {}'.format(mpicomm.rank, camera))
        channel = camera[0].lower()
        assert channel in (
            'b', 'r',
            'z'), "Unknown camera {} doesn't start with b,r,z".format(camera)

        #- Read inputs for this camera
        if args.psf is None:
            psf = desimodel.io.load_psf(channel)
            if args.ccd_npix_x is not None:
                psf.npix_x = args.ccd_npix_x
            if args.ccd_npix_y is not None:
                psf.npix_y = args.ccd_npix_y

        if args.cosmics:
            if args.cosmics_file is None:
                cosmics_file = io.find_cosmics(camera,
                                               simspec.header['EXPTIME'],
                                               cosmics_dir=args.cosmics_dir)
                log.info('cosmics templates {}'.format(cosmics_file))
            else:
                cosmics_file = args.cosmics_file

            shape = (psf.npix_y, psf.npix_x)
            cosmics = io.read_cosmics(cosmics_file, args.expid, shape=shape)
        else:
            cosmics = None

        #- Do the actual simulation
        image, rawpix, truepix = desisim.pixsim.simulate(camera,
                                                         simspec,
                                                         psf,
                                                         fibers=fibers,
                                                         nspec=args.nspec,
                                                         ncpu=args.ncpu,
                                                         cosmics=cosmics,
                                                         wavemin=args.wavemin,
                                                         wavemax=args.wavemax,
                                                         preproc=False)

        #- Synchronize with other MPI threads before continuing with output
        mpicomm.Barrier()

        #- Loop over MPI ranks, letting each one take its turn writing output
        for rank in range(mpicomm.size):
            if rank == mpicomm.rank:
                desispec.io.write_raw(args.rawfile,
                                      rawpix,
                                      camera=camera,
                                      header=image.meta,
                                      primary_header=simspec.header)
                log.info('Wrote {} image to {}'.format(camera, args.rawfile))
                io.write_simpix(args.simpixfile,
                                truepix,
                                camera=camera,
                                meta=simspec.header)
                log.info('Wrote {} image to {}'.format(camera,
                                                       args.simpixfile))
                mpicomm.Barrier()
            else:
                mpicomm.Barrier()

    #- Call preproc separately from pixsim.simulate to ensure that it is
    #- done identically to real data.
    #- Note: This does lose the advantage of parallel preprocessing; we could
    #- change this if that is problematic
    if args.preproc and mpicomm.rank == 0:
        log.info('Preprocessing raw -> pix files')
        from desispec.scripts import preproc
        preproc_opts = ['--infile', args.rawfile, '--outdir', args.preproc_dir]
        preproc_opts += ['--cameras', ','.join(args.cameras)]
        preproc.main(preproc.parse(preproc_opts))

    if mpicomm.rank == 0:
        log.info('Finished pixsim {} expid {} at {}'.format(
            args.night, args.expid, asctime()))
Exemplo n.º 5
0
def main(args, comm=None):
    if args.verbose:
        import logging
        log.setLevel(logging.DEBUG)

    rank = 0
    nproc = 1
    if comm is not None:
        rank = comm.rank
        nproc = comm.size

    if rank == 0:
        log.info('Starting pixsim at {}'.format(asctime()))

    #- Pre-flight check that these cameras haven't been done yet
    if (rank == 0) and (not args.overwrite) and os.path.exists(args.rawfile):
        log.debug('Checking if cameras are already in output file')
        from astropy.io import fits
        fx = fits.open(args.rawfile)
        oops = False
        for camera in args.cameras:
            if camera.upper() in fx:
                log.error('Camera {} already in {}'.format(camera, 
                    args.rawfile))
                oops = True
        fx.close()
        if oops:
            log.fatal('Exiting due to repeat cameras already in output file')
            if comm is not None:
                comm.Abort()
            else:
                sys.exit(1)

    ncamera = len(args.cameras)

    comm_group = comm
    comm_rank = None
    group = 0
    ngroup = 1
    group_rank = 0
    if comm is not None:
        if args.mpi_camera > 1:
            ngroup = int(comm.size / args.mpi_camera)
            group = int(comm.rank / args.mpi_camera)
            group_rank = comm.rank % args.mpi_camera
            comm_group = comm.Split(color=group, key=group_rank)
            comm_rank = comm.Split(color=group_rank, key=group)
        else:
            group = comm.rank
            ngroup = comm.size
            comm_group = MPI.COMM_SELF
            comm_rank = comm

    mycameras = np.array_split(np.arange(ncamera, dtype=np.int32), 
        ngroup)[group]

    rawtemp = "{}.tmp".format(args.rawfile)
    simpixtemp = "{}.tmp".format(args.simpixfile)

    if rank == 0:
        if args.overwrite and os.path.exists(args.rawfile):
            log.debug('removing {}'.format(args.rawfile))
            os.remove(args.rawfile)

        if args.overwrite and os.path.exists(args.simpixfile):
            log.debug('removing {}'.format(args.simpixfile))
            os.remove(args.simpixfile)

        # cleanup stale temp files
        if os.path.isfile(rawtemp):
            os.remove(rawtemp)
    
    if comm is not None:
        comm.barrier()

    psf = None
    if args.psf is not None:
        from specter.psf import load_psf
        psf = load_psf(args.psf)

    simspec = None
    if rank == 0:
        simspec = io.read_simspec(args.simspec)
    if comm is not None:
        # Broadcast one array at a time, since this is a 
        # very large object.
        flv = None
        wv = None
        pht = None
        if simspec is not None:
            flv = simspec.flavor
            wv = simspec.wave
            pht = simspec.phot
        flv = comm.bcast(flv, root=0)
        wv = comm.bcast(wv, root=0)
        pht = comm.bcast(pht, root=0)
        if simspec is None:
            simspec = SimSpec(flv, wv, pht)
        simspec.flux = comm.bcast(simspec.flux, root=0)
        simspec.skyflux = comm.bcast(simspec.skyflux, root=0)
        simspec.skyphot = comm.bcast(simspec.skyphot, root=0)
        simspec.metadata = comm.bcast(simspec.metadata, root=0)
        simspec.fibermap = comm.bcast(simspec.fibermap, root=0)
        simspec.obs = comm.bcast(simspec.obs, root=0)
        simspec.header = comm.bcast(simspec.header, root=0)

    fibers = None
    if args.fibermap:
        if rank == 0:
            fibermap = desispec.io.read_fibermap(args.fibermap)
            fibers = fibermap['FIBER']
            if args.nspec is not None:
                fibers = fibers[0:args.nspec]
        if comm is not None:
            fibers = comm.bcast(fibers, root=0)

    # Use original seed to generate different random seeds for each camera
    np.random.seed(args.seed)
    seeds = np.random.randint(0, 2**32-1, size=ncamera)

    image = {}
    rawpix = {}
    truepix = {}

    for c in mycameras:
        camera = args.cameras[c]
        if group_rank == 0:
            log.debug('Processing camera {}'.format(camera))
        channel = camera[0].lower()

        # Set the seed for this camera (regardless of which process is
        # performing the simulation).
        np.random.seed(seeds[c])

        # Get the random cosmic expids.  The actual values will be
        # remapped internally with the modulus operator.
        cosexpid = np.random.randint(0, 100, size=1)[0]

        # Read inputs for this camera.  Unfortunately psf
        # objects are not serializable, so we read it on all
        # processes.
        if args.psf is None:
            psf = desimodel.io.load_psf(channel)
            if args.ccd_npix_x is not None:
                psf.npix_x = args.ccd_npix_x
            if args.ccd_npix_y is not None:
                psf.npix_y = args.ccd_npix_y

        cosmics = None
        if args.cosmics:
            if group_rank == 0:
                if args.cosmics_file is None:
                    cosmics_file = io.find_cosmics(camera, 
                        simspec.header['EXPTIME'],
                        cosmics_dir=args.cosmics_dir)
                    log.info('cosmics templates {}'.format(cosmics_file))
                else:
                    cosmics_file = args.cosmics_file

                shape = (psf.npix_y, psf.npix_x)
                cosmics = io.read_cosmics(cosmics_file, cosexpid, 
                    shape=shape)
            if comm_group is not None:
                cosmics = comm_group.bcast(cosmics, root=0)

        #- Do the actual simulation
        image[camera], rawpix[camera], truepix[camera] = \
            desisim.pixsim.simulate(camera, simspec, psf, fibers=fibers,
            nspec=args.nspec, ncpu=args.ncpu, cosmics=cosmics,
            wavemin=args.wavemin, wavemax=args.wavemax, preproc=False,
            comm=comm_group)

        if args.psf is None:
            del psf

    # Wait for all processes to finish their cameras
    if comm is not None:
        comm.barrier()

    # Write the cameras in order.  Only the rank zero process in each
    # group has the data.
    for c in np.arange(ncamera, dtype=np.int32):
        camera = args.cameras[c]
        if c in mycameras:
            if group_rank == 0:
                desispec.io.write_raw(rawtemp, rawpix[camera], 
                    camera=camera, header=image[camera].meta, 
                    primary_header=simspec.header)
                log.info('Wrote {} image to {}'.format(camera, args.rawfile))
                io.write_simpix(simpixtemp, truepix[camera], 
                    camera=camera, meta=simspec.header)
                log.info('Wrote {} image to {}'.format(camera, 
                    args.simpixfile))
        if comm is not None:
            comm.barrier()

    # Move temp files into place
    if rank == 0:
        os.rename(simpixtemp, args.simpixfile)
        os.rename(rawtemp, args.rawfile)
    if comm is not None:
        comm.barrier()

    # Apply preprocessing
    if args.preproc:
        if rank == 0:
            log.info('Preprocessing raw -> pix files')
        from desispec.scripts import preproc
        if len(mycameras) > 0:
            if group_rank == 0:
                for c in mycameras:
                    camera = args.cameras[c]
                    pixfile = desispec.io.findfile('pix', night=args.night,
                        expid=args.expid, camera=camera)
                    preproc_opts = ['--infile', args.rawfile, '--outdir',
                        args.preproc_dir, '--pixfile', pixfile]
                    preproc_opts += ['--cameras', camera]
                    preproc.main(preproc.parse(preproc_opts))

    if comm is not None:
        comm.barrier()
    
    # Python is terrible with garbage collection, but at least
    # encourage it...
    del image
    del rawpix
    del truepix

    if rank == 0:
        log.info('Finished pixsim {} expid {} at {}'.format(args.night, args.expid, asctime()))
Exemplo n.º 6
0
def simulate(night, expid, camera, nspec=None, verbose=False, ncpu=None,
    trimxy=False, cosmics=None):
    """
    Run pixel-level simulation of input spectra
    
    Args:
        night (string) : YEARMMDD
        expid (integer) : exposure id
        camera (str) : e.g. b0, r1, z9
        nspec (int) : number of spectra to simulate
        verbose (boolean) : if True, print status messages
        ncpu (int) : number of CPU cores to use in parallel
        trimxy (boolean) : trim image to just pixels with input signal
        cosmics (str) : filename with dark images with cosmics to add

    Reads:
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/simspec-{expid}.fits
        
    Writes:
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/simpix-{camera}-{expid}.fits
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/pix-{camera}-{expid}.fits
    """
    if verbose:
        print "Reading input files"

    channel = camera[0].lower()
    ispec = int(camera[1])
    assert channel in 'brz'
    assert 0 <= ispec < 10

    #- Load DESI parameters
    params = desimodel.io.load_desiparams()
    nfibers = params['spectro']['nfibers']

    #- Load simspec file
    simfile = io.findfile('simspec', night=night, expid=expid)
    simspec = io.read_simspec(simfile)
    wave = simspec.wave[channel]
    if simspec.skyphot is not None:
        phot = simspec.phot[channel] + simspec.skyphot[channel]
    else:
        phot = simspec.phot[channel]

    if ispec*nfibers >= simspec.nspec:
        print "ERROR: camera {} not in the {} spectra in {}/{}".format(
            camera, simspec.nspec, night, os.path.basename(simfile))
        return

    #- Load PSF
    psf = desimodel.io.load_psf(channel)

    #- Trim to just the spectra for this spectrograph
    if nspec is None:
        ii = slice(nfibers*ispec, nfibers*(ispec+1))
    else:
        ii = slice(nfibers*ispec, nfibers*ispec + nspec)

    phot = phot[ii]

    #- check if simulation has less than 500 input spectra
    if phot.shape[0] < nspec:
        nspec = phot.shape[0]

    #- Project to image and append that to file
    if verbose:
        print "Projecting photons onto CCD"
        
    img = parallel_project(psf, wave, phot, ncpu=ncpu)
    
    if trimxy:
        xmin, xmax, ymin, ymax = psf.xyrange((0,nspec), wave)
        img = img[0:ymax, 0:xmax]
        # img = img[ymin:ymax, xmin:xmax]
        # hdr['CRVAL1'] = xmin+1
        # hdr['CRVAL2'] = ymin+1

    #- Prepare header
    hdr = simspec.header
    tmp = '/'.join(simfile.split('/')[-3:])  #- last 3 elements of path
    hdr['SIMFILE'] = (tmp, 'Input simulation file')

    #- Strip unnecessary keywords
    for key in ('EXTNAME', 'LOGLAM', 'AIRORVAC', 'CRVAL1', 'CDELT1'):
        if key in hdr:
            del hdr[key]

    #- Write noiseless output
    simpixfile = io.findfile('simpix', night=night, expid=expid, camera=camera)
    io.write_simpix(simpixfile, img, meta=hdr)

    #- Add cosmics from library of dark images
    #- in this case, don't add readnoise since the dark image already has it
    if cosmics is not None:
        cosmics = io.read_cosmics(cosmics, expid, shape=img.shape)
        pix = np.random.poisson(img) + cosmics.pix
        readnoise = cosmics.meta['RDNOISE']
        ivar = 1.0/(pix.clip(0) + readnoise**2)
        mask = cosmics.mask
    #- Or just add noise
    else:
        params = desimodel.io.load_desiparams()
        channel = camera[0].lower()
        readnoise = params['ccd'][channel]['readnoise']
        pix = np.random.poisson(img) + np.random.normal(scale=readnoise, size=img.shape)
        ivar = 1.0/(pix.clip(0) + readnoise**2)
        mask = np.zeros(img.shape, dtype=np.uint16)
    
    #- Metadata to be included in pix file header is in the fibermap header
    #- TODO: this is fragile; consider updating fibermap to use astropy Table
    #- that includes the header rather than directly assuming FITS as the
    #- underlying format.
    fibermapfile = desispec.io.findfile('fibermap', night=night, expid=expid)
    fm, fmhdr = desispec.io.read_fibermap(fibermapfile, header=True)
    meta = dict()
    try:
        meta['TELRA']  = simspec.header['TELRA']
        meta['TELDEC'] = simspec.header['TELDEC']
    except KeyError:  #- temporary backwards compatibilty
        meta['TELRA']  = fmhdr['TELERA']
        meta['TELDEC'] = fmhdr['TELEDEC']
        
    meta['TILEID'] = simspec.header['TILEID']
    meta['DATE-OBS'] = simspec.header['DATE-OBS']
    meta['FLAVOR'] = simspec.header['FLAVOR']
    meta['EXPTIME'] = simspec.header['EXPTIME']
    meta['AIRMASS'] = simspec.header['AIRMASS']

    image = Image(pix, ivar, mask, readnoise=readnoise, camera=camera, meta=meta)
    pixfile = desispec.io.findfile('pix', night=night, camera=camera, expid=expid)
    desispec.io.write_image(pixfile, image)

    if verbose:
        print "Wrote "+pixfile
Exemplo n.º 7
0
def main(args, comm=None):
    if args.verbose:
        import logging
        log.setLevel(logging.DEBUG)
    #we do this so we can use operator.itemgetter
    import operator
    #we do this so our print statements can have timestamps
    import time

    rank = 0
    nproc = 1
    if comm is not None:
        import mpi4py
        rank = comm.rank
        nproc = comm.size
    else:
        if args.ncpu == 0:
            import multiprocessing as mp
            args.ncpu = mp.cpu_count() // 2

    if rank == 0:
        log.info('Starting pixsim at {}'.format(asctime()))

    #- Pre-flight check that these cameras haven't been done yet
    if (rank == 0) and (not args.overwrite) and os.path.exists(args.rawfile):
        log.debug('Checking if cameras are already in output file')
        from astropy.io import fits
        fx = fits.open(args.rawfile)
        oops = False
        for camera in args.cameras:
            if camera.upper() in fx:
                log.error('Camera {} already in {}'.format(
                    camera, args.rawfile))
                oops = True
        fx.close()
        if oops:
            log.fatal('Exiting due to repeat cameras already in output file')
            if comm is not None:
                comm.Abort()
            else:
                sys.exit(1)

    #assign communicators, groups, etc, used for both mpi and multiprocessing
    ncamera = len(args.cameras)
    comm_group = comm
    comm_rank = None
    group = 0
    ngroup = 1
    group_rank = 0
    if comm is not None:
        if args.mpi_camera > 1:
            ngroup = int(comm.size / args.mpi_camera)
            group = int(comm.rank / args.mpi_camera)
            group_rank = comm.rank % args.mpi_camera
            comm_group = comm.Split(color=group, key=group_rank)
            comm_rank = comm.Split(color=group_rank, key=group)
        else:
            group = comm.rank
            ngroup = comm.size
            comm_group = MPI.COMM_SELF
            comm_rank = comm

    #check to make sure that the number of frames is evenly divisible by nubmer of nodes
    #if not, abort and provide a helpful error message
    #otherwise due to barrier logic the program will hang
    if comm is not None:
        if comm.rank == 0:
            if ncamera % ngroup != 0:
                msg = 'Number of frames (ncamera) must be evenly divisible by number of nodes (N)'
                log.error(msg)
                raise ValueError(msg)
                comm.Abort()

    # Remove outputs and or temp files
    rawtemp = "{}.tmp".format(args.rawfile)
    simpixtemp = "{}.tmp".format(args.simpixfile)

    if rank == 0:
        if args.overwrite and os.path.exists(args.rawfile):
            log.debug('removing {}'.format(args.rawfile))
            os.remove(args.rawfile)

        if args.overwrite and os.path.exists(args.simpixfile):
            log.debug('removing {}'.format(args.simpixfile))
            os.remove(args.simpixfile)

        # cleanup stale temp files
        if os.path.isfile(rawtemp):
            os.remove(rawtemp)
        if os.path.isfile(simpixtemp):
            os.remove(simpixtemp)

    if comm is not None:
        comm.barrier()

    #load psf for both mpi and multiprocessing
    psf = None
    if args.psf is not None:
        from specter.psf import load_psf
        psf = load_psf(args.psf)

    # create list of tuples
    camera_channel_args = []
    camera_channel_list = []
    if comm is not None:
        #need to parse these together so we preserve the order
        for item in args.cameras:
            #0th element is camera (b,r,z)
            #1st element is channel (0-9)
            camera_channel_args.append((item[0], item[1]))
            #divide cameras among nodes
            #first sort so zs get distributed first
            #reverse the order so z comes first
            #sort in reverse order so we get z first
            camera_channel_list = np.asarray(
                sorted(camera_channel_args,
                       key=operator.itemgetter(0),
                       reverse=True))
    #also handle multiprocessing case
    if comm is None:
        camera_channel_list = []
        for item in args.cameras:
            camera_channel_list.append((item[0], item[1]))

    #if mpi and N>1, divide cameras between nodes
    if comm is not None:
        if ngroup > 1:
            node_cameras = None
            for i in range(ngroup):
                if i == group:
                    node_cameras = camera_channel_list[i::ngroup]
        if ngroup == 1:
            node_cameras = camera_channel_list
    #also handle multiprocessing case
    if comm is None:
        node_cameras = camera_channel_list

    group_spectro = np.array([int(c[1]) for c in node_cameras], dtype=np.int32)

    #preallocate things needed for multiprocessing (we handle mpi later)
    if comm is None:
        simspec = None
        cosmics = None
        image = {}
        rawpix = {}
        truepix = {}
        lastcamera = None

    # Read the fibermap
    #this is cheap, do here both mpi non mpi
    fibers = None
    if rank == 0:
        if args.fibermap is not None:
            fibermap = desispec.io.read_fibermap(args.fibermap)
            fibers = np.array(fibermap['FIBER'], dtype=np.int32)
        else:
            # Get the fiber list from the simspec file
            from astropy.io import fits
            from astropy.table import Table
            fx = fits.open(args.simspec, memmap=True)
            if 'FIBERMAP' in fx:
                fibermap = Table(fx['FIBERMAP'].data)
                fibers = np.array(fibermap['FIBER'], dtype=np.int32)
            else:
                # Get the number of fibers from one of the photon HDUs
                fibers = np.arange(fx['PHOT_B'].header['NAXIS2'],
                                   dtype=np.int32)
            fx.close()
        if args.nspec > 0:
            fibers = fibers[0:args.nspec]
    if comm is not None:
        fibers = comm.bcast(fibers, root=0)

    #do multiprocessing fibers now, handle mpi inside the loop
    if comm is None:
        fs = np.in1d(fibers // 500, group_spectro)
        group_fibers = fibers[fs]

    # Use original seed to generate different random seeds for each camera
    np.random.seed(args.seed)
    seeds = np.random.randint(0, 2**32 - 1, size=ncamera)
    seed_counter = 0  #need to initialize counter

    #this loop handles the mpi case
    if comm is not None:
        #now that each communicator (group) knows which cameras it
        #is supposed to process, we can get started
        previous_camera = 'a'  #need to initialize
        for i in range(len(node_cameras)):  #may be different in each group
            if i > 1:
                #keep track in case we can avoid re-broadcasting stuff
                previous_camera = node_cameras[i - 1]
            current_camera = node_cameras[i]
            camera = current_camera[0] + current_camera[1]
            channel = current_camera[0]

            #since we clear, we need to pre-allocate every time
            simspec = None
            image = {}
            rawpix = {}
            truepix = {}

            #since we handle only one spectra at a time, just load the one we need
            simspec = io.read_simspec_mpi(args.simspec,
                                          comm_group,
                                          channel,
                                          spectrographs=group_spectro[i])

            if group_rank == 0:
                log.debug('Processing camera {}'.format(camera))

            # Set the seed for this camera (regardless of which process is
            # performing the simulation).
            np.random.seed(seeds[(seed_counter)])

            # Get the random cosmic expids.  The actual values will be
            # remapped internally with the modulus operator.
            cosexpid = np.random.randint(0, 100, size=1)[0]

            # Read inputs for this camera.  Unfortunately psf
            # objects are not serializable, so we read it on all
            # processes.

            if args.psf is None:
                psf = desimodel.io.load_psf(channel)
                if args.ccd_npix_x is not None:
                    psf.npix_x = args.ccd_npix_x
                if args.ccd_npix_y is not None:
                    psf.npix_y = args.ccd_npix_y

            #if we have already loaded the right cosmics camera, don't bcast again
            cosmics = None
            if args.cosmics:
                if current_camera[0] != previous_camera[0]:
                    cosmics = None
                    if group_rank == 0:
                        if args.cosmics_file is None:
                            cosmics_file = io.find_cosmics(
                                camera,
                                simspec.header['EXPTIME'],
                                cosmics_dir=args.cosmics_dir)
                            log.info(
                                'cosmics templates {}'.format(cosmics_file))
                        else:
                            cosmics_file = args.cosmics_file

                        shape = (psf.npix_y, psf.npix_x)
                        cosmics = io.read_cosmics(cosmics_file,
                                                  cosexpid,
                                                  shape=shape)

                    cosmics = comm_group.bcast(cosmics, root=0)

            if group_rank == 0:
                group_size = comm_group.size
                log.info("Group {} ({} processes) simulating camera "
                         "{}".format(group, group_size, camera))

            #calc group_fibers inside the loop
            fs = np.in1d(fibers // 500, group_spectro[i])
            group_fibers = fibers[fs]

            image[camera], rawpix[camera], truepix[camera] = \
                simulate(camera, simspec, psf, fibers=group_fibers,
                    ncpu=args.ncpu, nspec=args.nspec, cosmics=cosmics,
                    wavemin=args.wavemin, wavemax=args.wavemax, preproc=False,
                    comm=comm_group)

            #to avoid overflowing the memory let's write after each instance of simulate
            #the data are already at rank 0 in each communicator
            #then have each communicator open and write to the file one at a time
            #this is the part that will fail if the number of nodes divided by number of frames
            #is not evenly divisible (hence the check above)
            for i in range(ngroup):  #number of total communicators
                if group == i:  #only one group at a time should open/write/close
                    if group_rank == 0:  #write only from rank 0 where the data are
                        #write the raw file using desispec write_raw in raw.py
                        desispec.io.write_raw(rawtemp,
                                              rawpix[camera],
                                              camera=camera,
                                              header=image[camera].meta,
                                              primary_header=simspec.header)
                        log.info('Wrote {} image to {} at time {}'.format(
                            camera, args.rawfile, time.asctime()))
                        #write the simpix file using desisim write_simpix in io.py
                        io.write_simpix(simpixtemp,
                                        truepix[camera],
                                        camera=camera,
                                        meta=simspec.header)
                        log.info('Wrote {} image to {} at time {}'.format(
                            camera, args.simpixfile, time.asctime()))
                        #delete for each group after we have written the output files
                    del rawpix, image, truepix, simspec
                #to avoid bugs this barrier statement must align with if group==i!!
                comm.Barrier()  #all ranks should be done writing

            if args.psf is None:
                del psf

            #iterate random number counter
            seed_counter = seed_counter + 1

    #initalize random number seeds
    seed_counter = 0  #need to initialize counter
    if comm is None:
        simspec = io.read_simspec(args.simspec)
        previous_camera = 'a'
        for i in range(len(node_cameras)):

            current_camera = node_cameras[i]
            camera = current_camera[0] + current_camera[1]
            channel = current_camera[0]

            if group_rank == 0:
                log.debug('Processing camera {}'.format(camera))

            # Set the seed for this camera (regardless of which process is
            # performing the simulation).
            np.random.seed(seeds[(seed_counter)])

            # Get the random cosmic expids.  The actual values will be
            # remapped internally with the modulus operator.
            cosexpid = np.random.randint(0, 100, size=1)[0]

            # Read inputs for this camera.  Unfortunately psf
            # objects are not serializable, so we read it on all
            # processes.
            if args.psf is None:
                #add this so that it won't fail if we enter two repeated channels
                if camera != previous_camera:
                    psf = desimodel.io.load_psf(channel)
                    if args.ccd_npix_x is not None:
                        psf.npix_x = args.ccd_npix_x
                    if args.ccd_npix_y is not None:
                        psf.npix_y = args.ccd_npix_y
                previous_camera = camera

            if args.cosmics:
                if group_rank == 0:
                    if args.cosmics_file is None:
                        cosmics_file = io.find_cosmics(
                            camera,
                            simspec.header['EXPTIME'],
                            cosmics_dir=args.cosmics_dir)
                        log.info('cosmics templates {}'.format(cosmics_file))
                    else:
                        cosmics_file = args.cosmics_file

                    shape = (psf.npix_y, psf.npix_x)
                    cosmics = io.read_cosmics(cosmics_file,
                                              cosexpid,
                                              shape=shape)

            #- Do the actual simulation
            # Each process in the group is responsible for a subset of the
            # fibers.

            if group_rank == 0:
                group_size = 1
                log.info("Group {} ({} processes) simulating camera "
                         "{}".format(group, group_size, camera))

            image[camera], rawpix[camera], truepix[camera] = \
                simulate(camera, simspec, psf, fibers=group_fibers,
                    ncpu=args.ncpu, nspec=args.nspec, cosmics=cosmics,
                    wavemin=args.wavemin, wavemax=args.wavemax, preproc=False,
                    comm=comm_group)

            #like we did in mpi, move the write in here so we don't have to
            #keep accumulating all the data
            if group_rank == 0:
                desispec.io.write_raw(rawtemp,
                                      rawpix[camera],
                                      camera=camera,
                                      header=image[camera].meta,
                                      primary_header=simspec.header)
                log.info('Wrote {} image to {}'.format(camera, args.rawfile))
                io.write_simpix(simpixtemp,
                                truepix[camera],
                                camera=camera,
                                meta=simspec.header)
                log.info('Wrote {} image to {}'.format(camera,
                                                       args.simpixfile))

            if args.psf is None:
                del psf

            #iterate random number counter
            seed_counter = seed_counter + 1

    #done with both mpi and multiprocessing

    # Move temp files into place
    if rank == 0:
        # Copy the original files into place if we are appending
        if not args.overwrite and os.path.exists(args.rawfile):
            shutil.copy2(args.rawfile, rawtemp)
        if not args.overwrite and os.path.exists(args.simpixfile):
            shutil.copy2(args.simpixfile, simpixtemp)

    if rank == 0:
        os.rename(simpixtemp, args.simpixfile)
        os.rename(rawtemp, args.rawfile)
    if comm is not None:
        comm.barrier()

    # Apply preprocessing
    if args.preproc:
        if rank == 0:
            log.info('Preprocessing raw -> pix files')
        from desispec.scripts import preproc
        if len(node_cameras) > 0:
            if group_rank == 0:
                for i in range(len(node_cameras)):
                    current_camera = node_cameras[i]
                    camera = current_camera[0] + current_camera[1]
                    #file should be preproc instead of pix
                    pixfile = desispec.io.findfile('preproc',
                                                   night=args.night,
                                                   expid=args.expid,
                                                   camera=camera)
                    if args.preproc_dir:
                        pixfile = os.path.join(args.preproc_dir,
                                               os.path.basename(pixfile))

                    pixdir = os.path.dirname(pixfile)
                    if not os.path.isdir(pixdir):
                        os.makedirs(pixdir)
                    preproc_opts = [
                        '--infile', args.rawfile, '--outfile', pixfile,
                        '--cameras', camera
                    ]
                    preproc.main(preproc.parse(preproc_opts))

    # Python is terrible with garbage collection, but at least
    # encourage it...
    if comm is None:
        del image
        del rawpix
        del truepix
    #have already deleted these for the mpi case

    if rank == 0:
        log.info('Finished pixsim {} expid {} at {}'.format(
            args.night, args.expid, asctime()))
Exemplo n.º 8
0
def simulate(night,
             expid,
             camera,
             nspec=None,
             verbose=False,
             ncpu=None,
             trimxy=False,
             cosmics=None):
    """
    Run pixel-level simulation of input spectra
    
    Args:
        night (string) : YEARMMDD
        expid (integer) : exposure id
        camera (str) : e.g. b0, r1, z9
        nspec (int) : number of spectra to simulate
        verbose (boolean) : if True, print status messages
        ncpu (int) : number of CPU cores to use in parallel
        trimxy (boolean) : trim image to just pixels with input signal
        cosmics (str) : filename with dark images with cosmics to add

    Reads:
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/simspec-{expid}.fits
        
    Writes:
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/simpix-{camera}-{expid}.fits
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/pix-{camera}-{expid}.fits
    """
    if verbose:
        print "Reading input files"

    channel = camera[0].lower()
    ispec = int(camera[1])
    assert channel in 'brz'
    assert 0 <= ispec < 10

    #- Load DESI parameters
    params = desimodel.io.load_desiparams()
    nfibers = params['spectro']['nfibers']

    #- Load simspec file
    simfile = io.findfile('simspec', night=night, expid=expid)
    simspec = io.read_simspec(simfile)
    wave = simspec.wave[channel]
    if simspec.skyphot is not None:
        phot = simspec.phot[channel] + simspec.skyphot[channel]
    else:
        phot = simspec.phot[channel]

    if ispec * nfibers >= simspec.nspec:
        print "ERROR: camera {} not in the {} spectra in {}/{}".format(
            camera, simspec.nspec, night, os.path.basename(simfile))
        return

    #- Load PSF
    psf = desimodel.io.load_psf(channel)

    #- Trim to just the spectra for this spectrograph
    if nspec is None:
        ii = slice(nfibers * ispec, nfibers * (ispec + 1))
    else:
        ii = slice(nfibers * ispec, nfibers * ispec + nspec)

    phot = phot[ii]

    #- check if simulation has less than 500 input spectra
    if phot.shape[0] < nspec:
        nspec = phot.shape[0]

    #- Project to image and append that to file
    if verbose:
        print "Projecting photons onto CCD"

    img = parallel_project(psf, wave, phot, ncpu=ncpu)

    if trimxy:
        xmin, xmax, ymin, ymax = psf.xyrange((0, nspec), wave)
        img = img[0:ymax, 0:xmax]
        # img = img[ymin:ymax, xmin:xmax]
        # hdr['CRVAL1'] = xmin+1
        # hdr['CRVAL2'] = ymin+1

    #- Prepare header
    hdr = simspec.header
    tmp = '/'.join(simfile.split('/')[-3:])  #- last 3 elements of path
    hdr['SIMFILE'] = (tmp, 'Input simulation file')

    #- Strip unnecessary keywords
    for key in ('EXTNAME', 'LOGLAM', 'AIRORVAC', 'CRVAL1', 'CDELT1'):
        if key in hdr:
            del hdr[key]

    #- Write noiseless output
    simpixfile = io.findfile('simpix', night=night, expid=expid, camera=camera)
    io.write_simpix(simpixfile, img, meta=hdr)

    #- Add cosmics from library of dark images
    #- in this case, don't add readnoise since the dark image already has it
    if cosmics is not None:
        cosmics = io.read_cosmics(cosmics, expid, shape=img.shape)
        pix = np.random.poisson(img) + cosmics.pix
        readnoise = cosmics.meta['RDNOISE']
        ivar = 1.0 / (pix.clip(0) + readnoise**2)
        mask = cosmics.mask
    #- Or just add noise
    else:
        params = desimodel.io.load_desiparams()
        channel = camera[0].lower()
        readnoise = params['ccd'][channel]['readnoise']
        pix = np.random.poisson(img) + np.random.normal(scale=readnoise,
                                                        size=img.shape)
        ivar = 1.0 / (pix.clip(0) + readnoise**2)
        mask = np.zeros(img.shape, dtype=np.uint16)

    #- Metadata to be included in pix file header is in the fibermap header
    #- TODO: this is fragile; consider updating fibermap to use astropy Table
    #- that includes the header rather than directly assuming FITS as the
    #- underlying format.
    fibermapfile = desispec.io.findfile('fibermap', night=night, expid=expid)
    fm, fmhdr = desispec.io.read_fibermap(fibermapfile, header=True)
    meta = dict()
    try:
        meta['TELRA'] = simspec.header['TELRA']
        meta['TELDEC'] = simspec.header['TELDEC']
    except KeyError:  #- temporary backwards compatibilty
        meta['TELRA'] = fmhdr['TELERA']
        meta['TELDEC'] = fmhdr['TELEDEC']

    meta['TILEID'] = simspec.header['TILEID']
    meta['DATE-OBS'] = simspec.header['DATE-OBS']
    meta['FLAVOR'] = simspec.header['FLAVOR']
    meta['EXPTIME'] = simspec.header['EXPTIME']
    meta['AIRMASS'] = simspec.header['AIRMASS']

    image = Image(pix,
                  ivar,
                  mask,
                  readnoise=readnoise,
                  camera=camera,
                  meta=meta)
    pixfile = desispec.io.findfile('pix',
                                   night=night,
                                   camera=camera,
                                   expid=expid)
    desispec.io.write_image(pixfile, image)

    if verbose:
        print "Wrote " + pixfile
Exemplo n.º 9
0
def simulate(night, expid, camera, nspec=None, verbose=False, ncpu=None, trimxy=None):
    """
    Run pixel-level simulation of input spectra
    
    Args:
        night : YEARMMDD string
        expid : integer exposure id
        camera : e.g. b0, r1, z9
        nspec (optional) : number of spectra to simulate
        verbose (optional) : if True, print status messages
        ncpu (optional) : number of CPU cores to use

    Reads:
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/simspec-{expid}.fits
        
    Writes:
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/simpix-{camera}-{expid}.fits
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/pix-{camera}-{expid}.fits
    """
    simdir = io.simdir(night)
    simfile = "{}/simspec-{:08d}.fits".format(simdir, expid)

    if verbose:
        print "Reading input files"

    channel = camera[0].upper()
    ispec = int(camera[1])
    assert channel in "BRZ"
    assert 0 <= ispec < 10

    # - Load DESI parameters
    params = desimodel.io.load_desiparams()
    nfibers = params["spectro"]["nfibers"]

    # - Check that this camera has simulated spectra
    hdr = fits.getheader(simfile, "PHOT_" + channel)
    nspec_in = hdr["NAXIS2"]
    if ispec * nfibers >= nspec_in:
        print "ERROR: camera {} not in the {} spectra in {}/{}".format(
            camera, nspec_in, night, os.path.basename(simfile)
        )
        return

    # - Load input photon data
    phot = fits.getdata(simfile, "PHOT_" + channel)
    try:
        phot += fits.getdata(simfile, "SKYPHOT_" + channel)
    except KeyError:
        pass  # - arcs and flats don't have SKYPHOT

    nwave = phot.shape[1]
    wave = hdr["CRVAL1"] + np.arange(nwave) * hdr["CDELT1"]

    # - Load PSF
    psf = desimodel.io.load_psf(channel)

    # - Trim to just the spectra for this spectrograph
    if nspec is None:
        ii = slice(nfibers * ispec, nfibers * (ispec + 1))
        phot = phot[ii]
    else:
        ii = slice(nfibers * ispec, nfibers * ispec + nspec)
        phot = phot[ii]

    # - check if simulation has less than 500 input spectra
    if phot.shape[0] < nspec:
        nspec = phot.shape[0]

    # - Project to image and append that to file
    if verbose:
        print "Projecting photons onto CCD"

    img = parallel_project(psf, wave, phot, ncpu=ncpu)

    if trimxy:
        xmin, xmax, ymin, ymax = psf.xyrange((0, nspec), wave)
        img = img[0:ymax, 0:xmax]
        # img = img[ymin:ymax, xmin:xmax]
        # hdr['CRVAL1'] = xmin+1
        # hdr['CRVAL2'] = ymin+1

    # - Add noise and write output files
    tmp = "/".join(simfile.split("/")[-3:])  # - last 3 elements of path
    hdr["SIMFILE"] = (tmp, "Input simulation file")
    pixfile = io.write_simpix(img, camera, "science", night, expid, header=hdr)

    if verbose:
        print "Wrote " + pixfile
Exemplo n.º 10
0
def simulate(night, expid, camera, nspec=None, verbose=False, ncpu=None, trimxy=None):
    """
    Run pixel-level simulation of input spectra
    
    Args:
        night : YEARMMDD string
        expid : integer exposure id
        camera : e.g. b0, r1, z9
        nspec (optional) : number of spectra to simulate
        verbose (optional) : if True, print status messages
        ncpu (optional) : number of CPU cores to use

    Reads:
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/simspec-{expid}.fits
        
    Writes:
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/simpix-{camera}-{expid}.fits
        $DESI_SPECTRO_SIM/$PIXPROD/{night}/pix-{camera}-{expid}.fits
    """
    simdir = io.simdir(night)
    simfile = '{}/simspec-{:08d}.fits'.format(simdir, expid)

    if verbose:
        print "Reading input files"

    channel = camera[0].upper()
    ispec = int(camera[1])
    assert channel in 'BRZ'
    assert 0 <= ispec < 10

    #- Load DESI parameters
    params = desimodel.io.load_desiparams()
    nfibers = params['spectro']['nfibers']

    #- Check that this camera has simulated spectra
    fx = fits.open(simfile)
    hdr = fx['PHOT_'+channel].header
    nspec_in = hdr['NAXIS2']
    if ispec*nfibers >= nspec_in:
        print "ERROR: camera {} not in the {} spectra in {}/{}".format(
            camera, nspec_in, night, os.path.basename(simfile))
        return

    #- Load input photon data
    phot = fx['PHOT_'+channel].data
    wave = fx['WAVE_'+channel].data
    if 'SKYPHOT_'+channel in fx:
        phot += fx['SKYPHOT_'+channel].data

    #- Load PSF
    psf = desimodel.io.load_psf(channel)

    #- Trim to just the spectra for this spectrograph
    if nspec is None:
        ii = slice(nfibers*ispec, nfibers*(ispec+1))
        phot = phot[ii]
    else:
        ii = slice(nfibers*ispec, nfibers*ispec + nspec)
        phot = phot[ii]

    #- check if simulation has less than 500 input spectra
    if phot.shape[0] < nspec:
        nspec = phot.shape[0]

    #- Project to image and append that to file
    if verbose:
        print "Projecting photons onto CCD"
        
    img = parallel_project(psf, wave, phot, ncpu=ncpu)
    
    if trimxy:
        xmin, xmax, ymin, ymax = psf.xyrange((0,nspec), wave)
        img = img[0:ymax, 0:xmax]
        # img = img[ymin:ymax, xmin:xmax]
        # hdr['CRVAL1'] = xmin+1
        # hdr['CRVAL2'] = ymin+1

    #- Prepare header
    hdr = fx[0].header
    tmp = '/'.join(simfile.split('/')[-3:])  #- last 3 elements of path
    hdr['SIMFILE'] = (tmp, 'Input simulation file')

    #- Strip unnecessary keywords
    for key in ('EXTNAME', 'LOGLAM', 'AIRORVAC', 'CRVAL1', 'CDELT1'):
        if key in hdr:
            del hdr[key]

    #- Add noise and write output files
    pixfile = io.write_simpix(img, camera, night, expid, header=hdr)

    fx.close()

    if verbose:
        print "Wrote "+pixfile