add gc collect (1db4fa7e) · Commits · Liu Dezi / csst_msc_sim

Catalog/C6_Catalog.py

+129 −26

Original line number	Diff line number	Diff line
		@@ -76,6 +76,14 @@ class C6_Catalog(CatalogBase):
		self.galaxy_SED_path = os.path.join(config["data_dir"], config["SED_templates_path"]["galaxy_SED"])
		self._load_SED_lib_gals()

		if "AGN_cat" in config["input_path"] and config["input_path"]["AGN_cat"] and not config["run_option"]["star_only"]:
		AGN_dir = config["input_path"]["AGN_cat"]
		self.AGN_path = os.path.join(config["data_dir"], config["input_path"]["AGN_cat"])
		self.AGN_SED_path = os.path.join(config["data_dir"], config["SED_templates_path"]["AGN_SED"])
		self.AGN_SED_wave_path = os.path.join(config['data_dir'], config["SED_templates_path"]["AGN_SED_WAVE"])
		self._load_SED_lib_AGN()


		if "rotateEll" in config["shear_setting"]:
		self.rotation = float(int(config["shear_setting"]["rotateEll"]/45.))
		else:
		@@ -131,6 +139,11 @@ class C6_Catalog(CatalogBase):
		self.lamb_gal = lamb['lamb'][()]
		self.pcs = pcs['pcs'][()]

		def _load_SED_lib_AGN(self):
		from astropy.io import fits
		self.SED_AGN = fits.open(self.AGN_SED_path)[0].data
		self.lamb_AGN = np.load(self.AGN_SED_wave_path)


		def _load_gals(self, gals, pix_id=None, cat_id=0):
		ngals = len(gals['ra'])
		@@ -172,7 +185,7 @@ class C6_Catalog(CatalogBase):
		# # (TEST)
		# if igals > 100:
		# break
		# if igals != 1186:
		# if igals < 3447:
		# continue

		param = self.initialize_param()
		@@ -250,8 +263,6 @@ class C6_Catalog(CatalogBase):

		if param['star'] == 0:
		obj = Galaxy(param, self.rotation, logger=self.logger)
		if param['star'] == 2:
		obj = Quasar(param, logger=self.logger)

		# Need to deal with additional output columns
		obj.additional_output_str = self.add_fmt%("n", 0., 0., 0.,
		@@ -334,6 +345,69 @@ class C6_Catalog(CatalogBase):

		self.objs.append(obj)

		def _load_AGNs(self):
		data = Table.read(self.AGN_path)
		ra_arr = data['ra']
		dec_arr = data['dec']
		nAGNs = len(data)
		if self.config["obs_setting"]["enable_astrometric_model"]:
		ra_list = ra_arr.tolist()
		dec_list = dec_arr.tolist()
		pmra_list = np.zeros(nAGNs).tolist()
		pmdec_list = np.zeros(nAGNs).tolist()
		rv_list = np.zeros(nAGNs).tolist()
		parallax_list = [1e-9] * nAGNs
		dt = datetime.utcfromtimestamp(self.pointing.timestamp)
		date_str = dt.date().isoformat()
		time_str = dt.time().isoformat()
		ra_arr, dec_arr = on_orbit_obs_position(
		input_ra_list=ra_list,
		input_dec_list=dec_list,
		input_pmra_list=pmra_list,
		input_pmdec_list=pmdec_list,
		input_rv_list=rv_list,
		input_parallax_list=parallax_list,
		input_nstars=nAGNs,
		input_x=self.pointing.sat_x,
		input_y=self.pointing.sat_y,
		input_z=self.pointing.sat_z,
		input_vx=self.pointing.sat_vx,
		input_vy=self.pointing.sat_vy,
		input_vz=self.pointing.sat_vz,
		input_epoch="J2000",
		input_date_str=date_str,
		input_time_str=time_str
		)
		for iAGNs in range(nAGNs):
		param = self.initialize_param()
		param['ra'] = ra_arr[iAGNs]
		param['dec'] = dec_arr[iAGNs]
		param['ra_orig'] = data['ra'][iAGNs]
		param['dec_orig'] = data['dec'][iAGNs]
		param['z'] = data['z'][iAGNs]
		param['appMag'] = data['appMag'][iAGNs]
		param['absMag'] = data['absMag'][iAGNs]

		# NOTE: this cut cannot be put before the SED type has been assigned
		if not self.chip.isContainObj(ra_obj=param['ra'], dec_obj=param['dec'], margin=200):
		continue

		# TEST no redening and no extinction
		param['av'] = 0.0
		param['redden'] = 0

		param['star'] = 2 # Quasar
		param['id'] = data['igmlos'][iAGNs]

		if param['star'] == 2:
		obj = Quasar(param, logger=self.logger)

		# Append additional output columns to the .cat file
		obj.additional_output_str = self.add_fmt%("n", 0., 0., 0.,
		0., 0., 0., 0., 0., 0., 0., 0., 0., -1, 0.)
		self.objs.append(obj)


		def _load(self, **kwargs):
		self.objs = []
		self.ids = 0
		@@ -365,6 +439,9 @@ class C6_Catalog(CatalogBase):
		for pix in self.pix_list:
		try:
		bundleID = get_bundleIndex(pix)
		# # (TEST C6):
		# if pix != 35421 or bundleID != 523:
		# continue
		file_path = os.path.join(self.galaxy_path, "galaxies_C6_bundle{:06}.h5".format(bundleID))
		gals_cat = h5.File(file_path, 'r')['galaxies']
		gals = gals_cat[str(pix)]
		@@ -374,6 +451,14 @@ class C6_Catalog(CatalogBase):
		traceback.print_exc()
		self.logger.error(str(e))
		print(e)
		if "AGN_cat" in self.config["input_path"] and self.config["input_path"]["AGN_cat"] and not self.config["run_option"]["star_only"]:
		try:
		self._load_AGNs()
		except Exception as e:
		traceback.print_exc()
		self.logger.error(str(e))
		print(e)

		if self.logger is not None:
		self.logger.info("maximum galaxy size: %.4f"%(self.max_size))
		self.logger.info("number of objects in catalog: %d"%(len(self.objs)))
		@@ -407,6 +492,7 @@ class C6_Catalog(CatalogBase):
		# dist_L_pc = (1 + obj.z) * comoving_dist(z=obj.z)[0]
		# factor = (10 / dist_L_pc)**2
		factor = 10*(-.4 self.cosmo.distmod(obj.z).value)
		if obj.type == 'galaxy':
		flux = np.matmul(self.pcs, obj.coeff) * factor
		# if np.any(flux < 0):
		# raise ValueError("Glaxy %s: negative SED fluxes"%obj.id)
		@@ -419,6 +505,14 @@ class C6_Catalog(CatalogBase):
		redden=obj.param["redden"]
		)
		wave, flux = sed_data[0], sed_data[1]
		elif obj.type == 'quasar':
		# flux = self.SED_AGN[int(obj.id)] * factor * 1e-17
		flux = self.SED_AGN[int(obj.id)] * 1e-17
		# if np.any(flux < 0):
		# raise ValueError("Glaxy %s: negative SED fluxes"%obj.id)
		flux[flux < 0] = 0.
		# sedcat = np.vstack((self.lamb_AGN, flux)).T
		wave = self.lamb_AGN
		# print("sed (erg/s/cm2/A) = ", sed_data)
		# np.savetxt(os.path.join(self.config["work_dir"], "%s_sed.txt"%(obj.id)), sedcat)
		else:
		@@ -430,17 +524,26 @@ class C6_Catalog(CatalogBase):
		all_sed = y * lamb / (cons.h.value * cons.c.value) * 1e-13
		sed = Table(np.array([lamb, all_sed]).T, names=('WAVELENGTH', 'FLUX'))

		# if obj.type == 'galaxy' or obj.type == 'quasar':
		# # integrate to get the magnitudes
		# sed_photon = np.array([sed['WAVELENGTH'], sed['FLUX']]).T
		# sed_photon = galsim.LookupTable(x=np.array(sed_photon[:, 0]), f=np.array(sed_photon[:, 1]), interpolant='nearest')
		# sed_photon = galsim.SED(sed_photon, wave_type='A', flux_type='1', fast=False)
		# interFlux = integrate_sed_bandpass(sed=sed_photon, bandpass=self.filt.bandpass_full)
		# obj.param['mag_use_normal'] = getABMAG(interFlux, self.filt.bandpass_full)
		# # print("mag_use_normal = %.3f"%obj.param['mag_use_normal'])
		# # mag = getABMAG(interFlux, self.filt.bandpass_full)
		# # print("mag diff = %.3f"%(mag - obj.param['mag_use_normal']))

		if obj.type == 'quasar':
		# integrate to get the magnitudes
		sed_photon = np.array([sed['WAVELENGTH'], sed['FLUX']]).T
		sed_photon = galsim.LookupTable(x=np.array(sed_photon[:, 0]), f=np.array(sed_photon[:, 1]), interpolant='nearest')
		sed_photon = galsim.SED(sed_photon, wave_type='A', flux_type='1', fast=False)
		interFlux = integrate_sed_bandpass(sed=sed_photon, bandpass=self.filt.bandpass_full)
		obj.param['mag_use_normal'] = getABMAG(interFlux, self.filt.bandpass_full)
		# if obj.param['mag_use_normal'] >= 30:
		# print("obj ID = %d"%obj.id)
		# print("mag_use_normal = %.3f"%obj.param['mag_use_normal'])
		# print("integrated flux = %.7f"%(interFlux))
		# print("app mag = %.3f"%obj.param['appMag'])
		# np.savetxt('./AGN_SED_test/sed_objID_%d.txt'%obj.id, np.transpose([self.lamb_AGN, self.SED_AGN[int(obj.id)]]))
		# print("obj ID = %d"%obj.id)
		# print("mag_use_normal = %.3f"%obj.param['mag_use_normal'])
		# print("integrated flux = %.7f"%(interFlux))
		# print("app mag = %.3f"%obj.param['appMag'])
		# print("abs mag = %.3f"%obj.param['absMag'])
		# mag = getABMAG(interFlux, self.filt.bandpass_full)
		# print("mag diff = %.3f"%(mag - obj.param['mag_use_normal']))
		del wave
		del flux
		return sed

ObservationSim/ObservationSim.py

+10 −0

Original line number	Diff line number	Diff line
		@@ -4,6 +4,7 @@ import mpi4py.MPI as MPI
		import galsim
		import logging
		import psutil
		import gc
		from astropy.io import fits
		from datetime import datetime

		@@ -334,12 +335,20 @@ class Observation(object):
		traceback.print_exc()
		chip_output.logger.error(e)
		pass

		# [C6 TEST]
		# print("check running:1: pointing-{:} chip-{:} pid-{:} memory-{:6.2}GB".format(pointing.id, chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) ), flush=True)
		# print('draw object %s'%obj.id)
		# print('mag = %.3f'%obj.param['mag_use_normal'])

		# Unload SED:
		obj.unload_SED()
		del obj
		gc.collect()

		del psf_model
		del self.cat
		gc.collect()

		# print("check running:1: pointing-{:} chip-{:} pid-{:} memory-{:6.2}GB".format(pointing.id, chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) ), flush=True)
		chip_output.logger.info("check running:1: pointing-%d chip-%d pid-%d memory-%6.2fGB"%(pointing.id, chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) ))
		@@ -474,3 +483,4 @@ class Observation(object):
		chip_output.logger.info("finished running chip#%d..."%(chip.chipID))
		for handler in chip_output.logger.handlers[:]:
		chip_output.logger.removeHandler(handler)
		gc.collect()