From b2d3cf1ce7bd28983341570e5dcb6bbbbd3b4b8b Mon Sep 17 00:00:00 2001 From: Patrick Sheehan Date: Thu, 18 Jun 2026 21:05:04 +0000 Subject: [PATCH 1/8] torch.no_grad for evaluating ML models as it speeds up calculations. --- pinballrt/dust.py | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/pinballrt/dust.py b/pinballrt/dust.py index 7ab4969..444f115 100644 --- a/pinballrt/dust.py +++ b/pinballrt/dust.py @@ -171,7 +171,8 @@ def ml_kabs(self, p=None, amax=None, nu=None, abundances=None, photon_list=None, samples += (torch.log10(nu),) samples = torch.transpose(torch.vstack(samples), 0, 1) - kabs = 10.**self.kabs_y_scaler.inverse_transform(self.kabs_model(self.kabs_x_scaler.transform(samples))).detach().flatten() + with torch.no_grad(): + kabs = 10.**self.kabs_y_scaler.inverse_transform(self.kabs_model(self.kabs_x_scaler.transform(samples))).detach().flatten() return kabs @@ -212,7 +213,8 @@ def ml_ksca(self, p=None, amax=None, nu=None, abundances=None, photon_list=None, samples += (torch.log10(nu),) samples = torch.transpose(torch.vstack(samples), 0, 1) - ksca = 10.**self.ksca_y_scaler.inverse_transform(self.ksca_model(self.ksca_x_scaler.transform(samples))).detach().flatten() + with torch.no_grad(): + ksca = 10.**self.ksca_y_scaler.inverse_transform(self.ksca_model(self.ksca_x_scaler.transform(samples))).detach().flatten() return ksca @@ -246,8 +248,9 @@ def ml_kext(self, p=None, amax=None, nu=None, abundances=None, photon_list=None, samples += (torch.log10(nu),) samples = torch.transpose(torch.vstack(samples), 0, 1) - return 10.**self.kabs_y_scaler.inverse_transform(self.kabs_model(self.kabs_x_scaler.transform(samples))).detach().flatten() + \ - 10.**self.ksca_y_scaler.inverse_transform(self.ksca_model(self.ksca_x_scaler.transform(samples))).detach().flatten() + with torch.no_grad(): + return 10.**self.kabs_y_scaler.inverse_transform(self.kabs_model(self.kabs_x_scaler.transform(samples))).detach().flatten() + \ + 10.**self.ksca_y_scaler.inverse_transform(self.ksca_model(self.ksca_x_scaler.transform(samples))).detach().flatten() def absorb(self, temperature): nphotons = frequency.numpy().size @@ -280,7 +283,8 @@ def random_nu_ml(self, p, amax, temperature, abundances=None): samples = torch.transpose(torch.vstack(samples), 0, 1) test_x = self.random_nu_x_scaler.transform(samples) - log10_nu = torch.clamp(self.random_nu_y_scaler.inverse_transform(self.random_nu_model(test_x).detach()), self.log10_nu_min, self.log10_nu_max) + with torch.no_grad(): + log10_nu = torch.clamp(self.random_nu_y_scaler.inverse_transform(self.random_nu_model(test_x).detach()), self.log10_nu_min, self.log10_nu_max) return 10.**log10_nu.numpy() @@ -322,9 +326,11 @@ def random_nu(self, photon_list, subset=None): test_x = TensorDataset(test_x) loader = DataLoader(test_x, batch_size=250000) - log10_nu = torch.cat([torch.clamp(self.random_nu_y_scaler.inverse_transform(self.random_nu_model(X).detach()), self.log10_nu_min, self.log10_nu_max) for X, in loader], 0) + with torch.no_grad(): + log10_nu = torch.cat([torch.clamp(self.random_nu_y_scaler.inverse_transform(self.random_nu_model(X).detach()), self.log10_nu_min, self.log10_nu_max) for X, in loader], 0) else: - log10_nu = torch.clamp(self.random_nu_y_scaler.inverse_transform(self.random_nu_model(test_x).detach()), self.log10_nu_min, self.log10_nu_max) + with torch.no_grad(): + log10_nu = torch.clamp(self.random_nu_y_scaler.inverse_transform(self.random_nu_model(test_x).detach()), self.log10_nu_min, self.log10_nu_max) nu = wp.from_torch(10.**torch.flatten(log10_nu)) @@ -343,7 +349,8 @@ def ml_planck_mean_opacity(self, p, amax, temperature, abundances=()): samples = torch.transpose(torch.vstack(samples), 0, 1) - return 10.**self.pmo_y_scaler.inverse_transform(self.pmo_model(self.pmo_x_scaler.transform(samples))).detach().flatten() + with torch.no_grad(): + return 10.**self.pmo_y_scaler.inverse_transform(self.pmo_model(self.pmo_x_scaler.transform(samples))).detach().flatten() def ml_step(self, photon_list, s, iphotons): nphotons = iphotons.size(0) From 53c1dabbdbeb10c1a794cdc87a2f18fd370e2122 Mon Sep 17 00:00:00 2001 From: Patrick Sheehan Date: Thu, 18 Jun 2026 21:05:46 +0000 Subject: [PATCH 2/8] Put source photon/ray emission into torch tensors on the appropriate device so they get an appropriate speedup on execution. --- pinballrt/sources.py | 85 +++++++++++++++++++++++--------------------- 1 file changed, 45 insertions(+), 40 deletions(-) diff --git a/pinballrt/sources.py b/pinballrt/sources.py index b8a1812..c466524 100644 --- a/pinballrt/sources.py +++ b/pinballrt/sources.py @@ -92,23 +92,29 @@ def __init__(self, luminosity, frequency, intensity, x=0., y=0., z=0.): self.random_nu_CPD /= self.random_nu_CPD[-1] def emit(self, nphotons, distance_unit, wavelength="random", simulation="thermal", device="cpu", timing={}): - theta = np.pi*np.random.rand(nphotons) - phi = 2*np.pi*np.random.rand(nphotons) - - position = np.hstack(((self.radius.to(distance_unit).value*np.sin(theta)*np.cos(phi))[:,np.newaxis], - (self.radius.to(distance_unit).value*np.sin(theta)*np.sin(phi))[:,np.newaxis], - (self.radius.to(distance_unit).value*np.cos(theta))[:,np.newaxis])) - - r_hat = np.array([np.sin(theta)*np.cos(phi), np.sin(theta)*np.sin(phi), np.cos(theta)]).T - theta_hat = np.array([np.cos(theta)*np.cos(phi), np.cos(theta)*np.sin(phi), -np.sin(theta)]).T - phi_hat = np.array([-np.sin(phi), np.cos(phi), np.zeros(nphotons)]).T - - cost = np.random.rand(nphotons) - sint = np.sqrt(1-cost**2) - phi = 2*np.pi*np.random.rand(nphotons) + theta = torch.pi*torch.rand(nphotons, device=device, dtype=torch.float32) + phi = 2*torch.pi*torch.rand(nphotons, device=device, dtype=torch.float32) + + position = torch.hstack((torch.unsqueeze(self.radius.to(distance_unit).value*torch.sin(theta)*torch.cos(phi), 1), + torch.unsqueeze(self.radius.to(distance_unit).value*torch.sin(theta)*torch.sin(phi), 1), + torch.unsqueeze(self.radius.to(distance_unit).value*torch.cos(theta), 1))) + + r_hat = torch.hstack([torch.unsqueeze(torch.sin(theta)*torch.cos(phi), 1), + torch.unsqueeze(torch.sin(theta)*torch.sin(phi), 1), + torch.unsqueeze(torch.cos(theta), 1)]) + theta_hat = torch.hstack([torch.unsqueeze(torch.cos(theta)*torch.cos(phi), 1), + torch.unsqueeze(torch.cos(theta)*torch.sin(phi), 1), + torch.unsqueeze(-torch.sin(theta), 1)]) + phi_hat = torch.hstack([torch.unsqueeze(-torch.sin(phi), 1), + torch.unsqueeze(torch.cos(phi), 1), + torch.unsqueeze(torch.zeros(nphotons, device=device), 1)]) + + cost = torch.rand(nphotons, device=device, dtype=torch.float32) + sint = torch.sqrt(1-cost**2) + phi = 2*torch.pi*torch.rand(nphotons, device=device, dtype=torch.float32) - direction = cost[:,np.newaxis]*r_hat + (sint*np.cos(phi))[:,np.newaxis]*phi_hat + (sint*np.sin(phi))[:,np.newaxis]*theta_hat - direction_frame = cost[:,np.newaxis]*r_hat + (sint*np.cos(phi))[:,np.newaxis]*phi_hat + (sint*np.sin(phi))[:,np.newaxis]*theta_hat + direction = torch.unsqueeze(cost, 1)*r_hat + torch.unsqueeze(torch.exp(sint*torch.cos(phi)), 1)*phi_hat + torch.unsqueeze(sint*torch.sin(phi), 1)*theta_hat + direction_frame = torch.unsqueeze(cost, 1)*r_hat + torch.unsqueeze(sint*torch.cos(phi), 1)*phi_hat + torch.unsqueeze(sint*torch.sin(phi), 1)*theta_hat if wavelength == "random": t1 = time.time() @@ -116,45 +122,46 @@ def emit(self, nphotons, distance_unit, wavelength="random", simulation="thermal t2 = time.time() timing["Random frequency generation"] = t2 - t1 else: - frequency = np.repeat((const.c / wavelength).to(u.GHz), nphotons).value + frequency = wp.from_torch(torch.ones(nphotons, device=device, dtype=torch.float32) * (const.c / wavelength).to(u.GHz).value) if simulation == "thermal": - photon_energy = np.repeat(self.luminosity.to(u.L_sun).value / nphotons, nphotons) + photon_energy = torch.ones(nphotons, device=device, dtype=torch.float32) * (self.luminosity.to(u.L_sun).value / nphotons) elif simulation == "scattering": - photon_energy = np.repeat((4.*np.pi**2*u.steradian*self.radius**2*self.intensity(frequency[0]*u.GHz)).to(distance_unit**2 * u.Jy).value / nphotons, nphotons) + photon_energy = torch.ones(nphotons, device=device, dtype=torch.float32) * ((4.*np.pi**2*u.steradian*self.radius**2*self.intensity(wp.to_torch(frequency)[0]*u.GHz)).to(distance_unit**2 * u.Jy).value / nphotons) with wp.ScopedDevice(device): photon_list = PhotonList() - photon_list.position = wp.array(position, dtype=wp.vec3) - photon_list.direction = wp.array(direction, dtype=wp.vec3) - photon_list.direction_frame = wp.array(direction_frame, dtype=wp.vec3) - photon_list.frequency = wp.array(frequency, dtype=float) - photon_list.energy = wp.array(photon_energy, dtype=float) + photon_list.position = wp.from_torch(position, dtype=wp.vec3) + photon_list.direction = wp.from_torch(direction, dtype=wp.vec3) + photon_list.direction_frame = wp.from_torch(direction_frame, dtype=wp.vec3) + photon_list.frequency = frequency + photon_list.energy = wp.from_torch(photon_energy) photon_list.in_grid = wp.ones(nphotons, dtype=bool) return photon_list def emit_rays(self, nu, distance_unit, ez, nrays, physical_pixel_size, device="cpu"): - theta = np.pi*np.random.rand(nrays) - phi = 2*np.pi*np.random.rand(nrays) + theta = torch.pi*torch.rand(nrays, device=device) + phi = 2*torch.pi*torch.rand(nrays, device=device) - position = np.hstack(((self.radius.to(distance_unit).value*np.sin(theta)*np.cos(phi))[:,np.newaxis], - (self.radius.to(distance_unit).value*np.sin(theta)*np.sin(phi))[:,np.newaxis], - (self.radius.to(distance_unit).value*np.cos(theta))[:,np.newaxis])) + position = torch.hstack((torch.unsqueeze(self.radius.to(distance_unit).value*torch.sin(theta)*torch.cos(phi), 1), + torch.unsqueeze(self.radius.to(distance_unit).value*torch.sin(theta)*torch.sin(phi), 1), + torch.unsqueeze(self.radius.to(distance_unit).value*torch.cos(theta), 1))) - direction = np.tile(ez, (nrays, 1)) + direction = torch.unsqueeze(torch.tensor(ez, dtype=torch.float32, device=device), 0).repeat(nrays, 1) - intensity = (np.tile(self.intensity(nu.data)*np.pi, (nrays, 1)) / nrays).to(u.Jy / u.steradian).value * ((self.radius / physical_pixel_size).decompose()**2).value - tau_intensity = np.zeros((nrays, nu.size), dtype=float) + intensity = (self.intensity(nu.data)*np.pi).to(u.Jy / u.steradian).value * ((self.radius / physical_pixel_size).decompose()**2).value + intensity = torch.unsqueeze(torch.tensor(intensity, dtype=torch.float32, device=device), 0).repeat(nrays, 1) + tau_intensity = torch.zeros((nrays, nu.size), dtype=torch.float32, device=device) with wp.ScopedDevice(device): ray_list = PhotonList() - ray_list.position = wp.array(position, dtype=wp.vec3) - ray_list.direction = wp.array(direction, dtype=wp.vec3) - ray_list.direction_frame = wp.array(direction, dtype=wp.vec3) + ray_list.position = wp.from_torch(position, dtype=wp.vec3) + ray_list.direction = wp.from_torch(direction, dtype=wp.vec3) + ray_list.direction_frame = wp.from_torch(direction, dtype=wp.vec3) ray_list.indices = wp.zeros(position.shape, dtype=int) - ray_list.intensity = wp.array2d(intensity, dtype=float) - ray_list.tau_intensity = wp.array2d(tau_intensity, dtype=float) + ray_list.intensity = wp.from_torch(intensity) + ray_list.tau_intensity = wp.from_torch(tau_intensity) ray_list.pixel_too_large = wp.zeros(nrays, dtype=bool) ray_list.density = wp.zeros(nrays, dtype=float) @@ -280,7 +287,7 @@ def emit(self, nphotons, distance_unit, wavelength="random", simulation="thermal with wp.ScopedDevice(device): # Flip directions to point inward - photon_list.direction = wp.array2d(-photon_list.direction.numpy(), dtype=wp.vec3) + photon_list.direction = wp.from_torch(-wp.to_torch(photon_list.direction), dtype=wp.vec3) # Check the distance to the outer wall of the grid and move photons just inside s = wp.zeros(nphotons, dtype=float) @@ -366,9 +373,7 @@ def set_grid(self, grid): self.log10_intensity_func = np.interp1d(np.log10(self.frequency.to(u.GHz).value), np.log10(self.spectrum.value), kind='linear') self.intensity = lambda nu: 10**self.log10_intensity_func(np.log10(nu.to(u.GHz).value)) * self.spectrum.unit - print("bloop") self.total_luminosity = ((self.grid.volume.cpu().numpy()*self.grid.distance_unit**3 * self.density).sum() *scipy.integrate.trapezoid(self.intensity(self.frequency), self.frequency)).to(u.L_sun) - print("blop") self.random_nu_CPD = scipy.integrate.cumulative_trapezoid(self.intensity(self.frequency), self.frequency, initial=0.) self.random_nu_CPD /= self.random_nu_CPD[-1] From d51df5b3b98239d48513b5eab01b227544a322ef Mon Sep 17 00:00:00 2001 From: Patrick Sheehan Date: Thu, 18 Jun 2026 21:18:23 +0000 Subject: [PATCH 3/8] Remove double calculation of photon opacities; doesn't provide a huge speedup, but still unnecessary. --- pinballrt/grids.py | 37 ++++++++----------------------------- 1 file changed, 8 insertions(+), 29 deletions(-) diff --git a/pinballrt/grids.py b/pinballrt/grids.py index 27b4e5c..bb55631 100644 --- a/pinballrt/grids.py +++ b/pinballrt/grids.py @@ -357,17 +357,12 @@ def photon_cell_properties(photon_list: PhotonList, @wp.kernel def update_frequency(photon_list: PhotonList, frequency: wp.array(dtype=float), - kabs: wp.array(dtype=float), - ksca: wp.array(dtype=float), iphotons: wp.array(dtype=int)): # pragma: no cover i = wp.tid() ip = iphotons[i] photon_list.frequency[ip] = frequency[i] - photon_list.kabs[ip] = kabs[i] - photon_list.ksca[ip] = ksca[i] - photon_list.albedo[ip] = ksca[i] / (kabs[i] + ksca[i]) @wp.kernel def random_direction(direction: wp.array(dtype=wp.vec3), @@ -417,29 +412,14 @@ def interact(self, photon_list: PhotonList, absorb, iabsorb, interact, iphotons, t1 = time.time() nabsorb = iabsorb.size(0) - #if not scattering and nabsorb > 0: - # photon_temperature = wp.zeros(nabsorb, dtype=float) - # wp.launch(kernel=self.photon_temperature, - # dim=(nabsorb,), - # inputs=[photon_list, self.grid.temperature, photon_temperature, iabsorb]) - t2 = time.time() - photon_temperature_time = t2 - t1 - - t1 = time.time() if not scattering and nabsorb > 0: new_frequency = self.dust.random_nu(photon_list, subset=absorb) - t2 = time.time() - absorb_random_nu_time = t2 - t1 - - t1 = time.time() - if not scattering and nabsorb > 0: + wp.launch(kernel=self.update_frequency, dim=(nabsorb,), - inputs=[photon_list, new_frequency, - self.dust.ml_kabs(photon_list=photon_list, nu=wp.to_torch(new_frequency), iphotons=iabsorb), - self.dust.ml_ksca(photon_list=photon_list, nu=wp.to_torch(new_frequency), iphotons=iabsorb), iabsorb]) + inputs=[photon_list, new_frequency, iabsorb]) t2 = time.time() - dust_interpolation_time = t2 - t1 + absorb_random_nu_time = t2 - t1 seed = np.random.randint(0, 100000) wp.launch(kernel=self.random_tau, dim=(nphotons,), inputs=[photon_list, iphotons, seed]) @@ -458,7 +438,7 @@ def interact(self, photon_list: PhotonList, absorb, iabsorb, interact, iphotons, dim=(nphotons,), inputs=[photon_list, self.grid, iphotons]) - return photon_temperature_time, dust_interpolation_time, photon_loc_time, absorb_random_nu_time + return photon_loc_time, absorb_random_nu_time def update_grid(self, timing={}): with wp.ScopedDevice(self.device): @@ -567,7 +547,7 @@ def ml_step(self, photon_list, s, iphotons): wp.launch(kernel=self.update_frequency, dim=(nphotons,), - inputs=[photon_list, wp.from_torch(frequency), self.dust.ml_kabs(photon_list=photon_list, nu=wp.from_torch(frequency, iphotons=iphotons)), self.dust.ml_ksca(photon_list=photon_list, nu=wp.from_torch(frequency, iphotons=iphotons)), iphotons]) + inputs=[photon_list, wp.from_torch(frequency), iphotons]) wp.launch(kernel=self.ml_rotate_direction, dim=(nphotons,), @@ -763,12 +743,11 @@ def propagate_photons(self, photon_list: PhotonList, use_ml_step=False, learning interaction_indices = iphotons_original[interaction] absorb = torch.logical_and(interaction, wp.to_torch(photon_list.absorb)) absorb_indices = iphotons_original[absorb] - tmp_photon_temp_time, tmp_dust_interpolation_time, tmp_photon_loc_time, tmp_absorb_random_nu_time = self.interact(photon_list, absorb, absorb_indices, interaction, interaction_indices, learning=learning) + tmp_photon_loc_time, tmp_absorb_random_nu_time = self.interact(photon_list, absorb, absorb_indices, interaction, interaction_indices, learning=learning) t2 = time.time() - absorb_time += t2 - t1 - tmp_dust_interpolation_time - tmp_photon_loc_time + absorb_time += t2 - t1 - tmp_photon_loc_time absorb_random_nu_time += tmp_absorb_random_nu_time #absorb_time += tmp_time - dust_interpolation_time += tmp_dust_interpolation_time photon_loc_time += tmp_photon_loc_time t1 = time.time() @@ -956,7 +935,7 @@ def propagate_photons_scattering(self, photon_list: PhotonList, inu: int, debug= interaction_indices = iphotons_original[interaction] absorb = torch.logical_and(interaction, wp.to_torch(photon_list.absorb)) absorb_indices = iphotons_original[absorb] - tmp_photon_temp_time, tmp_dust_interpolation_time, tmp_photon_loc_time, tmp_absorb_random_nu_time = self.interact(photon_list, absorb, absorb_indices, interaction, interaction_indices, scattering=True) + tmp_photon_loc_time, tmp_absorb_random_nu_time = self.interact(photon_list, absorb, absorb_indices, interaction, interaction_indices, scattering=True) t2 = time.time() absorb_time += t2 - t1 - tmp_photon_loc_time #absorb_time += tmp_time From bb030d58a013ffc6b0f32cbb5bfabc6e4b3a11f0 Mon Sep 17 00:00:00 2001 From: Patrick Sheehan Date: Fri, 19 Jun 2026 03:11:51 +0000 Subject: [PATCH 4/8] Ensure that the emit_rays function uses random numbers sampled with the proper dtype to prevent downstream failures. --- pinballrt/sources.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pinballrt/sources.py b/pinballrt/sources.py index c466524..1814b31 100644 --- a/pinballrt/sources.py +++ b/pinballrt/sources.py @@ -141,8 +141,8 @@ def emit(self, nphotons, distance_unit, wavelength="random", simulation="thermal return photon_list def emit_rays(self, nu, distance_unit, ez, nrays, physical_pixel_size, device="cpu"): - theta = torch.pi*torch.rand(nrays, device=device) - phi = 2*torch.pi*torch.rand(nrays, device=device) + theta = np.pi*torch.rand(nrays, device=device, dtype=torch.float32) + phi = 2*np.pi*torch.rand(nrays, device=device, dtype=torch.float32) position = torch.hstack((torch.unsqueeze(self.radius.to(distance_unit).value*torch.sin(theta)*torch.cos(phi), 1), torch.unsqueeze(self.radius.to(distance_unit).value*torch.sin(theta)*torch.sin(phi), 1), From 6e950f1f678f22b42f78b7f58f9c4d5c8669c3c9 Mon Sep 17 00:00:00 2001 From: Patrick Sheehan Date: Thu, 18 Jun 2026 22:13:44 -0500 Subject: [PATCH 5/8] On GPUs, emitting photons from source requires that the frequency be brought back to the CPU for calculating the intensity. --- pinballrt/sources.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/pinballrt/sources.py b/pinballrt/sources.py index 1814b31..74d8e72 100644 --- a/pinballrt/sources.py +++ b/pinballrt/sources.py @@ -122,12 +122,17 @@ def emit(self, nphotons, distance_unit, wavelength="random", simulation="thermal t2 = time.time() timing["Random frequency generation"] = t2 - t1 else: - frequency = wp.from_torch(torch.ones(nphotons, device=device, dtype=torch.float32) * (const.c / wavelength).to(u.GHz).value) + frequency = wp.from_torch(torch.ones(nphotons, device=device, dtype=torch.float32) * \ + (const.c / wavelength).to(u.GHz).value) if simulation == "thermal": - photon_energy = torch.ones(nphotons, device=device, dtype=torch.float32) * (self.luminosity.to(u.L_sun).value / nphotons) + photon_energy = torch.ones(nphotons, device=device, dtype=torch.float32) * \ + (self.luminosity.to(u.L_sun).value / nphotons) elif simulation == "scattering": - photon_energy = torch.ones(nphotons, device=device, dtype=torch.float32) * ((4.*np.pi**2*u.steradian*self.radius**2*self.intensity(wp.to_torch(frequency)[0]*u.GHz)).to(distance_unit**2 * u.Jy).value / nphotons) + photon_energy = torch.ones(nphotons, device=device, dtype=torch.float32) * \ + ((4.*np.pi**2*u.steradian*self.radius**2*self.intensity( + wp.to_torch(frequency)[0].cpu()*u.GHz) + ).to(distance_unit**2 * u.Jy).value / nphotons) with wp.ScopedDevice(device): photon_list = PhotonList() @@ -150,7 +155,8 @@ def emit_rays(self, nu, distance_unit, ez, nrays, physical_pixel_size, device="c direction = torch.unsqueeze(torch.tensor(ez, dtype=torch.float32, device=device), 0).repeat(nrays, 1) - intensity = (self.intensity(nu.data)*np.pi).to(u.Jy / u.steradian).value * ((self.radius / physical_pixel_size).decompose()**2).value + intensity = (self.intensity(nu.data)*np.pi).to(u.Jy / u.steradian).value * \ + ((self.radius / physical_pixel_size).decompose()**2).value intensity = torch.unsqueeze(torch.tensor(intensity, dtype=torch.float32, device=device), 0).repeat(nrays, 1) tau_intensity = torch.zeros((nrays, nu.size), dtype=torch.float32, device=device) From 855dddd248a41f1a9f6a66074bf03434da005901 Mon Sep 17 00:00:00 2001 From: Patrick Sheehan Date: Sun, 31 May 2026 13:06:25 +0000 Subject: [PATCH 6/8] Only update photon opacities if frequency changes or if cell dust properties change; also don't update opacities after frequency change AND after cell properties change - only do once. --- pinballrt/camera.py | 1 + pinballrt/grids.py | 51 +++++++++++++++++++++++++++++++++++--------- pinballrt/photons.py | 2 ++ pinballrt/sources.py | 1 + 4 files changed, 45 insertions(+), 10 deletions(-) diff --git a/pinballrt/camera.py b/pinballrt/camera.py index 701ec1d..4df887b 100644 --- a/pinballrt/camera.py +++ b/pinballrt/camera.py @@ -67,6 +67,7 @@ def emit_rays(self, x, y, nu, nx, ny, pixel_size): ray_list.p = wp.zeros(xflat.size, dtype=float) if self.grid.n_dust_abundances > 0: ray_list.dust_abundances = wp.array2d(np.zeros((xflat.size, self.grid.n_dust_abundances)), dtype=float) + ray_list.opacities_out_of_date = wp.zeros(xflat.size, dtype=bool) ray_list.radius = wp.array(np.zeros(xflat.shape), dtype=float) if isinstance(self.grid, LogUniformSphericalGrid): diff --git a/pinballrt/grids.py b/pinballrt/grids.py index bb55631..98de131 100644 --- a/pinballrt/grids.py +++ b/pinballrt/grids.py @@ -348,11 +348,27 @@ def photon_cell_properties(photon_list: PhotonList, photon_list.temperature[ip] = grid.temperature[ix, iy, iz] photon_list.density[ip] = grid.dust_density[ix, iy, iz] - photon_list.amax[ip] = grid.amax[ix, iy, iz] - photon_list.p[ip] = grid.p[ix, iy, iz] - + + new_amax = grid.amax[ix, iy, iz] + new_p = grid.p[ix, iy, iz] + + updated_dust_properties = False + if photon_list.amax[ip] != new_amax or photon_list.p[ip] != new_p: + updated_dust_properties = True + + photon_list.amax[ip] = new_amax + photon_list.p[ip] = new_p + for i in range(n_dust_abundances): - photon_list.dust_abundances[ip][i] = grid.dust_abundances[i, ix, iy, iz] + new_abundance = grid.dust_abundances[i, ix, iy, iz] + + if photon_list.dust_abundances[ip][i] != new_abundance: + updated_dust_properties = True + + photon_list.dust_abundances[ip][i] = new_abundance + + if updated_dust_properties: + photon_list.opacities_out_of_date[ip] = True @wp.kernel def update_frequency(photon_list: PhotonList, @@ -363,6 +379,7 @@ def update_frequency(photon_list: PhotonList, ip = iphotons[i] photon_list.frequency[ip] = frequency[i] + photon_list.opacities_out_of_date[ip] = True @wp.kernel def random_direction(direction: wp.array(dtype=wp.vec3), @@ -571,6 +588,8 @@ def set_photon_opacities(photon_list: PhotonList, photon_list.ksca[ip] = ksca[i] photon_list.albedo[ip] = ksca[i] / (kabs[i] + ksca[i]) + photon_list.opacities_out_of_date[ip] = False + def propagate_photons(self, photon_list: PhotonList, use_ml_step=False, learning=False, debug=False, timing={}, position=0, time_limit=np.inf): with wp.ScopedDevice(self.device): nphotons = photon_list.position.numpy().shape[0] @@ -604,6 +623,7 @@ def propagate_photons(self, photon_list: PhotonList, use_ml_step=False, learning photon_list.ksca = wp.zeros(nphotons, dtype=float) photon_list.albedo = wp.zeros(nphotons, dtype=float) photon_list.absorb = wp.zeros(nphotons, dtype=bool) + photon_list.opacities_out_of_date = wp.zeros(nphotons, dtype=bool) progress_bar = tqdm(total=nphotons, position=position, leave=True) @@ -765,12 +785,14 @@ def propagate_photons(self, photon_list: PhotonList, use_ml_step=False, learning inputs=[photon_list, self.grid, iphotons, self.n_dust_abundances]) t1 = time.time() + iphotons_opacities = iphotons_original[torch.logical_and(wp.to_torch(photon_list.in_grid), + wp.to_torch(photon_list.opacities_out_of_date))] wp.launch(kernel=self.set_photon_opacities, - dim=(nphotons,), + dim=(iphotons_opacities.size(0),), inputs=[photon_list, - wp.from_torch(self.dust.ml_kabs(photon_list=photon_list, iphotons=iphotons)), - wp.from_torch(self.dust.ml_ksca(photon_list=photon_list, iphotons=iphotons)), - iphotons]) + wp.from_torch(self.dust.ml_kabs(photon_list=photon_list, iphotons=iphotons_opacities)), + wp.from_torch(self.dust.ml_ksca(photon_list=photon_list, iphotons=iphotons_opacities)), + iphotons_opacities]) t2 = time.time() dust_interpolation_time += t2 - t1 @@ -822,6 +844,8 @@ def update_photon_opacities(photon_list: PhotonList, photon_list.ksca[ip] = grid.ksca[inu, ix, iy, iz] photon_list.albedo[ip] = photon_list.ksca[ip] / (photon_list.kabs[ip] + photon_list.ksca[ip]) + photon_list.opacities_out_of_date[ip] = False + def propagate_photons_scattering(self, photon_list: PhotonList, inu: int, debug=False, timing={}, position=0): with wp.ScopedDevice(self.device): nphotons = photon_list.position.numpy().shape[0] @@ -849,6 +873,7 @@ def propagate_photons_scattering(self, photon_list: PhotonList, inu: int, debug= photon_list.ksca = wp.zeros(nphotons, dtype=float) photon_list.albedo = wp.zeros(nphotons, dtype=float) photon_list.absorb = wp.zeros(nphotons, dtype=bool) + photon_list.opacities_out_of_date = wp.zeros(nphotons, dtype=bool) progress_bar = tqdm(total=nphotons, position=position, leave=True) @@ -954,9 +979,11 @@ def propagate_photons_scattering(self, photon_list: PhotonList, inu: int, debug= inputs=[photon_list, self.grid, iphotons, self.n_dust_abundances]) t1 = time.time() + iphotons_opacities = iphotons_original[torch.logical_and(wp.to_torch(photon_list.in_grid), + wp.to_torch(photon_list.opacities_out_of_date))] wp.launch(kernel=self.update_photon_opacities, - dim=(nphotons,), - inputs=[photon_list, self.grid, inu, iphotons]) + dim=(iphotons_opacities.size(0),), + inputs=[photon_list, self.grid, inu, iphotons_opacities]) t2 = time.time() dust_interpolation_time += t2 - t1 @@ -1233,6 +1260,7 @@ def propagate_rays_from_source(self, ray_list: PhotonList, frequency): if self.n_dust_abundances > 0: ray_list.dust_abundances = wp.zeros((nrays, self.n_dust_abundances), dtype=float) + ray_list.opacities_out_of_date = wp.zeros(nrays, dtype=bool) wp.launch(kernel=self.photon_cell_properties, dim=(nrays,), @@ -1341,6 +1369,7 @@ def emit(self, nphotons, wavelength="random", scattering=False, learning=False, photon_list.p = wp.zeros(nphotons, dtype=float) if self.n_dust_abundances > 0: photon_list.dust_abundances = wp.zeros((nphotons, self.n_dust_abundances), dtype=float) + photon_list.opacities_out_of_date = wp.zeros(nphotons, dtype=bool) if not learning: wp.launch(kernel=self.photon_cell_properties, @@ -1671,6 +1700,7 @@ def emit(self, nphotons, wavelength="random", scattering=False, learning=False, photon_list.p = wp.array(np.zeros(nphotons), dtype=float) if self.n_dust_abundances > 0: photon_list.dust_abundances = wp.zeros((nphotons, self.n_dust_abundances), dtype=float) + photon_list.opacities_out_of_date = wp.zeros(nphotons, dtype=bool) if not learning: wp.launch(kernel=self.photon_cell_properties, @@ -2153,6 +2183,7 @@ def emit(self, nphotons, wavelength="random", scattering=False, learning=False, photon_list.p = wp.array(np.zeros(nphotons), dtype=float) if self.n_dust_abundances > 0: photon_list.dust_abundances = wp.zeros((nphotons, self.n_dust_abundances), dtype=float) + photon_list.opacities_out_of_date = wp.zeros(nphotons, dtype=bool) if not learning: wp.launch(kernel=self.photon_cell_properties, diff --git a/pinballrt/photons.py b/pinballrt/photons.py index e43d8eb..37681d2 100644 --- a/pinballrt/photons.py +++ b/pinballrt/photons.py @@ -25,6 +25,8 @@ class PhotonList: p: wp.array(dtype=float) dust_abundances: wp.array2d(dtype=float) + opacities_out_of_date: wp.array(dtype=bool) + tau: wp.array(dtype=float) total_tau_abs: wp.array(dtype=float) diff --git a/pinballrt/sources.py b/pinballrt/sources.py index 74d8e72..848dbb8 100644 --- a/pinballrt/sources.py +++ b/pinballrt/sources.py @@ -501,6 +501,7 @@ def random_nu(self, nphotons, cell_coords): photon_list.p = wp.zeros(nphotons, dtype=float) if self.grid.n_dust_abundances > 0: photon_list.dust_abundances = wp.zeros((nphotons, self.grid.n_dust_abundances), dtype=float) + photon_list.opacities_out_of_date = wp.zeros(nphotons, dtype=bool) wp.launch(kernel=self.grid.photon_cell_properties, dim=(nphotons,), From 9ee7668ab85e89ac2929d2b99d13fb752a4723fb Mon Sep 17 00:00:00 2001 From: Patrick Sheehan Date: Fri, 19 Jun 2026 05:40:35 -0500 Subject: [PATCH 7/8] Enable/disable showing the progress bar. --- pinballrt/grids.py | 29 +++++++++++++++++++---------- pinballrt/model.py | 35 ++++++++++++++++++++++------------- 2 files changed, 41 insertions(+), 23 deletions(-) diff --git a/pinballrt/grids.py b/pinballrt/grids.py index 98de131..8153d27 100644 --- a/pinballrt/grids.py +++ b/pinballrt/grids.py @@ -590,7 +590,8 @@ def set_photon_opacities(photon_list: PhotonList, photon_list.opacities_out_of_date[ip] = False - def propagate_photons(self, photon_list: PhotonList, use_ml_step=False, learning=False, debug=False, timing={}, position=0, time_limit=np.inf): + def propagate_photons(self, photon_list: PhotonList, use_ml_step=False, learning=False, debug=False, timing={}, position=0, time_limit=np.inf, + progress=True): with wp.ScopedDevice(self.device): nphotons = photon_list.position.numpy().shape[0] iphotons_original = torch.arange(nphotons, dtype=torch.int32, device=wp.device_to_torch(wp.get_device())) @@ -625,7 +626,8 @@ def propagate_photons(self, photon_list: PhotonList, use_ml_step=False, learning photon_list.absorb = wp.zeros(nphotons, dtype=bool) photon_list.opacities_out_of_date = wp.zeros(nphotons, dtype=bool) - progress_bar = tqdm(total=nphotons, position=position, leave=True) + if progress: + progress_bar = tqdm(total=nphotons, position=position, leave=True) iphotons = iphotons_original[wp.to_torch(photon_list.in_grid)] nphotons = iphotons.size(0) @@ -752,7 +754,8 @@ def propagate_photons(self, photon_list: PhotonList, use_ml_step=False, learning t1 = time.time() iphotons = iphotons_original[wp.to_torch(photon_list.in_grid)] - progress_bar.update(iphotons_original.size(0) - iphotons.size(0) - nphotons_done) + if progress: + progress_bar.update(iphotons_original.size(0) - iphotons.size(0) - nphotons_done) nphotons_done = iphotons_original.size(0) - iphotons.size(0) nphotons = iphotons.size(0) t2 = time.time() @@ -772,7 +775,8 @@ def propagate_photons(self, photon_list: PhotonList, use_ml_step=False, learning t1 = time.time() iphotons = iphotons_original[wp.to_torch(photon_list.in_grid)] - progress_bar.update(iphotons_original.size(0) - iphotons.size(0) - nphotons_done) + if progress: + progress_bar.update(iphotons_original.size(0) - iphotons.size(0) - nphotons_done) nphotons_done = iphotons_original.size(0) - iphotons.size(0) nphotons = iphotons.size(0) t2 = time.time() @@ -796,7 +800,8 @@ def propagate_photons(self, photon_list: PhotonList, use_ml_step=False, learning t2 = time.time() dust_interpolation_time += t2 - t1 - progress_bar.close() + if progress: + progress_bar.close() timing["next_wall_time"] = next_wall_time timing["dust_interpolation_time"] = dust_interpolation_time @@ -846,7 +851,7 @@ def update_photon_opacities(photon_list: PhotonList, photon_list.opacities_out_of_date[ip] = False - def propagate_photons_scattering(self, photon_list: PhotonList, inu: int, debug=False, timing={}, position=0): + def propagate_photons_scattering(self, photon_list: PhotonList, inu: int, debug=False, timing={}, position=0, progress=True): with wp.ScopedDevice(self.device): nphotons = photon_list.position.numpy().shape[0] iphotons_original = torch.arange(nphotons, dtype=torch.int32, device=wp.device_to_torch(wp.get_device())) @@ -875,7 +880,8 @@ def propagate_photons_scattering(self, photon_list: PhotonList, inu: int, debug= photon_list.absorb = wp.zeros(nphotons, dtype=bool) photon_list.opacities_out_of_date = wp.zeros(nphotons, dtype=bool) - progress_bar = tqdm(total=nphotons, position=position, leave=True) + if progress: + progress_bar = tqdm(total=nphotons, position=position, leave=True) iphotons = iphotons_original[wp.to_torch(photon_list.in_grid)] nphotons = iphotons.size(0) @@ -949,7 +955,8 @@ def propagate_photons_scattering(self, photon_list: PhotonList, inu: int, debug= t1 = time.time() iphotons = iphotons_original[torch.logical_and(wp.to_torch(photon_list.in_grid), wp.to_torch(photon_list.total_tau_abs) < 30.)] - progress_bar.update(iphotons_original.size(0) - iphotons.size(0) - nphotons_done) + if progress: + progress_bar.update(iphotons_original.size(0) - iphotons.size(0) - nphotons_done) nphotons_done = iphotons_original.size(0) - iphotons.size(0) nphotons = iphotons.size(0) t2 = time.time() @@ -967,7 +974,8 @@ def propagate_photons_scattering(self, photon_list: PhotonList, inu: int, debug= t1 = time.time() iphotons = iphotons_original[torch.logical_and(wp.to_torch(photon_list.in_grid), wp.to_torch(photon_list.total_tau_abs) < 30.)] - progress_bar.update(iphotons_original.size(0) - iphotons.size(0) - nphotons_done) + if progress: + progress_bar.update(iphotons_original.size(0) - iphotons.size(0) - nphotons_done) nphotons_done = iphotons_original.size(0) - iphotons.size(0) nphotons = iphotons.size(0) t2 = time.time() @@ -987,7 +995,8 @@ def propagate_photons_scattering(self, photon_list: PhotonList, inu: int, debug= t2 = time.time() dust_interpolation_time += t2 - t1 - progress_bar.close() + if progress: + progress_bar.close() timing["next_wall_time"] = next_wall_time timing["dust_interpolation_time"] = dust_interpolation_time diff --git a/pinballrt/model.py b/pinballrt/model.py index dbd13d8..029e0f3 100644 --- a/pinballrt/model.py +++ b/pinballrt/model.py @@ -116,7 +116,7 @@ def add_sources(self, sources): self.grid_list[device].add_sources(sources) def thermal_mc(self, nphotons, use_ml_step=False, Qthresh=2.0, Delthresh=1.1, p=99., device="cpu", - return_timing=False, nbatch=1): + return_timing=False, nbatch=1, progress=True): """ Perform a thermal Monte Carlo simulation. @@ -147,7 +147,9 @@ def thermal_mc(self, nphotons, use_ml_step=False, Qthresh=2.0, Delthresh=1.1, p= while count < 10: iter_timing = {} - print("Iteration", count) + if progress: + print("Iteration", count) + treallyold = told.copy() told = self.grid_list[device].grid.temperature.numpy().copy() @@ -160,7 +162,8 @@ def thermal_mc(self, nphotons, use_ml_step=False, Qthresh=2.0, Delthresh=1.1, p= SeedSequence(np.random.randint(10000)).spawn(njobs), [nphotons]*njobs, [njobs]*njobs, - [use_ml_step]*njobs)) + [use_ml_step]*njobs, + [progress]*njobs)) results = [r for r in result] total_energy = [r[0] for r in results] iter_timing["photon propagation"] = dict(zip([str(i) for i in range(njobs)], [r[1] for r in results])) @@ -193,18 +196,21 @@ def thermal_mc(self, nphotons, use_ml_step=False, Qthresh=2.0, Delthresh=1.1, p= Del = max(Q/Qold, Qold/Q) - print(count, Q, Del) + if progress: + print(count, Q, Del) + if Q < Qthresh and Del < Delthresh: break else: - print(count) + if progress: + print(count) count += 1 if return_timing: return timing - def scattering_mc(self, nphotons, wavelengths, device="cpu", return_timing=False, nbatch=1, set_grid_opacities=True): + def scattering_mc(self, nphotons, wavelengths, device="cpu", return_timing=False, nbatch=1, set_grid_opacities=True, progress=True): """ Perform a scattering Monte Carlo simulation. @@ -250,7 +256,8 @@ def scattering_mc(self, nphotons, wavelengths, device="cpu", return_timing=False [nphotons]*njobs, [njobs]*njobs, [wavelength]*njobs, - [i]*njobs,)) + [i]*njobs, + [progress]*njobs)) results = [r for r in result] total_scattering = [r[0] for r in results] iter_timing["photon propagation"] = dict(zip([str(i) for i in range(njobs)], [r[1] for r in results])) @@ -283,7 +290,8 @@ def scattering_mc(self, nphotons, wavelengths, device="cpu", return_timing=False return timing def make_image(self, npix=100, pixel_size=None, channels=None, rest_frequency=None, incl=0, pa=0, distance=1*u.pc, - include_dust=True, include_gas=True, include_sources=True, nphotons=100000, device="cpu", return_timing=False): + include_dust=True, include_gas=True, include_sources=True, nphotons=100000, device="cpu", return_timing=False, + progress=True): """ Create an image from the dust distribution. @@ -356,7 +364,8 @@ def make_image(self, npix=100, pixel_size=None, channels=None, rest_frequency=No self.grid_list[device].set_grid_opacities(nu) if include_dust: - timing["scattering"] = self.scattering_mc(nphotons, lam, device=device, set_grid_opacities=False, return_timing=True) + timing["scattering"] = self.scattering_mc(nphotons, lam, device=device, set_grid_opacities=False, return_timing=True, + progress=progress) # Now set up the image proper. @@ -456,20 +465,20 @@ def make_spectrum(self, lam=np.array([1.])*u.micron, incl=0, pa=0, distance=1*u. return spectrum def thermal_mc_task(args): - grid, position, s, nphotons, njobs, use_ml_step = args + grid, position, s, nphotons, njobs, use_ml_step, progress = args seed(s.generate_state(1)[0]) iter_timing = {} photon_list = grid.emit(int(nphotons / njobs), timing=iter_timing) - grid.propagate_photons(photon_list, use_ml_step=use_ml_step, timing=iter_timing, position=position) + grid.propagate_photons(photon_list, use_ml_step=use_ml_step, timing=iter_timing, position=position, progress=progress) return grid.grid.energy.numpy(), iter_timing def scattering_mc_task(args): - grid, position, s, nphotons, njobs, wavelength, i = args + grid, position, s, nphotons, njobs, wavelength, i, progress = args seed(s.generate_state(1)[0]) iter_timing = {} photon_list = grid.emit(int(nphotons / njobs), wavelength, scattering=True, timing=iter_timing) - grid.propagate_photons_scattering(photon_list, i, timing=iter_timing, position=position) + grid.propagate_photons_scattering(photon_list, i, timing=iter_timing, position=position, progress=progress) return grid.scattering, iter_timing From 13024c9ef505c56f141f4768d329c23b39b04fad Mon Sep 17 00:00:00 2001 From: Patrick Sheehan Date: Sat, 20 Jun 2026 07:55:12 -0500 Subject: [PATCH 8/8] Make sure direction and direction frame access different memory now that they are coming from torch. Also remove spurious exp that somehow ended up in the direction. --- pinballrt/sources.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pinballrt/sources.py b/pinballrt/sources.py index 848dbb8..aa190ef 100644 --- a/pinballrt/sources.py +++ b/pinballrt/sources.py @@ -113,7 +113,7 @@ def emit(self, nphotons, distance_unit, wavelength="random", simulation="thermal sint = torch.sqrt(1-cost**2) phi = 2*torch.pi*torch.rand(nphotons, device=device, dtype=torch.float32) - direction = torch.unsqueeze(cost, 1)*r_hat + torch.unsqueeze(torch.exp(sint*torch.cos(phi)), 1)*phi_hat + torch.unsqueeze(sint*torch.sin(phi), 1)*theta_hat + direction = torch.unsqueeze(cost, 1)*r_hat + torch.unsqueeze(sint*torch.cos(phi), 1)*phi_hat + torch.unsqueeze(sint*torch.sin(phi), 1)*theta_hat direction_frame = torch.unsqueeze(cost, 1)*r_hat + torch.unsqueeze(sint*torch.cos(phi), 1)*phi_hat + torch.unsqueeze(sint*torch.sin(phi), 1)*theta_hat if wavelength == "random": @@ -164,7 +164,7 @@ def emit_rays(self, nu, distance_unit, ez, nrays, physical_pixel_size, device="c ray_list = PhotonList() ray_list.position = wp.from_torch(position, dtype=wp.vec3) ray_list.direction = wp.from_torch(direction, dtype=wp.vec3) - ray_list.direction_frame = wp.from_torch(direction, dtype=wp.vec3) + ray_list.direction_frame = wp.from_torch(copy.deepcopy(direction), dtype=wp.vec3) ray_list.indices = wp.zeros(position.shape, dtype=int) ray_list.intensity = wp.from_torch(intensity) ray_list.tau_intensity = wp.from_torch(tau_intensity)