def proton_dP(self, theta, resolution_um, energy_MeV):
        th = np.pi / 180 * theta
        N = [np.sin(th), 0, np.cos(th)]

        resolution = Q(resolution_um, 'um')
        px_s = int(self.Ws / resolution) + 1
        px_t = int(self.Wt / resolution) + 1
        px = [px_s, px_t]

        Bx = yt.off_axis_projection(self.F,
                                    self.C,
                                    N,
                                    self.W,
                                    px,
                                    'magnetic_field_x',
                                    north_vector=[0, 1, 0])
        By = yt.off_axis_projection(self.F,
                                    self.C,
                                    N,
                                    self.W,
                                    px,
                                    'magnetic_field_y',
                                    north_vector=[0, 1, 0])
        Bz = yt.off_axis_projection(self.F,
                                    self.C,
                                    N,
                                    self.W,
                                    px,
                                    'magnetic_field_z',
                                    north_vector=[0, 1, 0])

        # S = [ cos(th), 0, -sin(th)]
        # T = [ 0, 1, 0 ]
        # U = N

        Bs = cm * ((Bx * np.cos(th) - Bz * np.sin(th)) / cm).to_equivalent(
            'gauss', 'CGS')
        Bt = cm * (By / cm).to_equivalent('gauss', 'CGS')

        s0 = np.linspace(-self.Ws, self.Ws, px_s) / 2
        t0 = np.linspace(-self.Wt, self.Wt, px_t) / 2

        [S0, T0] = np.meshgrid(s0, t0, indexing='ij')

        # The momenta here will be converted to PIC units (P / (m_e c))
        # F ~= N x B = [
        dPs = -(q * Bt) / (m * c**2)
        dPt = +(q * Bs) / (m * c**2)

        E_normalised = Q(energy_MeV, 'MeV') / (m * c**2)
        P0 = np.sqrt((E_normalised + M / m)**2 - (M / m)**2)

        return S0, T0, dPs, dPt, P0
    def __init__(self,
                 f,
                 resolution_um,
                 smooth_um=0,
                 scale=1,
                 dtheta=5,
                 cachefile='cache.npz'):
        if type(f) is str:
            self.TS = yt.load(f)
        else:
            self.TS = f

        self.i = 0
        try:
            self.F = self.TS[self.i]
        except:
            self.F = self.TS
            self.TS = None
            self.times = np.array([F.current_time.to('ps')])
        else:
            self.times = np.array([F.current_time.to('ps') for F in self.TS])

        self.dtheta = dtheta
        self.n_angles = int(360 / self.dtheta)
        if 360 % self.dtheta != 0:
            print('Warning: dtheta doesn\'t divide a full rotation')

        self.C = self.F.domain_center

        Ws = np.sqrt(self.F.domain_width[0]**2 + self.F.domain_width[2]**2)
        Wt = self.F.domain_width[1]
        self.W = [Ws, Wt, Ws]

        resolution = Q(resolution_um, 'um')
        px_s = int(Ws / resolution) + 1
        px_t = int(Wt / resolution) + 1
        self.px = [px_s, px_t]

        self.s0 = np.linspace(-Ws, Ws, px_s) / 2
        self.t0 = np.linspace(-Wt, Wt, px_t) / 2

        self.theta = 0
        self.strength_scale = scale
        self.smoothing = smooth_um / resolution_um

        self.cachefile = cachefile
        try:
            self._cache = np.load(self.cachefile,
                                  allow_pickle=True)['cache'][()]
        except FileNotFoundError:
            self._cache = {}
import numpy as np

import matplotlib.pyplot as plt

import yt
from yt import YTQuantity as Q

from scipy.ndimage.filters import gaussian_filter

q = yt.physical_constants.charge_proton
M = yt.physical_constants.mass_hydrogen
m = yt.physical_constants.mass_electron
c = yt.physical_constants.speed_of_light

cm = Q(1, 'cm')
gauss = Q(1, 'gauss')


class FieldsYT():
    '''
    Implements a radiography routine that rotates around the z axis.
    theta = 0 looks along the x axis; theta = 90 along the y-axis
    '''
    def __init__(self,
                 f,
                 resolution_um,
                 smooth_um=0,
                 scale=1,
                 dtheta=5,
                 cachefile='cache.npz'):
        if type(f) is str:
 def set_distance(self, L_cm):
     self.L = Q(L_cm, 'cm')
 def set_fringe(self, fringe_um, fringe_angle):
     fringe_spacing = Q(fringe_um, 'um')
     self.k_s = (2 * pi / fringe_spacing).to('1/cm') * np.array(
         [np.cos(fringe_angle), np.sin(fringe_angle)])
 def set_fringe(self, fringe_um):
     fringe_spacing = Q(fringe_um, 'um')
     self.k_s = (2 * pi / fringe_spacing).to('1/cm')
 def set_wavelength(self, wavelength_nm):
     lbd = Q(wavelength_nm, 'nm')
     self.k = (2 * pi / lbd).to('1/cm')
     self.omega = (c * self.k).to('1/s')
     self.n_c = (m * eps_0 * self.omega**2 / e**2).to('cm**-3')
import numpy as np
from numpy import pi

import matplotlib.pyplot as plt

from scipy.interpolate import RegularGridInterpolator

import yt
from yt import YTQuantity as Q

c = yt.physical_constants.speed_of_light
m = yt.physical_constants.mass_electron
e = yt.physical_constants.charge_proton.to_equivalent('C', 'SI')
eps_0 = yt.physical_constants.eps_0
cm = Q(1, 'cm')


class Laser:
    def __init__(self, wavelength_nm=800):
        self.set_wavelength(wavelength_nm)

    def set_wavelength(self, wavelength_nm):
        lbd = Q(wavelength_nm, 'nm')
        self.k = (2 * pi / lbd).to('1/cm')
        self.omega = (c * self.k).to('1/s')
        self.n_c = (m * eps_0 * self.omega**2 / e**2).to('cm**-3')

    def get_critical_density(self):
        return self.n_c