fix: handle blooming overflow at image boundaries (7e65f7cd) · Commits · csst-sims / csst_msc_sim

observation_sim/instruments/chip/effects.py

+23 −18

Original line number	Diff line number	Diff line
		@@ -404,15 +404,22 @@ def MakeTrail(
		):
		"""
		direction: "up" or "down". For "up", bleeds along Y-decreasing direction; for "down", bleeds along Y-increasing direction.
		Excess charge crossing an image edge leaves the simulated image. A trail
		starting at an edge can still propagate in the inward direction.
		"""
		yi, xi = satuyxtuple
		height = imgarr.shape[0]
		if direction not in ("up", "down"):
		raise ValueError(f"Unknown blooming direction: {direction}")
		step = -1 if direction == "up" else 1
		aa = np.log(charge / fullwell) ** 3 # scale length of the bleeding trail
		yy = 1

		while charge > 0:
		if yi < 0 or yi > imgarr.shape[0] - 1:
		if yi < 0 or yi >= height:
		break
		if yi == 0 or yi == imgarr.shape[0] - 1:
		next_y = yi + step
		if not 0 <= next_y < height:
		imgarr[yi, xi] = fullwell
		break

		@@ -434,19 +441,18 @@ def MakeTrail(
		if aa <= 1:
		while imgarr[yi, xi] >= fullwell:
		imgarr[yi, xi] = fullwell
		if direction == "up":
		imgarr[yi - 1, xi] += charge
		charge = imgarr[yi - 1, xi] - fullwell
		yi -= 1
		# if yi < 0:
		if yi < 0 or yi == (imgarr.shape[0] // 2 - 1):
		next_y = yi + step
		if not 0 <= next_y < height:
		# The remaining charge leaves the image. End both loops
		# without indexing past the bottom or wrapping at the top.
		charge = 0
		break
		elif direction == "down":
		imgarr[yi + 1, xi] += charge
		charge = imgarr[yi + 1, xi] - fullwell
		yi += 1
		# if yi > imgarr.shape[0]:
		if yi > imgarr.shape[0] or yi == (imgarr.shape[0] // 2):
		imgarr[next_y, xi] += charge
		charge = imgarr[next_y, xi] - fullwell
		yi = next_y
		if (direction == "up" and yi == (height // 2 - 1)) or (
		direction == "down" and yi == (height // 2)
		):
		break
		else:
		# calculate bleeding trail:
		@@ -530,8 +536,7 @@ def ChargeFlow(imgarr, fullwell=9e4):
		trailcutfrac=0.9,
		)
		except Exception as e:
		print(e, "@pix ", (yi + 1, xi + 1))
		return imgarr
		raise RuntimeError(f"Blooming failed at pixel (y={yi}, x={xi})") from e
		return imgarr

tests/test_blooming.py

0 → 100644

+119 −0

Original line number	Diff line number	Diff line
		from types import SimpleNamespace

		import galsim
		import numpy as np
		import pytest

		from observation_sim.instruments.chip import effects
		from observation_sim.sim_steps.add_pattern_noise import add_blooming


		FULL_WELL = 90000.0


		@pytest.fixture(autouse=True)
		def equal_charge_split(monkeypatch):
		# Equal upward/downward charge makes the expected edge losses exact.
		monkeypatch.setattr(effects.np.random, "random", lambda: 0.5)


		def _apply_blooming(pixels):
		chip = SimpleNamespace(img=galsim.Image(pixels), full_well=FULL_WELL)
		simulator = SimpleNamespace(
		chip_output=SimpleNamespace(Log_info=lambda message: None)
		)
		add_blooming(simulator, chip, None, None, None, None, None)
		return chip.img.array


		@pytest.mark.parametrize("dtype", [np.float32, np.float64])
		@pytest.mark.parametrize("source_y, edge_y, direction", [(1, 0, "up"), (8, 9, "down")])
		def test_trail_stops_at_image_edge_without_wrapping(source_y, edge_y, direction, dtype):
		pixels = np.zeros((10, 3), dtype=dtype)
		pixels[source_y, 1] = FULL_WELL

		result = effects.MakeTrail(
		pixels,
		(source_y, 1),
		charge=1.5 * FULL_WELL,
		fullwell=FULL_WELL,
		direction=direction,
		)

		expected = np.zeros_like(pixels)
		expected[[source_y, edge_y], 1] = FULL_WELL
		np.testing.assert_array_equal(result, expected)
		# Of the 1.5 wells sent toward the edge, one fits and half a well leaves.
		assert result.sum() + 0.5 * FULL_WELL == 2.5 * FULL_WELL


		@pytest.mark.parametrize("dtype", [np.float32, np.float64])
		@pytest.mark.parametrize("source_y, inward_step", [(1, 1), (8, -1)])
		def test_blooming_processes_later_bright_pixels_after_edge_overflow(
		source_y, inward_step, dtype, capsys
		):
		pixels = np.zeros((10, 3), dtype=dtype)
		pixels[source_y, 0] = 4 * FULL_WELL
		pixels[source_y, 1] = 2 * FULL_WELL
		initial_charge = pixels.sum()

		result = _apply_blooming(pixels)

		expected = np.zeros_like(pixels)
		expected[source_y, :2] = FULL_WELL
		expected[source_y - inward_step, 0] = FULL_WELL
		expected[source_y + inward_step, 0] = FULL_WELL
		expected[source_y + 2 * inward_step, 0] = 0.5 * FULL_WELL
		expected[[source_y - 1, source_y + 1], 1] = 0.5 * FULL_WELL
		np.testing.assert_array_equal(result, expected)
		assert result.max() == FULL_WELL
		assert result.sum() + 0.5 * FULL_WELL == initial_charge
		assert capsys.readouterr().out == ""


		@pytest.mark.parametrize("dtype", [np.float32, np.float64])
		@pytest.mark.parametrize("source_y, inward_y", [(0, 1), (9, 8)])
		def test_blooming_from_edge_pixel_still_flows_inward(source_y, inward_y, dtype):
		pixels = np.zeros((10, 3), dtype=dtype)
		pixels[source_y, 1] = 2 * FULL_WELL
		initial_charge = pixels.sum()

		result = _apply_blooming(pixels)

		expected = np.zeros_like(pixels)
		expected[source_y, 1] = FULL_WELL
		expected[inward_y, 1] = 0.5 * FULL_WELL
		np.testing.assert_array_equal(result, expected)
		assert result.sum() + 0.5 * FULL_WELL == initial_charge


		def test_blooming_conserves_charge_when_trails_fit_inside_image():
		pixels = np.zeros((20, 3))
		pixels[5, 1] = 4 * FULL_WELL
		initial_charge = pixels.sum()

		result = _apply_blooming(pixels)

		expected = np.zeros_like(pixels)
		expected[4:7, 1] = FULL_WELL
		expected[[3, 7], 1] = 0.5 * FULL_WELL
		np.testing.assert_array_equal(result, expected)
		assert result.sum() == initial_charge


		def test_blooming_reports_unexpected_errors_instead_of_returning_partial_image(
		monkeypatch,
		):
		def fail_trail(args, *kwargs):
		raise IndexError("unexpected trail failure")

		monkeypatch.setattr(effects, "MakeTrail", fail_trail)
		pixels = np.zeros((10, 3))
		pixels[1, 1] = 2 * FULL_WELL

		with pytest.raises(
		RuntimeError, match=r"Blooming failed at pixel \(y=1, x=1\)"
		) as exc:
		_apply_blooming(pixels)

		assert isinstance(exc.value.__cause__, IndexError)