first commit (f751d1a2) · Commits · Fang Yuedong / injection_pipeline

Grid.py

0 → 100644

+187 −0

Original line number	Diff line number	Diff line
		import math
		import numpy as np
		import matplotlib.pyplot as plt

		class BaseGrid(object):
		_valid_grid_types = ['RectGrid', 'HexGrid']
		_valid_mixed_types = ['MixedGrid']

		class Grid(BaseGrid):
		def __init__(self, grid_spacing, wcs, Npix_x=10000, Npix_y=10000, pixelscale=0.074, rot_angle=None, pos_offset=None, angle_unit='rad'):
		self.grid_spacing = grid_spacing
		self.im_gs = grid_spacing * (1.0 / pixelscale) # pixels
		self.pixelscale = pixelscale
		self.Npix_x, self.Npix_y = Npix_x, Npix_y
		self.wcs = wcs
		self.rot_angle = rot_angle # rotation angle, in rad
		self.angle_unit = angle_unit
		if pos_offset:
		self.pos_offset = np.array(pos_offset)
		else:
		self.pos_offset = np.array([0., 0.])

		# May have to modify grid corners if there is a rotation
		if rot_angle:
		dx = Npix_x / 2.
		dy = Npix_y / 2.
		if angle_unit == 'deg':
		theta = np.deg2rad(rot_angle)
		else:
		theta = rot_angle
		self.startx = (0.-dx) * np.cos(theta) - (Npix_y-dy) * np.sin(theta) + dx
		self.endx = (Npix_x-dx) * np.cos(theta) - (0.-dy) * np.sin(theta) + dx
		self.starty = (0.-dx) * np.cos(theta) + (0.-dy) * np.sin(theta) + dx
		self.endy = (Npix_x-dx) * np.cos(theta) + (Npix_y-dy) * np.sin(theta) + dx
		else:
		self.startx, self.endx= 0., Npix_x
		self.starty, self.endy= 0., Npix_y

		def rotate_grid(self, theta, offset=None, angle_unit='rad'):

		if angle_unit == 'deg':
		theta = np.deg2rad(theta)
		elif angle_unit != 'rad':
		raise ValueError('`angle_unit` can only be `deg` or `rad`! ' +
		'Passed unit of {}'.format(angle_unit))

		if not offset: offset = [0., 0.]

		c, s = np.cos(theta), np.sin(theta)
		R = np.array(((c,-s), (s, c)))

		offset_grid = np.array([self.im_ra - offset[0], self.im_dec - offset[1]])
		translate = np.empty_like(offset_grid)
		translate[0,:] = offset[0]
		translate[1,:] = offset[1]

		rotated_grid = np.dot(R, offset_grid) + translate

		self.im_pos = rotated_grid.T
		self.im_ra, self.im_dec = self.im_pos[0,:], self.im_pos[1,:]

		def cut2buffer(self):
		'''
		Remove objects outside of tile (and buffer).
		We must sample points in the buffer zone in the beginning due to
		possible rotations.
		'''
		b = self.im_gs
		in_region = np.where( (self.im_pos[:,0]>b) & (self.im_pos[:,0]<self.Npix_x-b) &
		(self.im_pos[:,1]>b) & (self.im_pos[:,1]<self.Npix_y-b) )
		self.im_pos = self.im_pos[in_region]
		self.im_ra = self.im_pos[:,0]
		self.im_dec = self.im_pos[:,1]

		# Get all image coordinate pairs
		self.pos = self.wcs.wcs_pix2world(self.im_pos, 1)
		self.ra = self.pos[:,0]
		self.dec = self.pos[:,1]

		class RectGrid(Grid):
		def __init__(self, grid_spacing, wcs, Npix_x=10000, Npix_y=10000, pixelscale=0.074,
		rot_angle=None, pos_offset=None, angle_unit='rad'):
		super(RectGrid, self).__init__(grid_spacing, wcs, Npix_x=Npix_x, Npix_y=Npix_y,
		pixelscale=pixelscale, rot_angle=rot_angle,
		pos_offset=pos_offset, angle_unit=angle_unit)
		self._create_grid()

		def _create_grid(self):
		im_gs = self.im_gs

		po = self.pos_offset
		im_po = po / self.pixelscale
		self.im_ra = np.arange(self.startx, self.endx, im_gs)
		self.im_dec = np.arange(self.starty, self.endy, im_gs)

		# Get all image coordinate pairs
		self.im_pos = np.array(np.meshgrid(self.im_ra, self.im_dec)).T.reshape(
		-1, 2)
		self.im_ra = self.im_pos[:,0]
		self.im_dec = self.im_pos[:,1]

		if self.rot_angle:
		self.rotate_grid(self.rot_angle, angle_unit=self.angle_unit,
		offset=[(self.Npix_x+im_po[0])/2., (self.Npix_y+im_po[1])/2.])

		self.cut2buffer()

		class HexGrid(Grid):
		def __init__(self, grid_spacing, wcs, Npix_x=10000, Npix_y=10000, pixelscale=0.074,
		rot_angle=None, pos_offset=None, angle_unit='rad'):
		super(HexGrid, self).__init__(grid_spacing, wcs, Npix_x=Npix_x, Npix_y=Npix_y,
		pixelscale=pixelscale, rot_angle=rot_angle,
		pos_offset=pos_offset, angle_unit=angle_unit)
		self._create_grid()

		def _create_grid(self):
		im_gs = self.im_gs

		po = self.pos_offset
		im_po = [p / self.pixelscale for p in po]
		self.im_pos = HexGrid.calc_hex_coords(self.startx, self.starty, self.endx, self.endy, im_gs)
		self.im_ra = self.im_pos[:, 0]
		self.im_dec = self.im_pos[:, 1]

		if self.rot_angle:
		self.rotate_grid(self.rot_angle, angle_unit=self.angle_unit,
		offset=[(self.Npix_x+im_po[0])/2., (self.Npix_y+im_po[1])/2.])
		self.cut2buffer()

		@classmethod
		def calc_hex_coords(cls, startx, starty, endx, endy, radius):
		# Geoemtric factors of given hexagon
		r = radius
		p = r * np.tan(np.pi / 6.) # side length / 2
		h = 4. * p
		dx = 2. * r
		dy = 2. * p

		row = 1

		xs = []
		ys = []

		while startx < endx:
		x = [startx, startx, startx + r, startx + dx, startx + dx, startx + r, startx + r]
		xs.append(x)
		startx += dx

		while starty < endy:
		y = [starty + p, starty + 3p, starty + h, starty + 3p, starty + p, starty, starty + dy]
		ys.append(y)
		starty += 2*p
		row += 1

		print(xs)
		print(ys)

		polygons = [zip(x, y) for x in xs for y in ys]
		polygons = [np.column_stack((x, y)) for x in xs for y in ys]
		# polygons = np.array(polygons)
		hexgrid = cls.polygons2coords(polygons)

		# Some hexagonal elements go beyond boundary; cut these out
		indx = np.where( (hexgrid[:,0]<endx) & (hexgrid[:,1]<endy) )
		return hexgrid[indx]

		@classmethod
		def polygons2coords(HexGrid, p):
		print(p)
		s = np.shape(p)
		print(s)
		L = s[0]*s[1]
		pp = np.array(p).reshape(L,2)
		c = np.vstack({tuple(row) for row in pp})
		# Some of the redundant coordinates are offset by ~1e-10 pixels
		return np.unique(c.round(decimals=6), axis=0)

		def _build_grid(grid_type, **kwargs):
		if grid_type in GRID_TYPES:
		return GRID_TYPES[grid_type](**kwargs)
		else:
		raise ValueError('There is not yet an implemnted default Grid of type {}'.format(grid_type))

		GRID_TYPES = {
		'RectGrid': RectGrid,
		'HexGrid': HexGrid
		}
		No newline at end of file

InjectionCatalog.py

0 → 100644

+124 −0

Original line number	Diff line number	Diff line
		import numpy as np
		import galsim

		import Grid
		import mathutil as util

		class InjectionCatalog(object):

		def __init__(self, image):
		self.rmin, self.rmax = image.ramin, image.ramax
		self.decmin, self.decmax = image.decmin, image.decmax
		self.ra_boundary_cross = image.ra_boundary_cross
		self.Npix_x, self.Npix_y = image.Npix_x, image.Npix_y
		self.pixel_scale = image.pixel_scale
		self.wcs = image.wcs

		# This is also set for a given image, as it may depend on its area
		self.objs_per_real = image.objs_per_real

		# TODO: output truth catalog after injection
		self.truth_outfile = {}
		pass

		def generate_positions(self, config):

		bg = Grid.BaseGrid()

		self.nobjects = self.objs_per_real
		ps = config['pos_sampling']
		pstype = ps['type']

		if pstype == 'uniform' and (self.objs_per_real is not None):
		ra = util.sample_uniform_ra(self.ramin, self.ramax, self.objs_per_real,
		boundary_cross=self.ra_boundary_cross)
		dec = util.sample_uniform_dec(self.decmin, self.decmax, self.objs_per_real, unit='deg')
		self.pos = np.column_stack((ra, dec))

		elif pstype in bg._valid_grid_types:
		grid_kwargs = self._build_grid_kwargs(pstype, ps)
		gtype = pstype

		img_grid = Grid._build_grid(gtype, **grid_kwargs)
		self.pos = img_grid.pos
		self.nobjects = np.shape(img_grid.pos)[0]

		elif pstype in bg._valid_mixed_types:
		pass

		else:
		raise ValueError('Position sampling type {} is not valid!'.format(gtype))

		# Generate object indices (in input catalog)
		# TODO
		# self.index = np.random.choice(xrange(input_nobjects), size=self.nobjects)

		# Generate object rotation angles, if desired
		if config['rotate_objs'] is True:
		rot = util.sample_uniform(0., 360., self.nobjects)
		self.rotate = np.array([str(r) + ' deg' for r in rot])
		else:
		self.rotate = None


		def _build_grid_kwargs(self, pstype, ps):
		if pstype == 'MixedGrid':
		gtype = ps['grid_type']
		else:
		gtype = pstype

		gs = ps['grid_spacing']

		try:
		r = ps['rotate']
		if (isinstance(r, str) and (r.lower() == 'random')):
		if gtype == 'RectGrid':
		self.grid_rot_angle = np.random.uniform(0., np.pi/2.)
		elif gtype == 'HexGrid':
		self.grid_rot_angle = np.random.uniform(0., np.pi/3.)
		else:
		unit = ps['angle_unit']
		if unit == 'deg':
		if (r >= 0.) and (r < 360.):
		self.grid_rot_angle = float(r)
		else:
		raise ValueError('Grid rotation of {} '.format(r) + 'deg is not valid!')
		else:
		if (r >= 0.) and (r < 2*np.pi):
		self.grid_rot_angle = float(r)
		else:
		raise ValueError('Grid rotation of {} '.format(r) + 'rad is not valid!')
		except KeyError:
		self.grid_rot_angle = 0.0

		# Offset grid if desired
		try:
		offset = ps['offset']
		if (isinstance(offset, str)) and (offset.lower() == 'random'):
		self.grid_offset = [np.random.uniform(-gs/2., gs/2.),
		np.random.uniform(-gs/2., gs/2.)]
		else:
		if isinstance(offset, list):
		self.grid_offset = list(offset)
		else:
		raise ValueError('Grid offset of {} '.format(offset) + 'is not valid!')
		except KeyError:
		self.grid_offset = [0., 0.]

		try:
		self.angle_unit = ps['angle_unit']
		except KeyError:
		self.angle_unit = 'rad'

		grid_kwargs = dict(grid_spacing=gs,
		wcs=self.wcs,
		Npix_x=self.Npix_x,
		Npix_y=self.Npix_y,
		pixelscale=self.pixel_scale,
		rot_angle=self.grid_rot_angle,
		angle_unit=self.angle_unit,
		pos_offset=self.grid_offset)
		return grid_kwargs

		def get_truth_outfile(self):
		pass

Injector.py

0 → 100644

+90 −0

Original line number	Diff line number	Diff line
		import galsim
		import galsim.config.stamp as stamp
		import logging
		import os
		import numpy as np

		import grid

		class AddOnImageBuilder(galsim.config.image_scattered.ScatteredImageBuilder):

		def setup(self, config, base, image_num, obj_num, ignore, logger):
		ignore = ignore + ['initial_image']
		return super(AddOnImageBuilder, self).setup(config, base, image_num, obj_num, ignore, logger)

		def addNoise(self, image, config, base, image_num, obj_num, current_var, logger):
		super(AddOnImageBuilder, self).addNoise(image, config, base, image_num, obj_num, current_var, logger)

		initial_image_name = galsim.config.ParseValue(config, 'initial_image', base, str)[0]
		initial_image = galsim.fits.read(initial_image_name)
		image += initial_image


		galsim.config.RegisterImageType('AddOnImage', AddOnImageBuilder())

		class InjectImageBuilder(AddOnImageBuilder):

		def setup(self, config, base, image_num, obj_num, ignore, logger):
		extra_ignore = ignore + ['tile_list', 'geom_file', 'tile_dir', 'config_dir', 'psf_dir',
		'version', 'run_name', 'bands', 'n_objects', 'n_realizations',
		'object_density', 'inj_objs_only', 'pos_sampling', 'realizations',
		'extinct_objs', 'rotate_objs']
		for key in config:
		if 'N_' in key:
		extra_ignore.append(key)
		full_xsize, full_ysize = super(InjectImageBuilder, self).setup(config, base, image_num, obj_num, extra_ignore, logger)

		# config = parse_inject_image_inputs(config, base)

		return full_xsize, full_ysize

		def addNoise(self, image, config, base, image_num, obj_num, current_var, logger):
		try:
		ioo = config['inj_objs_only']
		if (type(ioo) is bool) and (ioo is True):
		return super(AddOnImageBuilder, self).addNoise(image,
		config,
		base,
		image_num,
		obj_num,
		current_var,
		logger)
		elif (isinstance(ioo, dict)) and (ioo['value'] is True):
		# Still want to use existing image if changed to be BKG
		if (ioo['noise']) and ('BKG' in ioo['noise']):
		return super(InjectImageBuilder, self).addNoise(image,
		config,
		base,
		image_num,
		obj_num,
		current_var,
		logger)
		else:
		return super(AddOnImageBuilder, self).addNoise(image,
		config,
		base,
		image_num,
		obj_num,
		current_var,
		logger)
		else:
		# Default is to add on top of initial images
		return super(InjectImageBuilder, self).addNoise(image,
		config,
		base,
		image_num,
		obj_num,
		current_var,
		logger)

		except KeyError:
		# Default is to add on top of initial images
		return super(InjectImageBuilder, self).addNoise(image,
		config,
		base,
		image_num,
		obj_num,
		current_var,
		logger)

		galsim.config.RegisterImageType('InjectImage', InjectImageBuilder())

InputCatalogs.py

0 → 100644

+162 −0

Original line number	Diff line number	Diff line
		import os
		import numpy as np
		import h5py as h5
		import random
		import galsim
		import astropy.constants as cons
		from astropy.table import Table
		from scipy import interpolate

		from ObservationSim.MockObject import CatalogBase, Star, Galaxy, Quasar
		from ObservationSim.MockObject._util import seds, sed_assign, extAv, tag_sed, getObservedSED

		try:
		import importlib.resources as pkg_resources
		except ImportError:
		# Try backported to PY<37 'importlib_resources'
		import importlib_resources as pkg_resources

		class SimCat(CatalogBase):
		def __init__(self, config, chip, nobjects=None):
		super().__init__()
		self.cat_dir = os.path.join(config["data_dir"], config["input_path"]["cat_dir"])
		self.config = config
		self.chip = chip
		self.seed_Av = config["random_seeds"]["seed_Av"]

		with pkg_resources.path('Catalog.data', 'SLOAN_SDSS.g.fits') as filter_path:
		self.normF_star = Table.read(str(filter_path))
		with pkg_resources.path('Catalog.data', 'lsst_throuput_g.fits') as filter_path:
		self.normF_galaxy = Table.read(str(filter_path))
		if "star_cat" in config["input_path"] and config["input_path"]["star_cat"]:
		star_file = config["input_path"]["star_cat"]
		star_SED_file = config["SED_templates_path"]["star_SED"]
		self.star_path = os.path.join(self.cat_dir, star_file)
		self.star_SED_path = os.path.join(config["data_dir"], star_SED_file)
		self._load_SED_lib_star()
		if "galaxy_cat" in config["input_path"] and config["input_path"]["galaxy_cat"]:
		galaxy_file = config["input_path"]["galaxy_cat"]
		self.galaxy_path = os.path.join(self.cat_dir, galaxy_file)
		self.galaxy_SED_path = os.path.join(config["data_dir"], config["SED_templates_path"]["galaxy_SED"])
		self._load_SED_lib_gals()

		self._load(nobjects=nobjects)

		def _load_SED_lib_star(self):
		self.tempSED_star = h5.File(self.star_SED_path,'r')

		def _load_SED_lib_gals(self):
		self.tempSed_gal, self.tempRed_gal = seds("galaxy.list", seddir=self.galaxy_SED_path)

		def load_norm_filt(self, obj):
		if obj.type == "star":
		return self.normF_star
		elif obj.type == "galaxy" or obj.type == "quasar":
		return self.normF_galaxy
		else:
		return None

		def _load_gals(self, gals, pix_id=None, nobjects=None):
		# Load how mnay objects?
		if nobjects is None:
		ngals = 5000
		else:
		ngals = nobjects
		self.rng_sedGal = random.Random()
		self.rng_sedGal.seed(pix_id) # Use healpix index as the random seed
		self.ud = galsim.UniformDeviate(pix_id)

		for igals in range(ngals):
		param = self.initialize_param()
		param['ra'] = gals['ra_true'][igals]
		param['dec'] = gals['dec_true'][igals]

		# param['mag_use_normal'] = gals['mag_true_g_lsst'][igals]
		# (TEST) use same magnitude
		# (there will be slight difference due to randomness in SED)
		param['mag_use_normal'] = 18

		param['z'] = gals['redshift_true'][igals]
		param['model_tag'] = 'None'
		param['gamma1'] = 0
		param['gamma2'] = 0
		param['kappa'] = 0
		param['delta_ra'] = 0
		param['delta_dec'] = 0

		hlrMajB = gals['size_bulge_true'][igals]
		hlrMinB = gals['size_minor_bulge_true'][igals]

		hlrMajD = gals['size_disk_true'][igals]
		hlrMinD = gals['size_minor_disk_true'][igals]
		aGal = gals['size_true'][igals]
		bGal = gals['size_minor_true'][igals]
		param['bfrac'] = gals['bulge_to_total_ratio_i'][igals]
		param['theta'] = gals['position_angle_true'][igals]
		param['hlr_bulge'] = np.sqrt(hlrMajB * hlrMinB)
		param['hlr_disk'] = np.sqrt(hlrMajD * hlrMinD)
		param['ell_bulge'] = (hlrMajB - hlrMinB)/(hlrMajB + hlrMinB)
		param['ell_disk'] = (hlrMajD - hlrMinD)/(hlrMajD + hlrMinD)
		param['ell_tot'] = (aGal - bGal) / (aGal + bGal)

		# Assign each galaxy a template SED
		param['sed_type'] = sed_assign(phz=param['z'], btt=param['bfrac'], rng=self.rng_sedGal)
		param['redden'] = self.tempRed_gal[param['sed_type']]
		param['av'] = self.avGal[int(self.ud()*self.nav)]
		if param['sed_type'] <= 5:
		param['av'] = 0.0
		param['redden'] = 0
		param['star'] = 0 # Galaxy
		if param['sed_type'] >= 29:
		param['av'] = 0.6 * param['av'] / 3.0 # for quasar, av=[0, 0.2], 3.0=av.max-av.im
		param['star'] = 2 # Quasar

		param['id'] = gals['galaxyID'][igals]

		if param['star'] == 0:
		obj = Galaxy(param)
		if param['star'] == 2:
		obj = Quasar(param)

		self.objs.append(obj)

		def _load(self, nobjects=None):
		# (TEST) use objects in healpix:
		pix = 48656
		self.nav = 15005
		self.avGal = extAv(self.nav, seed=self.seed_Av)
		self.objs = []

		gals_cat = h5.File(self.galaxy_path, 'r')['galaxies']
		gals = gals_cat[str(pix)]
		self._load_gals(gals, pix_id=pix, nobjects=nobjects)
		del gals

		def load_sed(self, obj, **kwargs):
		if obj.type == 'star':
		_, wave, flux = tag_sed(
		h5file=self.tempSED_star,
		model_tag=obj.param['model_tag'],
		teff=obj.param['teff'],
		logg=obj.param['logg'],
		feh=obj.param['feh']
		)
		elif obj.type == 'galaxy' or obj.type == 'quasar':
		sed_data = getObservedSED(
		sedCat=self.tempSed_gal[obj.sed_type],
		redshift=obj.z,
		av=obj.param["av"],
		redden=obj.param["redden"]
		)
		wave, flux = sed_data[0], sed_data[1]
		else:
		raise ValueError("Object type not known")
		speci = interpolate.interp1d(wave, flux)
		lamb = np.arange(2000, 18001 + 0.5, 0.5)
		y = speci(lamb)
		# erg/s/cm2/A --> photo/s/m2/A
		all_sed = y * lamb / (cons.h.value * cons.c.value) * 1e-13
		sed = Table(np.array([lamb, all_sed]).T, names=('WAVELENGTH', 'FLUX'))
		del wave
		del flux
		return sed
		No newline at end of file

README.md

+1 −92

File changed.

Preview size limit exceeded, changes collapsed.