try optimize RAM usage of drawing methods (0877a23d) · Commits · csst-sims / csst_msc_sim

docs/object_rendering_memory.md

0 → 100644

+88 −0

Original line number	Diff line number	Diff line
		# Object-rendering memory safeguards

		Object draws in `Galaxy`, `MockObject` (including stars), and `Stamp` use
		`observation_sim.rendering`. Before allocating a stamp, the renderer estimates
		its dimensions and the Fourier grid, including local WCS, pixel integration,
		centering offsets, `stepk`, `maxk`, and FFT-friendly rounding.

		The default per-draw limits are a 4096-pixel FFT side and 16,777,216 output
		pixels. These are allocation-size limits, not a process or node RAM limit.
		PSF construction, catalogs, detector arrays, spectroscopic image stacks and
		full-detector PSF convolution need their own memory budget. MPI ranks multiply
		the concurrently required memory.

		## Rendering policy

		Normal draws retain their original folding threshold and drawing method.
		If the predicted grid or stamp exceeds a limit, the renderer tries increasing
		folding thresholds from `[0.005, 0.008, 0.01]`, using only entries larger than
		the original threshold. Every successful retry logs the attempted thresholds,
		stamp sizes, FFT sizes, `stepk`, `maxk`, and object context. The original bright
		object threshold of `5e-8` is not globally replaced.

		Relaxing the folding threshold changes the numerical accuracy and, for an
		automatically sized stamp, its extent. Validate wings, aperture fluxes, size,
		ellipticity and relevant science measurements for objects using retries.
		Set `folding_threshold_retries: []` for runs that must retain the original
		accuracy even if this means failing at the allocation limit.

		The existing photometric large-galaxy branch (either component radius >3 arcsec)
		already omits PSF convolution. This branch now draws its analytic profile using
		pixel-integrated real-space rendering, in 256-by-256 tiles covering the
		intersection of the original automatic stamp and detector. It preserves the
		detector origin, subpixel centering, WCS, shear, magnification and original
		stamp footprint. It does not drop an additional PSF or truncate the physical
		Sérsic profile. It does not apply this strategy to spectroscopic stamps, whose
		flux can be dispersed to other detector locations.

		`real_space` is not forced on PSF-convolved profiles. GalSim can fall back to an
		FFT when real-space convolution is unsupported. `no_pixel` also remains an
		FFT operation for non-analytic composite profiles and is guarded accordingly.

		GalSim 2.8.4 normally warns and proceeds for oversized FFTs. The helper converts
		only `GalSimFFTSizeWarning` to an exception during its draw, as a backup to the
		allocation-free preflight. It does not change the global
		`galsim.errors.raise_fft_size_error` setting. Exhausted retries raise
		`RenderMemoryError`, which propagates through object and exposure handlers.
		The affected run fails instead of silently omitting the object or writing a
		completed science image with missing flux. Discard partial products from a
		failed run before rerunning with revised settings.

		## Configuration

		Defaults apply without configuration changes. Override them under the
		observation configuration's existing `call_sequence.scie_obs` section:

		```yaml
		call_sequence:
		scie_obs:
		# Keep the other scie_obs settings in your configuration.
		rendering:
		maximum_fft_size: 4096
		maximum_stamp_pixels: 16777216
		tile_size: 256
		folding_threshold_retries: [0.005, 0.008, 0.01]
		```

		Increasing a limit permits larger allocations. Set limits with enough headroom
		for detector/PSF data and the number of simultaneous MPI processes. Reducing
		`tile_size` bounds temporary real-space images more tightly, but does not
		change FFT memory for other branches.

		## Validation

		Run `python -m pytest -q tests/test_rendering.py` in the `csst_sim` environment.
		The tests compare ordinary draws with GalSim, compare tiled and full-stamp
		real-space images including detector edges and nontrivial WCS, exercise the
		original large-Sérsic example through `Galaxy.drawObj_multiband`, and verify
		that oversized draws are rejected before calling `drawImage`. The preflight
		uses GalSim's internal centering helpers; rerun these tests when upgrading the
		pinned GalSim dependency.

		On the development machine, the implemented tiled path rendered the original
		`n=4`, `half_light_radius=13.771078`, `g1=-0.01157`, `g2=0.00185` example onto
		a 64-by-64 detector at scale 0.074 in 0.077 seconds with 87 MiB peak process RSS.
		Its flux was 0.1419769898056984 and central pixel 0.0004157892254. This is a
		small synthetic benchmark, not a full-detector or production-PSF performance
		claim. A large detector can still require substantial real-space integration
		time, even though tile buffers stay small.

observation_sim/ObservationSim.py

+4 −0

Original line number	Diff line number	Diff line
		import os
		import numpy as np
		import mpi4py.MPI as MPI
		from observation_sim.rendering import RenderMemoryError
		import galsim
		import psutil
		import gc
		@@ -204,6 +205,9 @@ class Observation(object):
		obs_param=obs_param,
		)
		chip_output.Log_info("Finished simulation step: %s" % (step))
		except RenderMemoryError as e:
		chip_output.Log_error(str(e))
		raise
		except Exception as e:
		traceback.print_exc()
		chip_output.Log_error(e)

observation_sim/mock_objects/Galaxy.py

+30 −7

Original line number	Diff line number	Diff line
		@@ -9,6 +9,11 @@ from observation_sim.mock_objects._util import (
		)
		from observation_sim.mock_objects.SpecDisperser import SpecDisperser
		from observation_sim.mock_objects.MockObject import MockObject
		from observation_sim.rendering import (
		RenderMemoryError,
		RenderPolicy,
		add_analytic_to_detector,
		)

		# import tracemalloc

		@@ -246,7 +251,19 @@ class Galaxy(MockObject):
		# gal = self.bfrac * bulge + (1.0 - self.bfrac - kfrac) * disk + kfrac * knots

		# stamp = gal.drawImage(wcs=chip_wcs_local, method='phot', offset=offset, save_photons=True)
		stamp = gal.drawImage(wcs=chip_wcs_local, offset=offset)
		if big_galaxy:
		updated = add_analytic_to_detector(
		gal,
		chip.img,
		wcs=chip_wcs_local,
		nominal_center=(x_nominal, y_nominal),
		offset=offset,
		policy=getattr(self, "render_policy", RenderPolicy()),
		context=getattr(self, "render_context", f"galaxy id={self.id}"),
		logger=self.logger,
		)
		return int(updated), pos_shear
		stamp = self._draw_image(gal, wcs=chip_wcs_local, offset=offset)
		if np.sum(np.isnan(stamp.array)) > 0:
		# ERROR happens
		return 2, pos_shear
		@@ -430,13 +447,15 @@ class Galaxy(MockObject):
		)
		star_p = galsim.Convolve(psf, gal)
		if nnx == 0:
		galImg = star_p.drawImage(
		galImg = self._draw_image(
		star_p,
		wcs=chip_wcs_local, offset=offset, method="no_pixel"
		)
		nnx = galImg.xmax - galImg.xmin + 1
		nny = galImg.ymax - galImg.ymin + 1
		else:
		galImg = star_p.drawImage(
		galImg = self._draw_image(
		star_p,
		nx=nnx,
		ny=nny,
		wcs=chip_wcs_local,
		@@ -454,17 +473,21 @@ class Galaxy(MockObject):
		galImg_List.append(galImg)
		for order in ["C", "D", "E"]:
		galImg_List.append(galImg)
		except RenderMemoryError:
		raise
		except:
		try:
		psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img)
		star_p = galsim.Convolve(psf, gal)
		galImg = star_p.drawImage(wcs=chip_wcs_local, offset=offset)
		galImg = self._draw_image(star_p, wcs=chip_wcs_local, offset=offset)
		galImg.setOrigin(0, 0)
		if np.sum(np.isnan(galImg.array)) > 0:
		# ERROR happens
		return 2, pos_shear
		for order in ["A", "B", "C", "D", "E"]:
		galImg_List.append(galImg)
		except RenderMemoryError:
		raise
		except Exception:
		continue
		# starImg = gal.drawImage(

observation_sim/mock_objects/MockObject.py

+28 −12

Original line number	Diff line number	Diff line
		@@ -14,6 +14,7 @@ from observation_sim.mock_objects._util import (
		from observation_sim.mock_objects.SpecDisperser import SpecDisperser

		from observation_sim.instruments.chip import chip_utils
		from observation_sim.rendering import RenderMemoryError, RenderPolicy, draw_image


		class MockObject(object):
		@@ -43,6 +44,15 @@ class MockObject(object):
		self.additional_output_str = ""
		self.logger = logger

		def _draw_image(self, profile, **kwargs):
		return draw_image(
		profile,
		policy=getattr(self, "render_policy", RenderPolicy()),
		context=getattr(self, "render_context", f"{self.type} id={getattr(self, 'id', '?')}"),
		logger=self.logger,
		**kwargs,
		)

		def getMagFilter(self, filt):
		if filt.survey_type == "spectroscopic":
		return self.param["mag_use_normal"]
		@@ -208,11 +218,12 @@ class MockObject(object):
		star = psf.withFlux(nphotons)

		if EXTRA:
		stamp = star.drawImage(
		stamp = self._draw_image(
		star,
		method="no_pixel", wcs=chip_wcs_local, offset=offset
		)
		else:
		stamp = star.drawImage(wcs=chip_wcs_local, offset=offset)
		stamp = self._draw_image(star, wcs=chip_wcs_local, offset=offset)
		if np.sum(np.isnan(stamp.array)) > 0:
		continue
		stamp.setCenter(x_nominal, y_nominal)
		@@ -411,7 +422,7 @@ class MockObject(object):
		except:
		psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img)

		psf_img = psf.drawImage(nx=100, ny=100, wcs=local_wcs)
		psf_img = self._draw_image(psf, nx=100, ny=100, wcs=local_wcs)

		psf_img_m = psf_img.array

		@@ -559,13 +570,15 @@ class MockObject(object):
		# star_p = galsim.Convolve(psf, star)
		star_p = psf.withFlux(tel.pupil_area * exptime)
		if nnx == 0:
		starImg = star_p.drawImage(
		starImg = self._draw_image(
		star_p,
		wcs=chip_wcs_local, offset=offset, method="no_pixel"
		)
		nnx = starImg.xmax - starImg.xmin + 1
		nny = starImg.ymax - starImg.ymin + 1
		else:
		starImg = star_p.drawImage(
		starImg = self._draw_image(
		star_p,
		nx=nnx,
		ny=nny,
		wcs=chip_wcs_local,
		@@ -580,11 +593,14 @@ class MockObject(object):
		starImg_List.append(starImg)
		for order in ["C", "D", "E"]:
		starImg_List.append(starImg)
		except RenderMemoryError:
		raise
		except:
		psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img)
		# star_p = galsim.Convolve(psf, star)
		star_p = psf.withFlux(tel.pupil_area * exptime)
		starImg = star_p.drawImage(
		starImg = self._draw_image(
		star_p,
		wcs=chip_wcs_local, offset=offset, method="no_pixel"
		)
		starImg.setOrigin(0, 0)
		@@ -819,7 +835,7 @@ class MockObject(object):
		star = star + star_temp

		pixelScale = 0.074
		stamp = star.drawImage(wcs=chip_wcs_local, offset=offset)
		stamp = self._draw_image(star, wcs=chip_wcs_local, offset=offset)
		# stamp = star.drawImage(nx=256, ny=256, scale=pixelScale)
		if np.sum(np.isnan(stamp.array)) > 0:
		return None

observation_sim/mock_objects/Stamp.py

+3 −2

Original line number	Diff line number	Diff line
		@@ -110,7 +110,7 @@ class Stamp(MockObject):
		else:
		gal = gal + gal_temp

		stamp = gal.drawImage(wcs=chip_wcs_local, offset=offset)
		stamp = self._draw_image(gal, wcs=chip_wcs_local, offset=offset)
		if np.sum(np.isnan(stamp.array)) > 0:
		# ERROR happens
		return 2, pos_shear
		@@ -244,7 +244,8 @@ class Stamp(MockObject):
		# # if fd_shear is not None:
		# # gal = gal.shear(fd_shear)

		starImg = gal.drawImage(
		starImg = self._draw_image(
		gal,
		wcs=chip_wcs_local, offset=offset, method="real_space"
		)