Skip to content

Commit

Permalink
Move axes equal into main functions
Browse files Browse the repository at this point in the history
  • Loading branch information
ColwynGulliford committed Oct 10, 2024
1 parent 7cccfd1 commit 9a5734c
Show file tree
Hide file tree
Showing 2 changed files with 104 additions and 160 deletions.
222 changes: 72 additions & 150 deletions docs/examples/fields/corrector_modeling.ipynb

Large diffs are not rendered by default.

42 changes: 32 additions & 10 deletions pmd_beamphysics/fields/corrector_modeling.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ def plot_3d_vector(v,
return ax


def bfield_from_thin_straight_wire(x, y, z, p1, p2, current, plot_wire=False, elev=45, azim=-45, ax=None):
def bfield_from_thin_straight_wire(x, y, z, p1, p2, current, plot_wire=False, elev=45, azim=-45, ax=None, return_axes=False):

"""
Calculate the magnetic field generated by a thin, straight current-carrying wire at specified points in 3D space.
Expand Down Expand Up @@ -139,11 +139,16 @@ def bfield_from_thin_straight_wire(x, y, z, p1, p2, current, plot_wire=False, el

if plot_wire:
ax = plot_3d_vector(p2-p1, p1, plot_arrow=True, color='k', elev=45, azim=-45, ax=ax)

ax.set_xlabel('x (m)')
ax.set_ylabel('y (m)')
ax.set_zlabel('z (m)')

return B[:,:,:,0], B[:,:,:,1], B[:,:,:,2]
if plot_wire and return_axes:
B[:,:,:,0], B[:,:,:,1], B[:,:,:,2], ax

else:
return B[:,:,:,0], B[:,:,:,1], B[:,:,:,2]


def bfield_from_thin_rectangular_coil(X, Y, Z, a, b, y0, current, plot_wire=False, elev=45, azim=-45, ax=None):
Expand Down Expand Up @@ -191,15 +196,22 @@ def bfield_from_thin_rectangular_coil(X, Y, Z, a, b, y0, current, plot_wire=Fals
p3 = np.array([+a/2, y0, +b/2])
p4 = np.array([-a/2, y0, +b/2])

Bx1, By1, Bz1 = bfield_from_thin_straight_wire(X, Y, Z, p1, p2, current, plot_wire=plot_wire, elev=elev, azim=azim, ax=ax)
Bx1, By1, Bz1 = bfield_from_thin_straight_wire(X, Y, Z, p1, p2, current, plot_wire=False, elev=elev, azim=azim)
Bx2, By2, Bz2 = bfield_from_thin_straight_wire(X, Y, Z, p2, p3, current, plot_wire=False, elev=elev, azim=azim)
Bx3, By3, Bz3 = bfield_from_thin_straight_wire(X, Y, Z, p3, p4, current, plot_wire=False, elev=elev, azim=azim)
Bx4, By4, Bz4 = bfield_from_thin_straight_wire(X, Y, Z, p4, p1, current, plot_wire=False, elev=elev, azim=azim)

if plot_wire:
ax = plt.gca()

Bx2, By2, Bz2 = bfield_from_thin_straight_wire(X, Y, Z, p2, p3, current, plot_wire=plot_wire, elev=elev, azim=azim, ax=ax)
Bx3, By3, Bz3 = bfield_from_thin_straight_wire(X, Y, Z, p3, p4, current, plot_wire=plot_wire, elev=elev, azim=azim, ax=ax)
Bx4, By4, Bz4 = bfield_from_thin_straight_wire(X, Y, Z, p4, p1, current, plot_wire=plot_wire, elev=elev, azim=azim, ax=ax)
ax = plot_3d_vector(p2-p1, p1, plot_arrow=True, color='k', elev=elev, azim=azim, ax=ax)
ax = plot_3d_vector(p3-p2, p2, plot_arrow=True, color='k', elev=elev, azim=azim, ax=ax)
ax = plot_3d_vector(p4-p3, p3, plot_arrow=True, color='k', elev=elev, azim=azim, ax=ax)
ax = plot_3d_vector(p1-p4, p4, plot_arrow=True, color='k', elev=elev, azim=azim, ax=ax)

ax.set_xlabel('x (m)')
ax.set_ylabel('y (m)')
ax.set_zlabel('z (m)')

return (Bx1+Bx2+Bx3+Bx4, By1+By2+By3+By4, Bz1+Bz2+Bz3+Bz4)


Expand Down Expand Up @@ -349,6 +361,8 @@ def get_arc_vectors(h, R, theta,

phi = (np.pi - theta)/2

#print(phi * 180/np.pi)

arc_e1 = np.matmul(rotate_around_e3(phi), np.array([1,0,0]))

assert np.isclose(np.dot(arc_e1, arc_e3), 0)
Expand Down Expand Up @@ -681,6 +695,10 @@ def make_saddle_dipole_corrector_fieldmesh(*,
X, Y, Z = np.meshgrid(xs, ys, zs, indexing='ij')

Bx, By, Bz = bfield_from_thin_saddle_corrector(X, Y, Z, L, R, theta, npts=npts, current=current, plot_wire=plot_wire)

if plot_wire:
ax = plt.gca()
set_axes_equal(ax)

dx = np.diff(xs)[0]
dy = np.diff(ys)[0]
Expand Down Expand Up @@ -869,6 +887,11 @@ def make_dipole_corrector_fieldmesh(*,
# Check that necessary parameters are provided
if R is None or L is None or theta is None:
raise ValueError("Parameters 'R', 'L', and 'theta' must be provided for saddle mode.")

#if plot_wire:
# ax = plt.gca()
# set_axes_equal(ax)

# Call the saddle dipole corrector function
return make_saddle_dipole_corrector_fieldmesh(R=R, L=L, theta=theta, current=current,
xmin=xmin, xmax=xmax, nx=nx,
Expand All @@ -877,8 +900,7 @@ def make_dipole_corrector_fieldmesh(*,
npts=npts, plot_wire=plot_wire)

else:
raise ValueError("Invalid mode. Choose either 'rectangular' or 'saddle'.")

raise ValueError("Invalid mode. Choose either 'rectangular' or 'saddle'.")


def make_thin_straight_wire_fieldmesh(p1, p2, *,
Expand Down

0 comments on commit 9a5734c

Please sign in to comment.