From 032ac5d6236fa016dc4fb4171664923ff0ea1190 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Thu, 4 Jan 2024 23:41:35 +0100 Subject: [PATCH] Convert beams (except for transform) to be batched --- cheetah/particles.py | 565 +++++++++++++++++-------- try_batched.ipynb | 951 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 1335 insertions(+), 181 deletions(-) create mode 100644 try_batched.ipynb diff --git a/cheetah/particles.py b/cheetah/particles.py index e31e31de..5a77d8be 100644 --- a/cheetah/particles.py +++ b/cheetah/particles.py @@ -2,6 +2,7 @@ import numpy as np import torch +from icecream import ic from scipy.constants import physical_constants from torch import nn from torch.distributions import MultivariateNormal @@ -220,10 +221,9 @@ def relativistic_gamma(self) -> torch.Tensor: @property def relativistic_beta(self) -> torch.Tensor: - relativistic_beta = ( - torch.sqrt(1 - 1 / (self.relativistic_gamma**2)) - if torch.abs(self.relativistic_gamma) > 0 - else torch.tensor(1.0) + relativistic_beta = torch.ones_like(self.relativistic_gamma) + relativistic_beta[torch.abs(self.relativistic_gamma) > 0] = torch.sqrt( + 1 - 1 / (self.relativistic_gamma[self.relativistic_gamma > 0] ** 2) ) return relativistic_beta @@ -314,7 +314,7 @@ def __init__( total_charge = ( total_charge if total_charge is not None - else torch.tensor(0.0, **factory_kwargs) + else torch.tensor([0.0], **factory_kwargs) ) self.total_charge = torch.as_tensor(total_charge, **factory_kwargs) self.energy = torch.as_tensor(energy, **factory_kwargs) @@ -340,22 +340,52 @@ def from_parameters( device=None, dtype=torch.float32, ) -> "ParameterBeam": + # Figure out if arguments were passed, figure out their shape + not_nones = [ + argument + for argument in [ + mu_x, + mu_xp, + mu_y, + mu_yp, + sigma_x, + sigma_xp, + sigma_y, + sigma_yp, + sigma_s, + sigma_p, + cor_x, + cor_y, + cor_s, + energy, + total_charge, + ] + if argument is not None + ] + shape = not_nones[0].shape if len(not_nones) > 0 else torch.Size([]) + if len(not_nones) > 1: + assert all( + argument.shape == shape for argument in not_nones + ), "Arguments must have the same shape." + # Set default values without function call in function signature - mu_x = mu_x if mu_x is not None else torch.tensor(0.0) - mu_xp = mu_xp if mu_xp is not None else torch.tensor(0.0) - mu_y = mu_y if mu_y is not None else torch.tensor(0.0) - mu_yp = mu_yp if mu_yp is not None else torch.tensor(0.0) - sigma_x = sigma_x if sigma_x is not None else torch.tensor(175e-9) - sigma_xp = sigma_xp if sigma_xp is not None else torch.tensor(2e-7) - sigma_y = sigma_y if sigma_y is not None else torch.tensor(175e-9) - sigma_yp = sigma_yp if sigma_yp is not None else torch.tensor(2e-7) - sigma_s = sigma_s if sigma_s is not None else torch.tensor(1e-6) - sigma_p = sigma_p if sigma_p is not None else torch.tensor(1e-6) - cor_x = cor_x if cor_x is not None else torch.tensor(0.0) - cor_y = cor_y if cor_y is not None else torch.tensor(0.0) - cor_s = cor_s if cor_s is not None else torch.tensor(0.0) - energy = energy if energy is not None else torch.tensor(1e8) - total_charge = total_charge if total_charge is not None else torch.tensor(0.0) + mu_x = mu_x if mu_x is not None else torch.full(shape, 0.0) + mu_xp = mu_xp if mu_xp is not None else torch.full(shape, 0.0) + mu_y = mu_y if mu_y is not None else torch.full(shape, 0.0) + mu_yp = mu_yp if mu_yp is not None else torch.full(shape, 0.0) + sigma_x = sigma_x if sigma_x is not None else torch.full(shape, 175e-9) + sigma_xp = sigma_xp if sigma_xp is not None else torch.full(shape, 2e-7) + sigma_y = sigma_y if sigma_y is not None else torch.full(shape, 175e-9) + sigma_yp = sigma_yp if sigma_yp is not None else torch.full(shape, 2e-7) + sigma_s = sigma_s if sigma_s is not None else torch.full(shape, 1e-6) + sigma_p = sigma_p if sigma_p is not None else torch.full(shape, 1e-6) + cor_x = cor_x if cor_x is not None else torch.full(shape, 0.0) + cor_y = cor_y if cor_y is not None else torch.full(shape, 0.0) + cor_s = cor_s if cor_s is not None else torch.full(shape, 0.0) + energy = energy if energy is not None else torch.full(shape, 1e8) + total_charge = ( + total_charge if total_charge is not None else torch.full(shape, 0.0) + ) mu = torch.stack( [ @@ -363,25 +393,26 @@ def from_parameters( mu_xp, mu_y, mu_yp, - torch.tensor(0.0), - torch.tensor(0.0), - torch.tensor(1.0), - ] + torch.full(shape, 0.0), + torch.full(shape, 0.0), + torch.full(shape, 1.0), + ], + dim=1, ) - cov = torch.zeros(7, 7) - cov[0, 0] = sigma_x**2 - cov[0, 1] = cor_x - cov[1, 0] = cor_x - cov[1, 1] = sigma_xp**2 - cov[2, 2] = sigma_y**2 - cov[2, 3] = cor_y - cov[3, 2] = cor_y - cov[3, 3] = sigma_yp**2 - cov[4, 4] = sigma_s**2 - cov[4, 5] = cor_s - cov[5, 4] = cor_s - cov[5, 5] = sigma_p**2 + cov = torch.zeros(shape[0], 7, 7) + cov[:, 0, 0] = sigma_x**2 + cov[:, 0, 1] = cor_x + cov[:, 1, 0] = cor_x + cov[:, 1, 1] = sigma_xp**2 + cov[:, 2, 2] = sigma_y**2 + cov[:, 2, 3] = cor_y + cov[:, 3, 2] = cor_y + cov[:, 3, 3] = sigma_yp**2 + cov[:, 4, 4] = sigma_s**2 + cov[:, 4, 5] = cor_s + cov[:, 5, 4] = cor_s + cov[:, 5, 5] = sigma_p**2 return cls( mu=mu, cov=cov, energy=energy, total_charge=total_charge, device=device @@ -404,18 +435,44 @@ def from_twiss( device=None, dtype=torch.float32, ) -> "ParameterBeam": + # Figure out if arguments were passed, figure out their shape + not_nones = [ + argument + for argument in [ + beta_x, + alpha_x, + emittance_x, + beta_y, + alpha_y, + emittance_y, + sigma_s, + sigma_p, + cor_s, + energy, + total_charge, + ] + if argument is not None + ] + shape = not_nones[0].shape if len(not_nones) > 0 else torch.Size([]) + if len(not_nones) > 1: + assert all( + argument.shape == shape for argument in not_nones + ), "Arguments must have the same shape." + # Set default values without function call in function signature - beta_x = beta_x if beta_x is not None else torch.tensor(0.0) - alpha_x = alpha_x if alpha_x is not None else torch.tensor(0.0) - emittance_x = emittance_x if emittance_x is not None else torch.tensor(0.0) - beta_y = beta_y if beta_y is not None else torch.tensor(0.0) - alpha_y = alpha_y if alpha_y is not None else torch.tensor(0.0) - emittance_y = emittance_y if emittance_y is not None else torch.tensor(0.0) - sigma_s = sigma_s if sigma_s is not None else torch.tensor(1e-6) - sigma_p = sigma_p if sigma_p is not None else torch.tensor(1e-6) - cor_s = cor_s if cor_s is not None else torch.tensor(0.0) - energy = energy if energy is not None else torch.tensor(1e8) - total_charge = total_charge if total_charge is not None else torch.tensor(0.0) + beta_x = beta_x if beta_x is not None else torch.full(shape, 0.0) + alpha_x = alpha_x if alpha_x is not None else torch.full(shape, 0.0) + emittance_x = emittance_x if emittance_x is not None else torch.full(shape, 0.0) + beta_y = beta_y if beta_y is not None else torch.full(shape, 0.0) + alpha_y = alpha_y if alpha_y is not None else torch.full(shape, 0.0) + emittance_y = emittance_y if emittance_y is not None else torch.full(shape, 0.0) + sigma_s = sigma_s if sigma_s is not None else torch.full(shape, 1e-6) + sigma_p = sigma_p if sigma_p is not None else torch.full(shape, 1e-6) + cor_s = cor_s if cor_s is not None else torch.full(shape, 0.0) + energy = energy if energy is not None else torch.full(shape, 1e8) + total_charge = ( + total_charge if total_charge is not None else torch.full(shape, 0.0) + ) sigma_x = torch.sqrt(emittance_x * beta_x) sigma_xp = torch.sqrt(emittance_x * (1 + alpha_x**2) / beta_x) @@ -451,10 +508,10 @@ def from_ocelot(cls, parray, device=None, dtype=torch.float32) -> "ParameterBeam total_charge = torch.tensor(np.sum(parray.q_array), dtype=torch.float32) return cls( - mu=mu, - cov=cov, - energy=energy, - total_charge=total_charge, + mu=mu.unsqueeze(0), + cov=cov.unsqueeze(0), + energy=energy.unsqueeze(0), + total_charge=total_charge.unsqueeze(0), device=device, dtype=dtype, ) @@ -474,10 +531,10 @@ def from_astra(cls, path: str, device=None, dtype=torch.float32) -> "ParameterBe total_charge = torch.tensor(np.sum(particle_charges), dtype=torch.float32) return cls( - mu=mu, - cov=cov, - energy=torch.tensor(energy, dtype=torch.float32), - total_charge=total_charge, + mu=mu.unsqueeze(0), + cov=cov.unsqueeze(0), + energy=torch.tensor(energy, dtype=torch.float32).unsqueeze(0), + total_charge=total_charge.unsqueeze(0), device=device, dtype=dtype, ) @@ -551,59 +608,59 @@ def transformed_to( @property def mu_x(self) -> torch.Tensor: - return self._mu[0] + return self._mu[:, 0] @property def sigma_x(self) -> torch.Tensor: - return torch.sqrt(self._cov[0, 0]) + return torch.sqrt(self._cov[:, 0, 0]) @property def mu_xp(self) -> torch.Tensor: - return self._mu[1] + return self._mu[:, 1] @property def sigma_xp(self) -> torch.Tensor: - return torch.sqrt(self._cov[1, 1]) + return torch.sqrt(self._cov[:, 1, 1]) @property def mu_y(self) -> torch.Tensor: - return self._mu[2] + return self._mu[:, 2] @property def sigma_y(self) -> torch.Tensor: - return torch.sqrt(self._cov[2, 2]) + return torch.sqrt(self._cov[:, 2, 2]) @property def mu_yp(self) -> torch.Tensor: - return self._mu[3] + return self._mu[:, 3] @property def sigma_yp(self) -> torch.Tensor: - return torch.sqrt(self._cov[3, 3]) + return torch.sqrt(self._cov[:, 3, 3]) @property def mu_s(self) -> torch.Tensor: - return self._mu[4] + return self._mu[:, 4] @property def sigma_s(self) -> torch.Tensor: - return torch.sqrt(self._cov[4, 4]) + return torch.sqrt(self._cov[:, 4, 4]) @property def mu_p(self) -> torch.Tensor: - return self._mu[5] + return self._mu[:, 5] @property def sigma_p(self) -> torch.Tensor: - return torch.sqrt(self._cov[5, 5]) + return torch.sqrt(self._cov[:, 5, 5]) @property def sigma_xxp(self) -> torch.Tensor: - return self._cov[0, 1] + return self._cov[:, 0, 1] @property def sigma_yyp(self) -> torch.Tensor: - return self._cov[2, 3] + return self._cov[:, 2, 3] def __repr__(self) -> str: return ( @@ -640,15 +697,14 @@ def __init__( factory_kwargs = {"device": device, "dtype": dtype} assert ( - len(particles) > 0 and particles.shape[1] == 7 + particles.shape[-2] > 0 and particles.shape[-1] == 7 ), "Particle vectors must be 7-dimensional." self.particles = particles.to(**factory_kwargs) - num_particles = len(self.particles) self.particle_charges = ( particle_charges.to(**factory_kwargs) if particle_charges is not None - else torch.zeros(num_particles, **factory_kwargs) + else torch.zeros(particles.shape[:2], **factory_kwargs) ) self.energy = energy.to(**factory_kwargs) @@ -696,52 +752,89 @@ def from_parameters( :param device: Device to move the beam's particle array to. If set to `"auto"` a CUDA GPU is selected if available. The CPU is used otherwise. """ + # Figure out if arguments were passed, figure out their shape + not_nones = [ + argument + for argument in [ + mu_x, + mu_xp, + mu_y, + mu_yp, + sigma_x, + sigma_xp, + sigma_y, + sigma_yp, + sigma_s, + sigma_p, + cor_x, + cor_y, + cor_s, + energy, + total_charge, + ] + if argument is not None + ] + shape = not_nones[0].shape if len(not_nones) > 0 else torch.Size([]) + if len(not_nones) > 1: + assert all( + argument.shape == shape for argument in not_nones + ), "Arguments must have the same shape." + # Set default values without function call in function signature num_particles = ( num_particles if num_particles is not None else torch.tensor(100_000) ) - mu_x = mu_x if mu_x is not None else torch.tensor(0.0) - mu_xp = mu_xp if mu_xp is not None else torch.tensor(0.0) - mu_y = mu_y if mu_y is not None else torch.tensor(0.0) - mu_yp = mu_yp if mu_yp is not None else torch.tensor(0.0) - sigma_x = sigma_x if sigma_x is not None else torch.tensor(175e-9) - sigma_xp = sigma_xp if sigma_xp is not None else torch.tensor(2e-7) - sigma_y = sigma_y if sigma_y is not None else torch.tensor(175e-9) - sigma_yp = sigma_yp if sigma_yp is not None else torch.tensor(2e-7) - sigma_s = sigma_s if sigma_s is not None else torch.tensor(1e-6) - sigma_p = sigma_p if sigma_p is not None else torch.tensor(1e-6) - cor_x = cor_x if cor_x is not None else torch.tensor(0.0) - cor_y = cor_y if cor_y is not None else torch.tensor(0.0) - cor_s = cor_s if cor_s is not None else torch.tensor(0.0) - energy = energy if energy is not None else torch.tensor(1e8) - total_charge = total_charge if total_charge is not None else torch.tensor(0.0) + mu_x = mu_x if mu_x is not None else torch.full(shape, 0.0) + mu_xp = mu_xp if mu_xp is not None else torch.full(shape, 0.0) + mu_y = mu_y if mu_y is not None else torch.full(shape, 0.0) + mu_yp = mu_yp if mu_yp is not None else torch.full(shape, 0.0) + sigma_x = sigma_x if sigma_x is not None else torch.full(shape, 175e-9) + sigma_xp = sigma_xp if sigma_xp is not None else torch.full(shape, 2e-7) + sigma_y = sigma_y if sigma_y is not None else torch.full(shape, 175e-9) + sigma_yp = sigma_yp if sigma_yp is not None else torch.full(shape, 2e-7) + sigma_s = sigma_s if sigma_s is not None else torch.full(shape, 1e-6) + sigma_p = sigma_p if sigma_p is not None else torch.full(shape, 1e-6) + cor_x = cor_x if cor_x is not None else torch.full(shape, 0.0) + cor_y = cor_y if cor_y is not None else torch.full(shape, 0.0) + cor_s = cor_s if cor_s is not None else torch.full(shape, 0.0) + energy = energy if energy is not None else torch.full(shape, 1e8) + total_charge = ( + total_charge if total_charge is not None else torch.full(shape, 0.0) + ) particle_charges = ( - torch.ones(num_particles, device=device, dtype=dtype) - * total_charge + torch.ones((shape[0], num_particles), device=device, dtype=dtype) + * total_charge.view(-1, 1) / num_particles ) mean = torch.stack( - [mu_x, mu_xp, mu_y, mu_yp, torch.tensor(0.0), torch.tensor(0.0)] + [mu_x, mu_xp, mu_y, mu_yp, torch.full(shape, 0.0), torch.full(shape, 0.0)], + dim=1, ) - cov = torch.zeros(6, 6) - cov[0, 0] = sigma_x**2 - cov[0, 1] = cor_x - cov[1, 0] = cor_x - cov[1, 1] = sigma_xp**2 - cov[2, 2] = sigma_y**2 - cov[2, 3] = cor_y - cov[3, 2] = cor_y - cov[3, 3] = sigma_yp**2 - cov[4, 4] = sigma_s**2 - cov[4, 5] = cor_s - cov[5, 4] = cor_s - cov[5, 5] = sigma_p**2 - - particles = torch.ones((num_particles, 7)) - distribution = MultivariateNormal(mean, covariance_matrix=cov) - particles[:, :6] = distribution.sample((num_particles,)) + cov = torch.zeros(shape[0], 6, 6) + cov[:, 0, 0] = sigma_x**2 + cov[:, 0, 1] = cor_x + cov[:, 1, 0] = cor_x + cov[:, 1, 1] = sigma_xp**2 + cov[:, 2, 2] = sigma_y**2 + cov[:, 2, 3] = cor_y + cov[:, 3, 2] = cor_y + cov[:, 3, 3] = sigma_yp**2 + cov[:, 4, 4] = sigma_s**2 + cov[:, 4, 5] = cor_s + cov[:, 5, 4] = cor_s + cov[:, 5, 5] = sigma_p**2 + + particles = torch.ones((shape[0], num_particles, 7)) + distributions = [ + MultivariateNormal(sample_mean, covariance_matrix=sample_cov) + for sample_mean, sample_cov in zip(mean, cov) + ] + particles[:, :, :6] = torch.stack( + [distribution.sample((num_particles,)) for distribution in distributions], + dim=0, + ) return cls( particles, @@ -769,21 +862,47 @@ def from_twiss( device=None, dtype=torch.float32, ) -> "ParticleBeam": + # Figure out if arguments were passed, figure out their shape + not_nones = [ + argument + for argument in [ + beta_x, + alpha_x, + emittance_x, + beta_y, + alpha_y, + emittance_y, + energy, + sigma_s, + sigma_p, + cor_s, + total_charge, + ] + if argument is not None + ] + shape = not_nones[0].shape if len(not_nones) > 0 else torch.Size([]) + if len(not_nones) > 1: + assert all( + argument.shape == shape for argument in not_nones + ), "Arguments must have the same shape." + # Set default values without function call in function signature num_particles = ( num_particles if num_particles is not None else torch.tensor(1_000_000) ) - beta_x = beta_x if beta_x is not None else torch.tensor(0.0) - alpha_x = alpha_x if alpha_x is not None else torch.tensor(0.0) - emittance_x = emittance_x if emittance_x is not None else torch.tensor(0.0) - beta_y = beta_y if beta_y is not None else torch.tensor(0.0) - alpha_y = alpha_y if alpha_y is not None else torch.tensor(0.0) - emittance_y = emittance_y if emittance_y is not None else torch.tensor(0.0) - energy = energy if energy is not None else torch.tensor(1e8) - sigma_s = sigma_s if sigma_s is not None else torch.tensor(1e-6) - sigma_p = sigma_p if sigma_p is not None else torch.tensor(1e-6) - cor_s = cor_s if cor_s is not None else torch.tensor(0.0) - total_charge = total_charge if total_charge is not None else torch.tensor(0.0) + beta_x = beta_x if beta_x is not None else torch.full(shape, 0.0) + alpha_x = alpha_x if alpha_x is not None else torch.full(shape, 0.0) + emittance_x = emittance_x if emittance_x is not None else torch.full(shape, 0.0) + beta_y = beta_y if beta_y is not None else torch.full(shape, 0.0) + alpha_y = alpha_y if alpha_y is not None else torch.full(shape, 0.0) + emittance_y = emittance_y if emittance_y is not None else torch.full(shape, 0.0) + energy = energy if energy is not None else torch.full(shape, 1e8) + sigma_s = sigma_s if sigma_s is not None else torch.full(shape, 1e-6) + sigma_p = sigma_p if sigma_p is not None else torch.full(shape, 1e-6) + cor_s = cor_s if cor_s is not None else torch.full(shape, 0.0) + total_charge = ( + total_charge if total_charge is not None else torch.full(shape, 0.0) + ) sigma_x = torch.sqrt(beta_x * emittance_x) sigma_xp = torch.sqrt(emittance_x * (1 + alpha_x**2) / beta_x) @@ -794,10 +913,10 @@ def from_twiss( return cls.from_parameters( num_particles=num_particles, - mu_x=torch.tensor(0.0), - mu_xp=torch.tensor(0.0), - mu_y=torch.tensor(0.0), - mu_yp=torch.tensor(0.0), + mu_x=torch.full(shape, 0.0), + mu_xp=torch.full(shape, 0.0), + mu_y=torch.full(shape, 0.0), + mu_yp=torch.full(shape, 0.0), sigma_x=sigma_x, sigma_xp=sigma_xp, sigma_y=sigma_y, @@ -850,40 +969,118 @@ def make_linspaced( :param device: Device to move the beam's particle array to. If set to `"auto"` a CUDA GPU is selected if available. The CPU is used otherwise. """ + # Figure out if arguments were passed, figure out their shape + not_nones = [ + argument + for argument in [ + mu_x, + mu_xp, + mu_y, + mu_yp, + sigma_x, + sigma_xp, + sigma_y, + sigma_yp, + sigma_s, + sigma_p, + energy, + total_charge, + ] + if argument is not None + ] + shape = not_nones[0].shape if len(not_nones) > 0 else torch.Size([]) + if len(not_nones) > 1: + assert all( + argument.shape == shape for argument in not_nones + ), "Arguments must have the same shape." # Set default values without function call in function signature num_particles = num_particles if num_particles is not None else torch.tensor(10) - mu_x = mu_x if mu_x is not None else torch.tensor(0.0) - mu_xp = mu_xp if mu_xp is not None else torch.tensor(0.0) - mu_y = mu_y if mu_y is not None else torch.tensor(0.0) - mu_yp = mu_yp if mu_yp is not None else torch.tensor(0.0) - sigma_x = sigma_x if sigma_x is not None else torch.tensor(175e-9) - sigma_xp = sigma_xp if sigma_xp is not None else torch.tensor(2e-7) - sigma_y = sigma_y if sigma_y is not None else torch.tensor(175e-9) - sigma_yp = sigma_yp if sigma_yp is not None else torch.tensor(2e-7) - sigma_s = sigma_s if sigma_s is not None else torch.tensor(0.0) - sigma_p = sigma_p if sigma_p is not None else torch.tensor(0.0) - energy = energy if energy is not None else torch.tensor(1e8) - total_charge = total_charge if total_charge is not None else torch.tensor(0.0) + mu_x = mu_x if mu_x is not None else torch.full(shape, 0.0) + mu_xp = mu_xp if mu_xp is not None else torch.full(shape, 0.0) + mu_y = mu_y if mu_y is not None else torch.full(shape, 0.0) + mu_yp = mu_yp if mu_yp is not None else torch.full(shape, 0.0) + sigma_x = sigma_x if sigma_x is not None else torch.full(shape, 175e-9) + sigma_xp = sigma_xp if sigma_xp is not None else torch.full(shape, 2e-7) + sigma_y = sigma_y if sigma_y is not None else torch.full(shape, 175e-9) + sigma_yp = sigma_yp if sigma_yp is not None else torch.full(shape, 2e-7) + sigma_s = sigma_s if sigma_s is not None else torch.full(shape, 0.0) + sigma_p = sigma_p if sigma_p is not None else torch.full(shape, 0.0) + energy = energy if energy is not None else torch.full(shape, 1e8) + total_charge = ( + total_charge if total_charge is not None else torch.full(shape, 0.0) + ) particle_charges = ( - torch.ones(num_particles, dtype=torch.float32) - * total_charge + torch.ones((shape[0], num_particles), device=device, dtype=dtype) + * total_charge.view(-1, 1) / num_particles ) - particles = torch.ones((num_particles, 7)) + particles = torch.ones((shape[0], num_particles, 7)) - particles[:, 0] = torch.linspace(mu_x - sigma_x, mu_x + sigma_x, num_particles) - particles[:, 1] = torch.linspace( - mu_xp - sigma_xp, mu_xp + sigma_xp, num_particles + particles[:, :, 0] = torch.stack( + [ + torch.linspace( + sample_mu_x - sample_sigma_x, + sample_mu_x + sample_sigma_x, + num_particles, + ) + for sample_mu_x, sample_sigma_x in zip(mu_x, sigma_x) + ], + dim=0, + ) + particles[:, :, 1] = torch.stack( + [ + torch.linspace( + sample_mu_xp - sample_sigma_xp, + sample_mu_xp + sample_sigma_xp, + num_particles, + ) + for sample_mu_xp, sample_sigma_xp in zip(mu_xp, sigma_xp) + ], + dim=0, ) - particles[:, 2] = torch.linspace(mu_y - sigma_y, mu_y + sigma_y, num_particles) - particles[:, 3] = torch.linspace( - mu_yp - sigma_yp, mu_yp + sigma_yp, num_particles + particles[:, :, 2] = torch.stack( + [ + torch.linspace( + sample_mu_y - sample_sigma_y, + sample_mu_y + sample_sigma_y, + num_particles, + ) + for sample_mu_y, sample_sigma_y in zip(mu_y, sigma_y) + ], + dim=0, + ) + particles[:, :, 3] = torch.stack( + [ + torch.linspace( + sample_mu_yp - sample_sigma_yp, + sample_mu_yp + sample_sigma_yp, + num_particles, + ) + for sample_mu_yp, sample_sigma_yp in zip(mu_yp, sigma_yp) + ], + dim=0, + ) + particles[:, :, 4] = torch.stack( + [ + torch.linspace( + -sample_sigma_s, sample_sigma_s, num_particles, device=device + ) + for sample_sigma_s in sigma_s + ], + dim=0, + ) + particles[:, :, 5] = torch.stack( + [ + torch.linspace( + -sample_sigma_p, sample_sigma_p, num_particles, device=device + ) + for sample_sigma_p in sigma_p + ], + dim=0, ) - particles[:, 4] = torch.linspace(-sigma_s, sigma_s, num_particles) - particles[:, 5] = torch.linspace(-sigma_p, sigma_p, num_particles) return cls( particles=particles, @@ -904,9 +1101,9 @@ def from_ocelot(cls, parray, device=None, dtype=torch.float32) -> "ParticleBeam" particle_charges = torch.tensor(parray.q_array) return cls( - particles=particles, - energy=torch.tensor(1e9 * parray.E), - particle_charges=particle_charges, + particles=particles.unsqueeze(0), + energy=torch.tensor(1e9 * parray.E).unsqueeze(0), + particle_charges=particle_charges.unsqueeze(0), device=device, dtype=dtype, ) @@ -921,9 +1118,9 @@ def from_astra(cls, path: str, device=None, dtype=torch.float32) -> "ParticleBea particles_7d[:, :6] = torch.from_numpy(particles) particle_charges = torch.from_numpy(particle_charges) return cls( - particles=particles_7d, - energy=torch.tensor(energy), - particle_charges=particle_charges, + particles=particles_7d.unsqueeze(0), + energy=torch.tensor(energy).unsqueeze(0), + particle_charges=particle_charges.unsqueeze(0), device=device, dtype=dtype, ) @@ -1043,115 +1240,121 @@ def __len__(self) -> int: @property def total_charge(self) -> torch.Tensor: - return torch.sum(self.particle_charges) + return torch.sum(self.particle_charges, dim=1) @property - def num_particles(self) -> torch.Tensor: - return len(self.particles) + def num_particles(self) -> int: + return self.particles.shape[1] @property def xs(self) -> Optional[torch.Tensor]: - return self.particles[:, 0] if self is not Beam.empty else None + return self.particles[:, :, 0] if self is not Beam.empty else None @xs.setter def xs(self, value: torch.Tensor) -> None: - self.particles[:, 0] = value + self.particles[:, :, 0] = value @property def mu_x(self) -> Optional[torch.Tensor]: - return self.xs.mean() if self is not Beam.empty else None + return self.xs.mean(dim=1) if self is not Beam.empty else None @property def sigma_x(self) -> Optional[torch.Tensor]: - return self.xs.std() if self is not Beam.empty else None + return self.xs.std(dim=1) if self is not Beam.empty else None @property def xps(self) -> Optional[torch.Tensor]: - return self.particles[:, 1] if self is not Beam.empty else None + return self.particles[:, :, 1] if self is not Beam.empty else None @xps.setter def xps(self, value: torch.Tensor) -> None: - self.particles[:, 1] = value + self.particles[:, :, 1] = value @property def mu_xp(self) -> Optional[torch.Tensor]: - return self.xps.mean() if self is not Beam.empty else None + return self.xps.mean(dim=1) if self is not Beam.empty else None @property def sigma_xp(self) -> Optional[torch.Tensor]: - return self.xps.std() if self is not Beam.empty else None + return self.xps.std(dim=1) if self is not Beam.empty else None @property def ys(self) -> Optional[torch.Tensor]: - return self.particles[:, 2] if self is not Beam.empty else None + return self.particles[:, :, 2] if self is not Beam.empty else None @ys.setter def ys(self, value: torch.Tensor) -> None: - self.particles[:, 2] = value + self.particles[:, :, 2] = value @property def mu_y(self) -> Optional[float]: - return self.ys.mean() if self is not Beam.empty else None + return self.ys.mean(dim=1) if self is not Beam.empty else None @property def sigma_y(self) -> Optional[torch.Tensor]: - return self.ys.std() if self is not Beam.empty else None + return self.ys.std(dim=1) if self is not Beam.empty else None @property def yps(self) -> Optional[torch.Tensor]: - return self.particles[:, 3] if self is not Beam.empty else None + return self.particles[:, :, 3] if self is not Beam.empty else None @yps.setter def yps(self, value: torch.Tensor) -> None: - self.particles[:, 3] = value + self.particles[:, :, 3] = value @property def mu_yp(self) -> Optional[torch.Tensor]: - return self.yps.mean() if self is not Beam.empty else None + return self.yps.mean(dim=1) if self is not Beam.empty else None @property def sigma_yp(self) -> Optional[torch.Tensor]: - return self.yps.std() if self is not Beam.empty else None + return self.yps.std(dim=1) if self is not Beam.empty else None @property def ss(self) -> Optional[torch.Tensor]: - return self.particles[:, 4] if self is not Beam.empty else None + return self.particles[:, :, 4] if self is not Beam.empty else None @ss.setter def ss(self, value: torch.Tensor) -> None: - self.particles[:, 4] = value + self.particles[:, :, 4] = value @property def mu_s(self) -> Optional[torch.Tensor]: - return self.ss.mean() if self is not Beam.empty else None + return self.ss.mean(dim=1) if self is not Beam.empty else None @property def sigma_s(self) -> Optional[torch.Tensor]: - return self.ss.std() if self is not Beam.empty else None + return self.ss.std(dim=1) if self is not Beam.empty else None @property def ps(self) -> Optional[torch.Tensor]: - return self.particles[:, 5] if self is not Beam.empty else None + return self.particles[:, :, 5] if self is not Beam.empty else None @ps.setter def ps(self, value: torch.Tensor) -> None: - self.particles[:, 5] = value + self.particles[:, :, 5] = value @property def mu_p(self) -> Optional[torch.Tensor]: - return self.ps.mean() if self is not Beam.empty else None + return self.ps.mean(dim=1) if self is not Beam.empty else None @property def sigma_p(self) -> Optional[torch.Tensor]: - return self.ps.std() if self is not Beam.empty else None + return self.ps.std(dim=1) if self is not Beam.empty else None @property def sigma_xxp(self) -> torch.Tensor: - return torch.mean((self.xs - self.mu_x) * (self.xps - self.mu_xp)) + return torch.mean( + (self.xs - self.mu_x.view(-1, 1)) * (self.xps - self.mu_xp.view(-1, 1)), + dim=1, + ) @property def sigma_yyp(self) -> torch.Tensor: - return torch.mean((self.ys - self.mu_y) * (self.yps - self.mu_yp)) + return torch.mean( + (self.ys - self.mu_y.view(-1, 1)) * (self.yps - self.mu_yp.view(-1, 1)), + dim=1, + ) def __repr__(self) -> str: return ( diff --git a/try_batched.ipynb b/try_batched.ipynb new file mode 100644 index 00000000..01012c30 --- /dev/null +++ b/try_batched.ipynb @@ -0,0 +1,951 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "\n", + "import cheetah" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "ParameterBeam(mu_x=tensor([8.2413e-07]), mu_xp=tensor([5.9885e-08]), mu_y=tensor([-1.7276e-06]), mu_yp=tensor([-1.1746e-07]), sigma_x=tensor([0.0002]), sigma_xp=tensor([3.6794e-06]), sigma_y=tensor([0.0002]), sigma_yp=tensor([3.6941e-06]), sigma_s=tensor([8.0116e-06]), sigma_p=tensor([0.0023]), energy=tensor([1.0732e+08])), total_charge=tensor([5.0000e-13]))" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# parameter_beam = cheetah.ParameterBeam.from_parameters(\n", + "# mu_x=torch.tensor([0.0, 1e-6]),\n", + "# sigma_x=torch.tensor([175e-9, 42e-8]),\n", + "# total_charge=torch.tensor([1e-12, 2e-12]),\n", + "# )\n", + "# parameter_beam = cheetah.ParameterBeam.from_twiss(\n", + "# beta_x=torch.tensor([61.47503078, 99.0]),\n", + "# alpha_x=torch.tensor([-1.21242463, -0.9]),\n", + "# emittance_x=torch.tensor([7.1971891e-13, 5.0e-13]),\n", + "# beta_y=torch.tensor([35.41897281, 60.0]),\n", + "# alpha_y=torch.tensor([0.66554622, 0.5]),\n", + "# emittance_y=torch.tensor([3.5866484e-15, 1.0e-15]),\n", + "# total_charge=torch.tensor([1e-12, 4e-12]),\n", + "# )\n", + "parameter_beam = cheetah.ParameterBeam.from_astra(\n", + " \"tests/resources/ACHIP_EA1_2021.1351.001\"\n", + ")\n", + "parameter_beam" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "ParticleBeam(n=100000, mu_x=tensor([8.2413e-07]), mu_xp=tensor([5.9885e-08]), mu_y=tensor([-1.7276e-06]), mu_yp=tensor([-1.1746e-07]), sigma_x=tensor([0.0002]), sigma_xp=tensor([3.6794e-06]), sigma_y=tensor([0.0002]), sigma_yp=tensor([3.6941e-06]), sigma_s=tensor([8.0116e-06]), sigma_p=tensor([0.0023]), energy=tensor([1.0732e+08])) total_charge=tensor([5.0000e-13]))" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# particle_beam = cheetah.ParticleBeam.from_parameters(\n", + "# mu_x=torch.tensor([0.0, 1e-6]),\n", + "# sigma_x=torch.tensor([175e-9, 42e-8]),\n", + "# total_charge=torch.tensor([1e-12, 2e-12]),\n", + "# )\n", + "# particle_beam = cheetah.ParticleBeam.from_twiss(\n", + "# beta_x=torch.tensor([61.47503078, 99.0]),\n", + "# alpha_x=torch.tensor([-1.21242463, -0.9]),\n", + "# emittance_x=torch.tensor([7.1971891e-13, 5.0e-13]),\n", + "# beta_y=torch.tensor([35.41897281, 60.0]),\n", + "# alpha_y=torch.tensor([0.66554622, 0.5]),\n", + "# emittance_y=torch.tensor([3.5866484e-15, 1.0e-15]),\n", + "# total_charge=torch.tensor([1e-12, 4e-12]),\n", + "# )\n", + "particle_beam = cheetah.ParticleBeam.from_astra(\n", + " \"tests/resources/ACHIP_EA1_2021.1351.001\"\n", + ")\n", + "particle_beam" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "100000" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.num_particles" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([5.0000e-13])" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.total_charge" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([5.0000e-13])" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.total_charge" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([8.2413e-07])" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.mu_x" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "torch.Size([1, 100000])" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.xs.shape" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([8.2413e-07])" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.mu_x" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([0.0002])" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.sigma_x" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([0.0002])" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.sigma_x" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([[ 5.0896e-08, 6.1503e-05, 2.3025e-08, ..., 2.0862e-04,\n", + " 5.6475e-05, -6.2661e-05]])" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.xs" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([[8.2413e-07]])" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.mu_x.view(-1, 1)" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([[-7.7323e-07, 6.0679e-05, -8.0110e-07, ..., 2.0780e-04,\n", + " 5.5651e-05, -6.3486e-05]])" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.xs - particle_beam.mu_x.view(-1, 1)" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([3.2265e-10])" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.emittance_x" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([3.2266e-10])" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.emittance_x" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([3.2312e-10])" + ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.emittance_y" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([3.2313e-10])" + ] + }, + "execution_count": 18, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.emittance_y" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([1.0000])" + ] + }, + "execution_count": 19, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.relativistic_beta" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([1.0000])" + ] + }, + "execution_count": 20, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.relativistic_beta" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([6.7759e-08])" + ] + }, + "execution_count": 21, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.normalized_emittance_x" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([6.7761e-08])" + ] + }, + "execution_count": 22, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.normalized_emittance_x" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([6.7858e-08])" + ] + }, + "execution_count": 23, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.normalized_emittance_y" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([6.7860e-08])" + ] + }, + "execution_count": 24, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.normalized_emittance_y" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'mu_x': tensor([8.2413e-07]),\n", + " 'mu_xp': tensor([5.9885e-08]),\n", + " 'mu_y': tensor([-1.7276e-06]),\n", + " 'mu_yp': tensor([-1.1746e-07]),\n", + " 'sigma_x': tensor([0.0002]),\n", + " 'sigma_xp': tensor([3.6794e-06]),\n", + " 'sigma_y': tensor([0.0002]),\n", + " 'sigma_yp': tensor([3.6941e-06]),\n", + " 'sigma_s': tensor([8.0116e-06]),\n", + " 'sigma_p': tensor([0.0023]),\n", + " 'energy': tensor([1.0732e+08])}" + ] + }, + "execution_count": 25, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.parameters" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'mu_x': tensor([8.2413e-07]),\n", + " 'mu_xp': tensor([5.9885e-08]),\n", + " 'mu_y': tensor([-1.7276e-06]),\n", + " 'mu_yp': tensor([-1.1746e-07]),\n", + " 'sigma_x': tensor([0.0002]),\n", + " 'sigma_xp': tensor([3.6794e-06]),\n", + " 'sigma_y': tensor([0.0002]),\n", + " 'sigma_yp': tensor([3.6941e-06]),\n", + " 'sigma_s': tensor([8.0116e-06]),\n", + " 'sigma_p': tensor([0.0023]),\n", + " 'energy': tensor([1.0732e+08])}" + ] + }, + "execution_count": 26, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.parameters" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([210.0120])" + ] + }, + "execution_count": 27, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.relativistic_gamma" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([210.0120])" + ] + }, + "execution_count": 28, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.relativistic_gamma" + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([94.8069])" + ] + }, + "execution_count": 29, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.beta_x" + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([94.8041])" + ] + }, + "execution_count": 30, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.beta_x" + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([-1.7257])" + ] + }, + "execution_count": 31, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.alpha_x" + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([-1.7256])" + ] + }, + "execution_count": 32, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.alpha_x" + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([94.9920])" + ] + }, + "execution_count": 33, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.beta_y" + ] + }, + { + "cell_type": "code", + "execution_count": 34, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([94.9891])" + ] + }, + "execution_count": 34, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.beta_y" + ] + }, + { + "cell_type": "code", + "execution_count": 35, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([-1.7355])" + ] + }, + "execution_count": 35, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.alpha_y" + ] + }, + { + "cell_type": "code", + "execution_count": 36, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([-1.7354])" + ] + }, + "execution_count": 36, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.alpha_y" + ] + }, + { + "cell_type": "code", + "execution_count": 37, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([5.7250e-06])" + ] + }, + "execution_count": 37, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.mu_s" + ] + }, + { + "cell_type": "code", + "execution_count": 38, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([5.7250e-06])" + ] + }, + "execution_count": 38, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.mu_s" + ] + }, + { + "cell_type": "code", + "execution_count": 39, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([5.5679e-10])" + ] + }, + "execution_count": 39, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.sigma_xxp" + ] + }, + { + "cell_type": "code", + "execution_count": 40, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([5.5679e-10])" + ] + }, + "execution_count": 40, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.sigma_xxp" + ] + }, + { + "cell_type": "code", + "execution_count": 41, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([5.6076e-10])" + ] + }, + "execution_count": 41, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "parameter_beam.sigma_yyp" + ] + }, + { + "cell_type": "code", + "execution_count": 42, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([5.6075e-10])" + ] + }, + "execution_count": 42, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam.sigma_yyp" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": 43, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "ParticleBeam(n=10, mu_x=tensor([0.0000e+00, 1.0000e-06]), mu_xp=tensor([0., 0.]), mu_y=tensor([0., 0.]), mu_yp=tensor([0., 0.]), sigma_x=tensor([1.1774e-07, 2.8258e-07]), sigma_xp=tensor([1.3456e-07, 1.3456e-07]), sigma_y=tensor([1.1774e-07, 1.1774e-07]), sigma_yp=tensor([1.3456e-07, 1.3456e-07]), sigma_s=tensor([0., 0.]), sigma_p=tensor([0., 0.]), energy=tensor([100000000., 100000000.])) total_charge=tensor([1.0000e-12, 2.0000e-12]))" + ] + }, + "execution_count": 43, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "particle_beam = cheetah.ParticleBeam.make_linspaced(\n", + " mu_x=torch.tensor([0.0, 1e-6]),\n", + " sigma_x=torch.tensor([175e-9, 42e-8]),\n", + " total_charge=torch.tensor([1e-12, 2e-12]),\n", + ")\n", + "particle_beam" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "cheetah-dev", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.18" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +}