diff --git a/defs.py b/defs.py index 1bfe794..98d32b4 100644 --- a/defs.py +++ b/defs.py @@ -13,23 +13,51 @@ from scipy.integrate import solve_ivp -# TODO put params and interpolation arrays in a class/dict +# TODO put params and interpolation arrays in a class/dict - Partially addressed by SimParams # parameters -strength: int = 1000 # mT magnetization -nres: int = 400 # resolution in z and r - -# field lines -offset: int = 10 # distance from the symmetry and wall -nlines: int = 19 # +1, so 20 in total since rr[1] is added by hand - +@dataclass +class SimParams: + """Holds global simulation parameters. + + Attributes: + strength: Magnetization strength in mT. + nres: Resolution for z and r grids. + offset: Distance from symmetry axis and wall for field line starting points. + nlines: Number of field lines to integrate (excluding the one on the separatrix). + """ + strength: int = 1000 # mT magnetization + nres: int = 400 # resolution in z and r + offset: int = 10 # distance from the symmetry and wall + nlines: int = 19 # +1, so 20 in total since rr[1] is added by hand + +sim_params = SimParams() # definitions def find_separatrix(z: np.ndarray, Bz: np.ndarray) -> Tuple[int, float]: - """find the integration end near the axis where the magnetic field changes topologically, - i.e, where B_z near the axis chnages sign""" + """Finds the separatrix based on the change in sign of Bz near the axis. + + The separatrix is a key topological feature in the magnetic field, + indicating a region where field lines change their connection. This function + identifies its location along the z-axis. + + Args: + z: A 2D numpy array representing the z-coordinates of the grid. + Expected shape (n_r_points, n_z_points). + Bz: A 2D numpy array representing the z-component of the magnetic field + on the grid. Expected shape (n_r_points, n_z_points). + + Returns: + A tuple (index, zsep): + index (int): The index along the z-axis (second dimension of input arrays) + where Bz near the axis first becomes positive. If no such + point is found, `sim_params.nres` is returned. + zsep (float): The z-coordinate corresponding to the found index. If no + such point is found, the maximum z-value from the input + `z` array (at r_index=1) is returned. + """ Bzline: np.ndarray = Bz[1, :] # 1 means one point in r above symmetry axis end_index: np.ndarray = np.where(Bzline > 0)[0] @@ -37,14 +65,32 @@ def find_separatrix(z: np.ndarray, Bz: np.ndarray) -> Tuple[int, float]: zmax: float = z[1, -1] print(zmax) print(z) - return nres, zmax + return sim_params.nres, zmax end_index_val: int = end_index[0] zsep: float = z[1, end_index_val] return end_index_val, zsep # index and z coordinate, where B_z is positive def add_lengths_to_df(df: pd.DataFrame, norm: float = 1) -> pd.DataFrame: - """add the various measures to the dataframe df""" + """Adds several derived length-related measures to a DataFrame. + + The function calculates new columns based on existing 'length', 'mr' (mirror ratio), + and 'r0' (initial radial position of field line) columns. + + Args: + df: Pandas DataFrame to which new columns will be added. + It must contain 'length', 'mr', and 'r0' columns. + norm: A normalization factor to apply to the 'length' column before + deriving other measures. Defaults to 1 (no normalization). + + Returns: + The input DataFrame `df` with added columns: + - 'length': Original 'length' divided by `norm`. + - 'l_mr': `length * sqrt(mr)`. + - 'r0_l': `r0 * length`. + - 'r0_l_mr': `r0 * length * sqrt(mr)`. + - 'r0_l_mr_exp_r0': `r0 * length * sqrt(mr) * exp(-r0)`. + """ df["length"] = df["length"] / norm df["l_mr"] = df["length"] * np.sqrt(df["mr"]) df["r0_l"] = df["r0"] * df["length"] @@ -62,7 +108,25 @@ def parallelism( L: float, p: float = 1, ) -> float: - """calculate average theta = arctan2(br / bz) within (0,0) and (p*L,p*R) at the nres positions""" + """Calculates the average angle theta = arctan2(Br_avg, Bz_avg) in a specified region. + + This angle represents the average field line direction relative to the z-axis + within a rectangular subgrid defined by (0,0) and (p*L, p*R). + + Args: + zz: 1D numpy array of z-coordinates for the grid. + rr: 1D numpy array of r-coordinates for the grid. + Binterp_z: A `scipy.interpolate.RegularGridInterpolator` for Bz component. + Binterp_r: A `scipy.interpolate.RegularGridInterpolator` for Br component. + R: The characteristic radial dimension (e.g., magnet radius). + L: The characteristic axial dimension (e.g., magnet length). + p: A factor (0 to 1) determining the subgrid size relative to R and L. + The calculation is performed within r < p*R and z < p*L. Defaults to 1. + + Returns: + The average angle theta in radians. Returns 0.0 if no grid points + fall within the specified subregion (to avoid division by zero). + """ rp: np.ndarray = rr[np.where(rr < p * R)] zp: np.ndarray = zz[np.where(zz < p * L)] @@ -81,12 +145,32 @@ def parallelism( br_sum += float(br_item) bz_sum += float(bz_item) i += 1 - + + if i == 0: # Avoid division by zero if rp or zp is empty + return 0.0 return float(np.arctan2(br_sum / i, bz_sum / i)) # global angle theta (for that region) def active_volume(z: np.ndarray, r: np.ndarray, L: float, R: float) -> float: - """active volume""" + """Calculates the relative active volume of a field line. + + The active volume is defined by revolving the field line (r(z)) around the + z-axis and is calculated using `np.trapz(r^2, z)`. This volume is then + normalized by the cylindrical volume defined by R and 0.5*L. + The input arrays `z` and `r` may be modified by appending a point + (0.5*L, R) if the field line ends before z = 0.5*L. + + Args: + z: 1D numpy array of z-coordinates along a field line. + r: 1D numpy array of r-coordinates along a field line. + L: Characteristic length of the system (e.g., magnet length). + R: Characteristic radius of the system (e.g., magnet radius). + + Returns: + The relative active volume, normalized by (pi * R^2 * 0.5*L). + Note: pi is implicitly handled as it cancels out if comparing to pi*R^2*L/2. + The function returns (integral(r^2 dz)) / (R^2 * L/2). + """ if z[-1] < 0.5 * L: # add last point to compare with the rectangle z = np.append(z, 0.5 * L) @@ -102,6 +186,33 @@ def active_volume(z: np.ndarray, r: np.ndarray, L: float, R: float) -> float: # basic plot +@dataclass +class RingCalculationOutput: + """Holds the output from the `ring_calculation` function. + + This dataclass serves as a structured container for the various results + computed during the magnetic field characterization of a ring magnet. + + Attributes: + r0_lines: List of initial radial positions (r0) for field line integration. + interp_r_func: Interpolator for the radial component of the B-field (Br). + interp_z_func: Interpolator for the axial component of the B-field (Bz). + pa_val: Calculated parallelism value (average field angle) for p=0.1. + zz_interp: 1D array of z-coordinates used for interpolation. + rr_interp: 1D array of r-coordinates used for interpolation. + zsep: z-coordinate of the separatrix. + magnet: The `magpylib.magnet.CylinderSegment` object representing the ring. + """ + r0_lines: List[float] + interp_r_func: RegularGridInterpolator + interp_z_func: RegularGridInterpolator + pa_val: float + zz_interp: np.ndarray + rr_interp: np.ndarray + zsep: float + magnet: magpy.magnet.CylinderSegment + + def field_plot( z: np.ndarray, rho: np.ndarray, @@ -112,7 +223,21 @@ def field_plot( dR: float, ax: Optional[plt.Axes] = None, ) -> None: - """plot the magnetic field of the ring magnet""" + """Plots the magnetic field streamplot of a ring magnet. + + Displays the magnetic field lines (streamplot) and the magnet's cross-section. + + Args: + z: 1D numpy array of z-coordinates for the grid. + rho: 1D numpy array of radial (rho or r) coordinates for the grid. + Bz: 2D numpy array of the z-component of the magnetic field on the grid. + Br: 2D numpy array of the r-component of the magnetic field on the grid. + R: Inner radius of the ring magnet. + L: Length of the ring magnet. + dR: Thickness (radial extent) of the ring magnet. + ax: Optional `matplotlib.axes.Axes` object to plot on. If None, a new + figure and axes are created. + """ Bamp: np.ndarray = np.linalg.norm(np.array([Bz, Br]), axis=0) # plot @@ -125,15 +250,11 @@ def field_plot( Bz, Br, density=1.1, - # color=Bamp, - # linewidth=np.sqrt(Bamp)*3, cmap="coolwarm", - # broken_streamlines=False ) # figure styling ax.set( - # title="Magnetic field of thin cylinder", ylabel="r / mm", xlabel="z / mm", aspect=1, @@ -148,51 +269,96 @@ def field_plot( # field calculation @dataclass class Ring: + """Represents a ring magnet and its calculated magnetic field properties. + + This class encapsulates the geometric parameters of a ring magnet and the + results obtained from the `ring_calculation` function, such as field + interpolators, separatrix location, and the magnet object itself. + + Attributes: + R: Inner radius of the ring magnet (mm). + L: Length of the ring magnet (mm). + dR: Thickness (radial extent) of the ring magnet (mm). + r0: List of initial radial positions (r0) for field line integration. + bri: Interpolator for the radial component of the B-field (Br). + bzi: Interpolator for the axial component of the B-field (Bz). + pa: Calculated parallelism value (average field angle). + z: 1D array of z-coordinates used for interpolation. + r: 1D array of r-coordinates used for interpolation. + zsep: z-coordinate of the separatrix. + ring: The `magpylib.magnet.CylinderSegment` object representing the ring. + """ R: float L: float dR: float r0: List[float] bri: RegularGridInterpolator bzi: RegularGridInterpolator - pa: List[float] + pa: float z: np.ndarray r: np.ndarray zsep: float ring: magpy.magnet.CylinderSegment def __init__(self, R: float = 20, L: float = 50, dR: float = 2, plt_on: bool = False, ax: Optional[plt.Axes] = None): + """Initializes a Ring object by calculating its magnetic field properties. + + Args: + R: Inner radius of the ring magnet in mm. Defaults to 20. + L: Length of the ring magnet in mm. Defaults to 50. + dR: Thickness (radial extent) of the ring magnet in mm. Defaults to 2. + plt_on: If True, generates a plot of the magnetic field during calculation. + Defaults to False. + ax: Optional `matplotlib.axes.Axes` object to plot on if `plt_on` is True. + If None, a new figure and axes are created for the plot. + """ self.R = R self.L = L self.dR = dR - r0, interp_r, interp_z, pa, zz, rr, zsep, magnet = ring_calculation(R, L, dR, plt_on, ax) - self.r0 = r0 - self.bri = interp_r - self.bzi = interp_z - self.pa = pa - self.z = zz - self.r = rr - self.zsep = zsep - self.ring = magnet + calc_output: RingCalculationOutput = ring_calculation(R, L, dR, plt_on, ax) + self.r0 = calc_output.r0_lines + self.bri = calc_output.interp_r_func + self.bzi = calc_output.interp_z_func + self.pa = calc_output.pa_val + self.z = calc_output.zz_interp + self.r = calc_output.rr_interp + self.zsep = calc_output.zsep + self.ring = calc_output.magnet def ring_calculation( R: float = 20, L: float = 50, dR: float = 2, plt_on: bool = False, ax: Optional[plt.Axes] = None -) -> Tuple[ - List[float], - RegularGridInterpolator, - RegularGridInterpolator, - List[float], - np.ndarray, - np.ndarray, - float, - magpy.magnet.CylinderSegment, -]: - """calculate the field and field line length(s) - return r0, interp_r, interp_z, pa, zz, rr, zsep, magnet""" +) -> RingCalculationOutput: + """Calculates magnetic field properties of a ring magnet. + + This function models a ring magnet, computes its magnetic field on a grid, + finds the separatrix, creates field interpolators, determines field line + starting positions, and calculates field parallelism. Optionally, it can + plot the field. + + Args: + R: Inner radius of the ring magnet in mm. Defaults to 20. + L: Length of the ring magnet in mm. Defaults to 50. + dR: Thickness (radial extent) of the ring magnet in mm. Defaults to 2. + plt_on: If True, generates a plot of the magnetic field. Defaults to False. + ax: Optional `matplotlib.axes.Axes` object to plot on if `plt_on` is True. + If None, a new figure and axes are created for the plot. + + Returns: + A `RingCalculationOutput` object containing the results: + - r0_lines: Initial radial positions for field lines. + - interp_r_func: Interpolator for Br. + - interp_z_func: Interpolator for Bz. + - pa_val: Parallelism value. + - zz_interp: z-coordinates for interpolation. + - rr_interp: r-coordinates for interpolation. + - zsep: z-coordinate of the separatrix. + - magnet: Magpylib magnet object. + """ # generate ring magnet (same as two cylinders, see notebook) magnet: magpy.magnet.CylinderSegment = magpy.magnet.CylinderSegment( - magnetization=(0, 0, strength), + magnetization=(0, 0, sim_params.strength), dimension=(R, R + dR, L, 0, 360), position=(0, 0, 0), ) @@ -204,8 +370,8 @@ def ring_calculation( # pre-compute and plot field of thin_cylinder (faster) # create grid - tr: np.ndarray = np.linspace(0, rmax, nres) - tz: np.ndarray = np.linspace(0, zmax, nres) + tr: np.ndarray = np.linspace(0, rmax, sim_params.nres) + tz: np.ndarray = np.linspace(0, zmax, sim_params.nres) grid: np.ndarray = np.array([[(rh, 0, zh) for zh in tz] for rh in tr]) # compute and plot field of thin_cylinder @@ -217,9 +383,8 @@ def ring_calculation( Br_grid: np.ndarray = np.ascontiguousarray(B[:, :, 0]) # r,z from grid end_index, zsep = find_separatrix(z_grid, Bz_grid) - if end_index == nres: + if end_index == sim_params.nres: print("!!! domain too small !!!") - # zmax = np.ceil(zsep+1) ## plot with switch on/off if plt_on: @@ -233,20 +398,27 @@ def ring_calculation( # define the r0 positions of the integrated field lines # rr[1] is added specifically for showing the separatrix r0_lines: List[float] = [ - rr_interp[i] for i in np.append(1, np.linspace(offset, nres / rmax * R - offset, nlines, dtype=int)) + rr_interp[i] for i in np.append(1, np.linspace(sim_params.offset, sim_params.nres / rmax * R - sim_params.offset, sim_params.nlines, dtype=int)) ] interp_r_func: RegularGridInterpolator = RegularGridInterpolator([rr_interp, zz_interp], Br_grid) interp_z_func: RegularGridInterpolator = RegularGridInterpolator([rr_interp, zz_interp], Bz_grid) # parallelism - pa_vals: List[float] = [] - for p_val in [0.1, 0.3, 0.5, 0.9]: - pa_vals.append(parallelism(zz_interp, rr_interp, interp_z_func, interp_r_func, R, L, p_val)) - # pa = 1 + # Calculate a single parallelism value with p=0.1 + pa_val: float = parallelism(zz_interp, rr_interp, interp_z_func, interp_r_func, R, L, p=0.1) # return for integration and plotting - return r0_lines, interp_r_func, interp_z_func, pa_vals, zz_interp, rr_interp, zsep, magnet + return RingCalculationOutput( + r0_lines=r0_lines, + interp_r_func=interp_r_func, + interp_z_func=interp_z_func, + pa_val=pa_val, + zz_interp=zz_interp, + rr_interp=rr_interp, + zsep=zsep, + magnet=magnet + ) def integration( @@ -257,14 +429,49 @@ def integration( rlines: List[float], interp_r: RegularGridInterpolator, interp_z: RegularGridInterpolator, - parallel: List[float], # This was identified as List[float] in ring_calculation, but seems to be a single float value in usage + parallel: float, zsep: float, plt_on: bool = False, ax: Optional[plt.Axes] = None, ) -> pd.DataFrame: + """Integrates magnetic field lines and calculates their properties. + + For a given set of starting radial positions (`rlines`), this function + traces magnetic field lines until they hit the magnet wall (r=R). + It uses `scipy.integrate.solve_ivp` with the "RK45" method. + The "DOP853" method was previously considered but showed instability for + some integration cases. Properties like field line length, mirror ratio, + and active volume are calculated and stored in a DataFrame. + + Args: + R: Inner radius of the magnet system (wall boundary). + L: Length of the magnet system. + dR: Thickness of the magnet (used for record-keeping in DataFrame). + df: Pandas DataFrame to append results to. + rlines: List of initial radial positions (r0) to start field line integration. + interp_r: Interpolator for the radial component of the B-field (Br). + interp_z: Interpolator for the axial component of the B-field (Bz). + parallel: Parallelism value (average field angle), for record-keeping. + zsep: z-coordinate of the separatrix, used to calculate `zsep_L`. + plt_on: If True, plots the integrated field lines. Defaults to False. + ax: Optional `matplotlib.axes.Axes` object to plot on if `plt_on` is True. + + Returns: + A Pandas DataFrame with the results of the field line integrations. + Each row corresponds to an integrated field line and includes columns: + 'R', 'L', 'dR', 'parallelism', 'zsep_L' (zsep - 0.5*L), 'va' (active volume), + 'r0' (initial radial position), 'mr' (mirror ratio), 'length' (field line length). + """ def field(t: float, y: List[float]) -> List[float]: - """return the interpolated magnetic field (Bz,Br) at point (z,r) - here used as the derivative f(y) = dy/dx in an ODE""" + """Defines the ODE system dy/dt = [-Bz/|B|, -Br/|B|] for field line tracing. + + Args: + t: Time (or path length variable, not explicitly used in field calculation). + y: Current position [z, r]. + + Returns: + The derivatives [dz/ds, dr/ds] proportional to [-Bz, -Br] normalized. + """ z_val, r_val = y bz_val: float = float(interp_z([r_val, z_val], method="linear")[0]) br_val: float = float(interp_r([r_val, z_val], method="linear")[0]) @@ -280,7 +487,15 @@ def field(t: float, y: List[float]) -> List[float]: ] # opposed direction is necessary for integration in positive z direction def hit_wall(t: float, y: List[float]) -> float: - """end criterion for field line integration/ODE""" + """Event function for `solve_ivp` to detect when a field line hits the wall. + + Args: + t: Time (or path length variable). + y: Current position [z, r]. + + Returns: + Value that is zero when the event occurs (r - R + eps = 0). + """ eps: float = 1e-1 # stop a bit before the wall, else motion integration parallel z is possible before the wall, # see scaling of R @@ -312,13 +527,14 @@ def hit_wall(t: float, y: List[float]) -> float: max_step=maxstep, atol=1e-4, rtol=1e-7, - ) # was DOP853 - # TODO some integrations with DOP853 seem to quit after some steps?? - + ) + # Note: The 'DOP853' solver was previously explored but found to be unstable + # for some integration scenarios, leading to premature termination. + # 'RK45' is currently used as a more robust default. + # The TODO comment regarding DOP853 can be removed as this note is added. # calculate the length of each field line s_vals: np.ndarray = np.sqrt(np.power(np.diff(sol.y[0]), 2) + np.power(np.diff(sol.y[1]), 2)) length_val: float = np.sum(s_vals) # not weighted length, that is done later - # print(s.shape, length) # bfield at wall (w) -> mirror ratio mr rw_val: float = float(sol.y[1][-1]) @@ -344,14 +560,10 @@ def hit_wall(t: float, y: List[float]) -> float: ax.plot(sol.y[0], sol.y[1], "red") # put results of each field line into df - # Ensure `parallel` is treated as a scalar if it's a list from `pa_vals` - # Taking the first element as a placeholder, this might need domain specific logic - current_parallel_val = parallel[0] if isinstance(parallel, list) and parallel else parallel - df = pd.concat( [ pd.DataFrame( - [[R, L, dR, current_parallel_val, zsep - 0.5 * L, va_val, r0_val, mr_value, length_val]], + [[R, L, dR, parallel, zsep - 0.5 * L, va_val, r0_val, mr_value, length_val]], # Use parallel directly columns=df.columns, ), df, diff --git a/ring_magnetic_field.ipynb b/ring_magnetic_field.ipynb index 420a8a1..2c8fde5 100644 --- a/ring_magnetic_field.ipynb +++ b/ring_magnetic_field.ipynb @@ -23,6 +23,38 @@ "plt.rc(\"legend\", fontsize=17) # using a size in points" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "parameters_cell", + "metadata": {}, + "outputs": [], + "source": [ + "# --- Parameter Definitions ---\n", + "ngeom_default = 41\n", + "ngeom_secondary = 18 # for R variation scan\n", + "\n", + "# Default geometric parameters for the first set of scans (varying L)\n", + "R_default_scan1 = 20 # Default radius in mm\n", + "dR_default_list_scan1 = [2] # Default dR list in mm\n", + "L_default_linspace_params_scan1 = (10, 90) # (start, end) for L linspace\n", + "\n", + "# Default geometric parameters for the second set of scans (varying R)\n", + "L_default_scan2 = 50 # Default length in mm\n", + "dR_default_list_scan2 = [2] # Default dR list in mm\n", + "R_default_linspace_params_scan2 = (6, 40) # (start, end) for R linspace\n", + "\n", + "# Filenames for saving data and plots\n", + "scan_l_vs_r_filename = \"scan_l_vs_r_data.csv\" # Renamed for clarity\n", + "parallelism_plot_filename = \"parallelism_L_scan.png\"\n", + "zsep_va_plot_filename = \"zsep_va_L_scan.png\"\n", + "lengths_plot_filename = \"lengths_L_scan.png\"\n", + "normed_lengths_plot_filename = \"normed_lengths_L_scan.png\"\n", + "rl_field_plot_filename = \"RL_field_examples.png\"\n", + "scan_r_vs_l_filename = \"scan_r_vs_l_data.csv\"\n", + "lengths_R_varied_plot_filename = \"lengths_R_scan.png\"" + ] + }, { "cell_type": "code", "execution_count": null, @@ -30,7 +62,7 @@ "metadata": {}, "outputs": [], "source": [ - "ngeom = 41" + "ngeom = ngeom_default" ] }, { @@ -50,20 +82,21 @@ " columns=[\"R\", \"L\", \"dR\", \"parallelism\", \"zsep_L\", \"va\", \"r0\", \"mr\", \"length\"]\n", ")\n", "\n", - "# fixed R\n", - "R = 20 # radius in mm\n", - "\n", - "# dR = 2 # width of the magnet in mm\n", - "# for dR in np.linspace(2, 10, 5): # length in mm\n", - "for dR in [2]: # length in mm\n", - " for L in np.linspace(10, 90, ngeom): # length in mm\n", - " # for L in np.linspace(40, 60, 11): # length in mm\n", - " # for L in [10, 50, 90]: # length in mm\n", - " rlines, Brinterp, Bzinterp, parallel, zz, rr, zsep, ring = ring_calculation(\n", - " R, L, dR, plt_on\n", - " ) # calculate the field\n", + "# Use parameterized values for the first scan (varying L)\n", + "R_scan1 = R_default_scan1\n", + "\n", + "for dR_scan1 in dR_default_list_scan1:\n", + " for L_scan1 in np.linspace(L_default_linspace_params_scan1[0], L_default_linspace_params_scan1[1], ngeom): # length in mm\n", + " # Comments for previous loop variations removed for clarity\n", + " calc_output = ring_calculation(R_scan1, L_scan1, dR_scan1, plt_on) # calculate the field\n", + " rlines = calc_output.r0_lines\n", + " Brinterp = calc_output.interp_r_func\n", + " Bzinterp = calc_output.interp_z_func\n", + " parallel = calc_output.pa_val\n", + " zsep = calc_output.zsep\n", + " # zz, rr, and ring (magnet object) from calc_output are not directly used here for integration input\n", " df = integration(\n", - " R, L, dR, df, rlines, Brinterp, Bzinterp, parallel, zsep, plt_on\n", + " R_scan1, L_scan1, dR_scan1, df, rlines, Brinterp, Bzinterp, parallel, zsep, plt_on\n", " ) # generate parameter df" ] }, @@ -123,7 +156,7 @@ "\n", "fig.tight_layout()\n", "plt.grid()\n", - "plt.savefig(\"parallelism.png\", dpi=300)" + "plt.savefig(parallelism_plot_filename, dpi=300)" ] }, { @@ -144,13 +177,18 @@ "ax.set_ylabel(r\"$\\widetilde{z}$ / mm\")\n", "ax.grid()\n", "\n", - "# select only the uppermost field line from the df\n", - "nn = 20 # == nline+1 # TODO parameter from class\n", + "# select only the uppermost field line from the df for each geometry\n", + "# The DataFrame `df` stores results for `ngeom` different geometries (L values).\n", + "# For each geometry, `sim_params.nlines + 1` field lines are integrated.\n", + "# `nn` represents the number of field lines per geometry.\n", + "nn = sim_params.nlines + 1 # Address TODO: Use parameter from sim_params\n", "line = []\n", - "for rrr in range(ngeom): # r range\n", - " line.append(0 + (nn) * rrr)\n", + "# This loop creates a list of indices to select the first field line (index 0 within its group)\n", + "# for each of the `ngeom` geometries.\n", + "for rrr in range(ngeom): # Iterate through each geometry scanned\n", + " line.append(0 + (nn) * rrr) # Index of the first field line for the rrr-th geometry\n", "\n", - "# plt.plot(df.L[line], 1 - (df.va[line] / (R*R*L/2)))\n", + "# plt.plot(df.L[line], 1 - (df.va[line] / (R_default_scan1*R_default_scan1*df.L[line]/2))) # Example with R_default_scan1\n", "lns2 = ax2.plot(df.L[line], df.va[line], \"r\", label=r\"$\\widetilde{V}$\") # /(R*R*0.5*L))\n", "# ax2.set_ylabel(\"$\\widetilde{V}$ = $V_{active}$ / $V_{total}$\")\n", "ax2.set_ylabel(r\"$\\widetilde{V}$\")\n", @@ -160,7 +198,7 @@ "ax.legend(lns, labs)\n", "plt.tight_layout()\n", "\n", - "plt.savefig(\"zsep_va.png\", dpi=300, bbox_inches=\"tight\")" + "plt.savefig(zsep_va_plot_filename, dpi=300, bbox_inches=\"tight\")" ] }, { @@ -170,7 +208,7 @@ "metadata": {}, "outputs": [], "source": [ - "df.to_csv(\"scan_l2r_mr.csv\")\n", + "df.to_csv(scan_l_vs_r_filename)\n", "\n", "df = add_lengths_to_df(df) # add lengths to df\n", "\n", @@ -212,20 +250,20 @@ " pb = df.R[0] * np.sqrt(6) # parallelB value for R = 20mm\n", "\n", " fig, ax = plt.subplots(3, 1, figsize=(8, 8), sharex=True)\n", - " nlines = 20\n", - " rhos = [df.r0[2], df.r0[10], df.r0[17]]\n", + " nlines_for_avg = sim_params.nlines + 1 # Number of lines used for averaging in dfsum\n", + " rhos = [df.r0[2], df.r0[10], df.r0[17]] # Example r0 values for specific line plots\n", " label = \"average\"\n", "\n", - " # plot all lengths for all parameters\n", - " ax[0].plot(dfsum.L, dfsum.length / nlines, label=label)\n", + " # plot all lengths for all parameters, averaged over field lines\n", + " ax[0].plot(dfsum.L, dfsum.length / nlines_for_avg, label=label)\n", " ax[0].axvline(x=pb, color=\"tab:olive\")\n", " ax[0].set_ylabel(r\"$\\ell$ / a.u.\")\n", - " ax[1].plot(dfsum.L, dfsum.mr / nlines, label=label)\n", + " ax[1].plot(dfsum.L, dfsum.mr / nlines_for_avg, label=label)\n", " ax[1].axvline(x=pb, color=\"tab:olive\")\n", " ax[1].set_ylabel(r\"$R_m$\")\n", "\n", - " # exponentially weighted lengths\n", - " ax[2].plot(dfsum.L, dfsum[\"r0_l_mr_exp_r0\"] / nlines, label=label)\n", + " # exponentially weighted lengths, averaged over field lines\n", + " ax[2].plot(dfsum.L, dfsum[\"r0_l_mr_exp_r0\"] / nlines_for_avg, label=label)\n", " ax[2].axvline(x=pb, color=\"tab:olive\")\n", " ax[2].set_ylabel(r\"$\\tau$ / a.u.\")\n", " # $\\hat{\\ell} = \\ell \\cdot \\sqrt{R_m} \\cdot \\rho_0 \\cdot e^{-\\rho_0}$\n", @@ -260,7 +298,7 @@ "\n", "\n", "plot_lengths(df, dfsum)\n", - "plt.savefig(\"lengths.png\", dpi=300)" + "plt.savefig(lengths_plot_filename, dpi=300)" ] }, { @@ -285,7 +323,7 @@ "ndfsum\n", "\n", "plot_lengths(ndf, ndfsum)\n", - "plt.savefig(\"normed_lengths.png\", dpi=300)" + "plt.savefig(normed_lengths_plot_filename, dpi=300)" ] }, { @@ -335,19 +373,22 @@ "axes = [ax1, ax2, ax3]\n", "label = [\"(a)\", \"(b)\", \"(c)\"]\n", "\n", - "# fixed R\n", - "R = 20 # radius in mm\n", - "for dR in [2]: # length in mm\n", - " for i, L in enumerate([10, 50, 90]): # length in mm\n", - " m = Ring(R, L, dR, plt_on, axes[i])\n", + "# fixed R, example values for L and dR for plotting field lines\n", + "R_example = R_default_scan1 # Use the same R as the first scan for consistency\n", + "dR_example_list = dR_default_list_scan1 # Use the same dR as the first scan\n", + "L_example_list = [10, 50, 90] # Example L values for specific field plots\n", + "\n", + "for dR_ex in dR_example_list:\n", + " for i, L_ex in enumerate(L_example_list):\n", + " m = Ring(R_example, L_ex, dR_ex, plt_on, axes[i])\n", " df = integration(\n", " m.R, m.L, m.dR, df, m.r0, m.bri, m.bzi, m.pa, m.zsep, plt_on, axes[i]\n", - " ) # generate parameter df\n", - " axes[i].plot(m.zsep, 0, \"k*\", markersize=14) # , color=\"tab:orange\")\n", + " ) # generate parameter df (note: df is being appended globally, might be unintended for this example section)\n", + " axes[i].plot(m.zsep, 0, \"k*\", markersize=14)\n", " axes[i].set_title(label[i], fontfamily=\"serif\", loc=\"left\", fontsize=\"medium\")\n", "\n", "\n", - "plt.savefig(\"RL_field.png\", dpi=300, bbox_inches=\"tight\")" + "plt.savefig(rl_field_plot_filename, dpi=300, bbox_inches=\"tight\")" ] }, { @@ -367,8 +408,8 @@ }, "outputs": [], "source": [ - "### magnetic geometry\n", - "ngeom = 18\n", + "### magnetic geometry - Scan 2: Varying R\n", + "ngeom_R_scan = ngeom_secondary # Use parameterized number of geometry points for R scan\n", "\n", "# to plot or not to plot (scans)\n", "plt_on = False\n", @@ -378,18 +419,21 @@ " columns=[\"R\", \"L\", \"dR\", \"parallelism\", \"zsep_L\", \"va\", \"r0\", \"mr\", \"length\"]\n", ")\n", "\n", - "# for all R,L\n", - "L = 50 # radius in mm\n", - "\n", - "# for dR in np.linspace(2, 10, 5): # width of the magnet in mm\n", - "for dR in [2]: # length in mm\n", - " for R in np.linspace(6, 40, ngeom): # length in mm\n", - " # for R in [6, 22, 40]: # length in mm\n", - " rlines, Brinterp, Bzinterp, parallel, zz, rr, zsep, ring = ring_calculation(\n", - " R, L, dR, plt_on\n", - " ) # calculate the field\n", + "# Use parameterized values for the second scan (varying R)\n", + "L_scan2 = L_default_scan2\n", + "\n", + "for dR_scan2 in dR_default_list_scan2:\n", + " for R_scan2 in np.linspace(R_default_linspace_params_scan2[0], R_default_linspace_params_scan2[1], ngeom_R_scan):\n", + " # Comments for previous loop variations removed for clarity\n", + " calc_output = ring_calculation(R_scan2, L_scan2, dR_scan2, plt_on) # calculate the field\n", + " rlines = calc_output.r0_lines\n", + " Brinterp = calc_output.interp_r_func\n", + " Bzinterp = calc_output.interp_z_func\n", + " parallel = calc_output.pa_val\n", + " zsep = calc_output.zsep\n", + " # zz, rr, and ring (magnet object) from calc_output are not directly used here for integration input\n", " dfr = integration(\n", - " R, L, dR, dfr, rlines, Brinterp, Bzinterp, parallel, zsep, plt_on\n", + " R_scan2, L_scan2, dR_scan2, dfr, rlines, Brinterp, Bzinterp, parallel, zsep, plt_on\n", " ) # generate parameter df" ] }, @@ -416,7 +460,8 @@ "\n", "# dfr\n", "dfrsum = dfr.groupby([\"R\", \"L\"]).sum().reset_index()\n", - "dfrsum" + "dfrsum.to_csv(scan_r_vs_l_filename) # Save the summarized data for this scan\n", + "dfrsum # Display the dataframe" ] }, { @@ -435,47 +480,53 @@ "\n", " fig, ax = plt.subplots(3, 1, figsize=(8, 8), sharex=True)\n", "\n", - " nlines = 20\n", - " # rhos = [df.r0[2], df.r0[10],df.r0[17]] # fixed positions\n", - " rhos = [2, 10, 17] # relative positions\n", - " rlabel = [\"0.85\", \"0.5\", \"0.1\"]\n", + " nlines_for_avg_R_scan = sim_params.nlines + 1 # Number of lines used for averaging in dfsum\n", + " # rhos defines indices of specific field lines to plot individually.\n", + " # These are 0-indexed within each group of 'nlines_for_avg_R_scan' lines.\n", + " rhos_indices = [2, 10, 17] # Example indices of r0 within each geometry's set of field lines\n", + " rlabel = [\"0.85\", \"0.5\", \"0.1\"] # Labels corresponding to relative r0/R positions\n", "\n", " label = \"average\"\n", "\n", - " # plot all lengths for all parameters\n", - " ax[0].plot(dfsum.R, dfsum.length / nlines, label=label)\n", + " # Plotting average values from dfsum\n", + " ax[0].plot(dfsum.R, dfsum.length / nlines_for_avg_R_scan, label=label)\n", " ax[0].axvline(x=pb, color=\"tab:olive\")\n", " ax[0].set_ylabel(r\"$\\ell$ / a.u.\")\n", - " ax[1].plot(dfsum.R, dfsum.mr / nlines, label=label)\n", + " ax[1].plot(dfsum.R, dfsum.mr / nlines_for_avg_R_scan, label=label)\n", " ax[1].axvline(x=pb, color=\"tab:olive\")\n", " ax[1].set_ylabel(\"$R_m$\")\n", "\n", - " # exponentially weighted lengths\n", - " ax[2].plot(dfsum.R, dfsum[\"r0_l_mr_exp_r0\"] / nlines, label=label)\n", + " # exponentially weighted lengths, averaged\n", + " ax[2].plot(dfsum.R, dfsum[\"r0_l_mr_exp_r0\"] / nlines_for_avg_R_scan, label=label)\n", " ax[2].axvline(x=pb, color=\"tab:olive\")\n", " ax[2].set_ylabel(r\"$\\tau$ / a.u.\")\n", "\n", - " # for selected lines\n", + " # Plotting specific field lines from the original dfr DataFrame\n", " ls = [\"-\", \":\", \"--\"]\n", - " for i, l in enumerate(rhos):\n", - " line = []\n", - " for rrr in range(ngeom): # number of R runs TODO adapt!! 18\n", - " line.append(l + (nlines) * rrr)\n", - " dff = df.iloc[line] # ???\n", - " # plot all lengths for all parameters\n", - "\n", + " for i, line_index_in_group in enumerate(rhos_indices):\n", + " # Construct a list of global indices for the specific field line across all R scans\n", + " # Each R scan (geometry) has 'nlines_for_avg_R_scan' rows in dfr.\n", + " # 'rrr' iterates through each R scan.\n", + " # 'line_index_in_group' is the 0-based index of the desired field line within that R scan's group.\n", + " line_global_indices = []\n", + " for rrr in range(ngeom_R_scan): # ngeom_R_scan is the number of R values scanned\n", + " line_global_indices.append(line_index_in_group + (nlines_for_avg_R_scan * rrr))\n", + " \n", + " # Select rows from dfr using these global indices\n", + " dff = df.iloc[line_global_indices]\n", + " \n", " ax[0].plot(\n", - " dff.R, dff.length, label=r\"$r_0/R = $\" + str(rlabel[i]), linestyle=ls[i]\n", + " dff.R, dff.length, label=r\"$r_0/R \\approx $ \" + str(rlabel[i]), linestyle=ls[i]\n", " )\n", - " ax[1].plot(dff.R, dff.mr, label=r\"$r_0/R = $\" + str(rlabel[i]), linestyle=ls[i])\n", - " ax[2].plot(dff.R, dff[\"r0_l_mr_exp_r0\"], label=str(rlabel[i]), linestyle=ls[i])\n", + " ax[1].plot(dff.R, dff.mr, label=r\"$r_0/R \\approx $ \" + str(rlabel[i]), linestyle=ls[i])\n", + " ax[2].plot(dff.R, dff[\"r0_l_mr_exp_r0\"], label=r\"$r_0/R \\approx $ \" + str(rlabel[i]), linestyle=ls[i])\n", "\n", " ax[0].legend()\n", " ax[2].set_xlabel(\"R / mm\")\n", "\n", "\n", "plot_lengths_R(dfr, dfrsum)\n", - "plt.savefig(\"lengths_R_varied.png\", dpi=300)" + "plt.savefig(lengths_R_varied_plot_filename, dpi=300)" ] }, { diff --git a/taylor_expansion.ipynb b/taylor_expansion.ipynb index cdf6d99..f10e210 100644 --- a/taylor_expansion.ipynb +++ b/taylor_expansion.ipynb @@ -1,5 +1,17 @@ { "cells": [ + { + "cell_type": "markdown", + "id": "intro_markdown", + "metadata": {}, + "source": [ + "## Symbolic Taylor Expansion of Magnetic Potential $\\Phi_m$\n", + "\n", + "This notebook performs a symbolic Taylor expansion of the magnetic scalar potential $\\Phi_m$ for a specific magnet configuration. The primary goals are:\n", + "1. To analyze the behavior of $\\Phi_m$ around the point $z=0$ (typically the center of the magnet system).\n", + "2. To find conditions (relationships between magnet dimensions $R$ and $L$) under which specific terms in the Taylor expansion vanish. This can be relevant for designing magnets with desired field properties, such as high homogeneity in a region (where lower-order gradient terms are minimized)." + ] + }, { "cell_type": "code", "execution_count": 1, @@ -9,7 +21,9 @@ "source": [ "## Taylor expansion of Phi_m\n", "z = var('z')\n", - "var('R,L')\n", + "var('R,L') # Define R (radius) and L (length) as symbolic variables\n", + "# Define the magnetic scalar potential Phi_m(z)\n", + "# This expression typically represents the potential from two loops or ends of a solenoid\n", "Phi_m(z) = R * (1 / (R^2 + (z+L/2)^2)^(1/2) - 1 / (R^2 + (z-L/2)^2)^(1/2))" ] }, @@ -21,6 +35,16 @@ "outputs": [], "source": [] }, + { + "cell_type": "markdown", + "id": "taylor_expansion_markdown", + "metadata": {}, + "source": [ + "### Compute Taylor Expansion\n", + "\n", + "The following step computes the Taylor series expansion of the defined magnetic potential $\\Phi_m(z)$ with respect to the variable $z$, centered around $z=0$. The expansion is carried out up to the $n$-th order (inclusive). The `.full_simplify()` method is used to simplify the resulting expression." + ] + }, { "cell_type": "code", "execution_count": 2, @@ -92,10 +116,26 @@ ], "source": [ "# z^3 term vanishes if the polynomial is zero\n", + "# The expression L^7*R + 2*L^5*R^3 - 32*L^3*R^5 - 96*L*R^7 is the coefficient of the z^3 term \n", + "# (multiplied by some constants) in the Taylor expansion of Phi_m.\n", "b = solve(L^7*R + 2*L^5*R^3 - 32*L^3*R^5 - 96*L*R^7,L)\n", "b" ] }, + { + "cell_type": "markdown", + "id": "solve_explanation_markdown", + "metadata": {}, + "source": [ + "### Solving for Vanishing $z^3$ Term Coefficient\n", + "\n", + "The previous cell solved for the conditions on $L$ (length) and $R$ (radius) that make the coefficient of the $z^3$ term in the Taylor expansion of $\\Phi_m$ equal to zero. \n", + "\n", + "The expression `L^7*R + 2*L^5*R^3 - 32*L^3*R^5 - 96*L*R^7` is proportional to the coefficient of the $z^3$ term. Setting this to zero and solving for $L$ in terms of $R$ (or vice versa) gives the geometric conditions for which the $z^3$ gradient of the potential (and thus certain components of the magnetic field related to $d^3\\Phi_m/dz^3$) vanishes at $z=0$.\n", + "\n", + "A key result, $L = \\sqrt{6}R$ (or $L/R = \\sqrt{6} \\approx 2.449$), is known as a condition for creating a Helmholtz-like coil configuration with improved field homogeneity near the center ($z=0$). When this condition is met, the $z^3$ term (and other odd-powered terms due to symmetry) in the expansion of $\\Phi_m$ (and consequently $B_z = -d\\Phi_m/dz$) vanishes, leading to a flatter field profile for $B_z$ around $z=0$. This is significant for applications requiring a uniform magnetic field." + ] + }, { "cell_type": "code", "execution_count": 5, diff --git a/test_defs.py b/test_defs.py index fc9722d..8389a45 100644 --- a/test_defs.py +++ b/test_defs.py @@ -8,8 +8,7 @@ from magpylib.magnet import CylinderSegment from defs import ( find_separatrix, - nres, - nlines, + sim_params, # Updated import ring_calculation, Ring, parallelism, @@ -33,8 +32,8 @@ def test_find_separatrix_found() -> None: def test_find_separatrix_not_found() -> None: z: np.ndarray = np.array([[0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4, 5]]) Bz: np.ndarray = np.array([[-1, -1, -0.5, -0.1, -0.2, -0.3], [-1, -1, -0.5, -0.1, -0.2, -0.3]]) - # nres is imported from defs - expected_index: int = nres # Accessing global nres from defs.py + # sim_params.nres is used instead of nres + expected_index: int = sim_params.nres # Accessing nres from sim_params expected_zsep: float = z[1, -1] # 5.0 index, zsep = find_separatrix(z, Bz) assert index == expected_index @@ -61,71 +60,85 @@ def test_find_separatrix_at_end() -> None: assert zsep == expected_zsep +from defs import ( + find_separatrix, + sim_params, # Updated import + ring_calculation, RingCalculationOutput, # Added RingCalculationOutput + Ring, + parallelism, + active_volume, + add_lengths_to_df, # Import necessary items from defs + integration, # Import the integration function +) +from pandas.testing import assert_frame_equal + + # Tests for ring_calculation def test_ring_calculation_basic_run_and_types() -> None: R: float = 20 L: float = 50 dR: float = 2 - r0, interp_r, interp_z, pa, zz, rr, zsep, magnet = ring_calculation( + # Updated to expect RingCalculationOutput + result: RingCalculationOutput = ring_calculation( R, L, dR, plt_on=False ) - assert isinstance(r0, list) - assert all(isinstance(x, float) for x in r0) - assert isinstance(interp_r, RegularGridInterpolator) - assert isinstance(interp_z, RegularGridInterpolator) - assert isinstance(pa, list) - # assert all(isinstance(x, float) for x in pa) # This assertion fails - assert isinstance(zz, np.ndarray) - assert isinstance(rr, np.ndarray) - assert isinstance(zsep, float) - assert isinstance(magnet, CylinderSegment) + assert isinstance(result, RingCalculationOutput) + assert isinstance(result.r0_lines, list) + assert all(isinstance(x, float) for x in result.r0_lines) + assert isinstance(result.interp_r_func, RegularGridInterpolator) + assert isinstance(result.interp_z_func, RegularGridInterpolator) + assert isinstance(result.pa_val, float) + assert isinstance(result.zz_interp, np.ndarray) + assert isinstance(result.rr_interp, np.ndarray) + assert isinstance(result.zsep, float) + assert isinstance(result.magnet, CylinderSegment) # Not None checks - assert r0 is not None - assert interp_r is not None - assert interp_z is not None - assert pa is not None - assert zz is not None - assert rr is not None - assert zsep is not None - assert magnet is not None + assert result.r0_lines is not None + assert result.interp_r_func is not None + assert result.interp_z_func is not None + assert result.pa_val is not None + assert result.zz_interp is not None + assert result.rr_interp is not None + assert result.zsep is not None + assert result.magnet is not None def test_ring_calculation_array_shapes() -> None: R: float = 20 L: float = 50 dR: float = 2 - r0, _, _, pa, zz, rr, _, _ = ring_calculation(R, L, dR) # plt_on defaults to False + result: RingCalculationOutput = ring_calculation(R, L, dR) # plt_on defaults to False - assert zz.ndim == 1 - assert zz.shape[0] == nres - assert rr.ndim == 1 - assert rr.shape[0] == nres - assert len(r0) == nlines + 1 # nlines is from defs.py - assert len(pa) == 4 + assert result.zz_interp.ndim == 1 + assert result.zz_interp.shape[0] == sim_params.nres + assert result.rr_interp.ndim == 1 + assert result.rr_interp.shape[0] == sim_params.nres + assert len(result.r0_lines) == sim_params.nlines + 1 # nlines from sim_params + assert isinstance(result.pa_val, float) def test_ring_calculation_r0_content() -> None: R: float = 20 L: float = 50 dR: float = 2 - r0, _, _, _, _, rr, _, _ = ring_calculation(R, L, dR) - assert r0[0] == rr[1] - min_rr_for_r0: float = rr[1] - max_rr_for_r0: float = rr[-2] - if len(r0) > 1: - assert all(min_rr_for_r0 <= x <= max_rr_for_r0 for x in r0[1:]) + result: RingCalculationOutput = ring_calculation(R, L, dR) + assert result.r0_lines[0] == result.rr_interp[1] + min_rr_for_r0: float = result.rr_interp[1] + max_rr_for_r0: float = result.rr_interp[-2] # Assuming nres is large enough + if len(result.r0_lines) > 1: + assert all(min_rr_for_r0 <= x <= max_rr_for_r0 for x in result.r0_lines[1:]) def test_ring_calculation_zsep_plausibility() -> None: R: float = 20 L: float = 50 dR: float = 2 - _, _, _, _, _, _, zsep, _ = ring_calculation(R, L, dR) - assert zsep > 0 + result: RingCalculationOutput = ring_calculation(R, L, dR) + assert result.zsep > 0 zmax_calc: float = 0.5 * L + dR + 0.5 * R - assert zsep <= zmax_calc + assert result.zsep <= zmax_calc def test_ring_calculation_plot_on_true(mocker: MockerFixture) -> None: @@ -153,40 +166,52 @@ def test_ring_instantiation_and_attributes() -> None: assert all(isinstance(x, float) for x in ring_instance.r0) assert isinstance(ring_instance.bri, RegularGridInterpolator) assert isinstance(ring_instance.bzi, RegularGridInterpolator) - assert isinstance(ring_instance.pa, list) - # if ring_instance.pa: # This assertion fails + assert isinstance(ring_instance.pa, float) # Changed from list + # # if ring_instance.pa: # This assertion fails - Removed # assert all(isinstance(x, float) for x in ring_instance.pa) assert isinstance(ring_instance.z, np.ndarray) assert isinstance(ring_instance.r, np.ndarray) assert isinstance(ring_instance.zsep, float) assert isinstance(ring_instance.ring, CylinderSegment) - assert len(ring_instance.r0) == nlines + 1 + assert len(ring_instance.r0) == sim_params.nlines + 1 assert ring_instance.z.ndim == 1 - assert ring_instance.z.shape[0] == nres + assert ring_instance.z.shape[0] == sim_params.nres assert ring_instance.r.ndim == 1 - assert ring_instance.r.shape[0] == nres - - -# def test_ring_plt_on_ax_passing(mocker: MockerFixture) -> None: -# R_val: float = 20 -# L_val: float = 50 -# dR_val: float = 2 -# mock_ax_object: Mock = mocker.Mock() -# mocked_ring_calc: Mock = mocker.patch('defs.ring_calculation') -# dummy_r0: List[float] = [1.0] * (nlines + 1) -# dummy_interp: Mock = mocker.Mock(spec=RegularGridInterpolator) -# dummy_interp.side_effect = lambda x: np.zeros((x.shape[0],1)) -# dummy_pa: List[float] = [0.1, 0.2, 0.3, 0.4] -# dummy_zz: np.ndarray = np.linspace(-10, 10, nres) -# dummy_rr: np.ndarray = np.linspace(0, 5, nres) -# dummy_zsep: float = 5.0 -# dummy_magnet: Mock = mocker.Mock(spec=CylinderSegment) -# mocked_ring_calc.return_value = (dummy_r0, dummy_interp, dummy_interp, dummy_pa, dummy_zz, dummy_rr, dummy_zsep, dummy_magnet) -# Ring(R=R_val, L=L_val, dR=dR_val, plt_on=True, ax=mock_ax_object) -# mocked_ring_calc.assert_called_once_with( -# R=R_val, L=L_val, dR=dR_val, -# plt_on=True, ax=mock_ax_object, -# param_dict=None, filename=None, Bz0=None, curr=None) + assert ring_instance.r.shape[0] == sim_params.nres + + +def test_ring_plt_on_ax_passing(mocker: MockerFixture) -> None: + R_val: float = 20 + L_val: float = 50 + dR_val: float = 2 + mock_ax_object: Mock = mocker.Mock() + mocked_ring_calc: Mock = mocker.patch('defs.ring_calculation') + # Use sim_params.nlines + dummy_r0: List[float] = [1.0] * (sim_params.nlines + 1) + dummy_interp: Mock = mocker.Mock(spec=RegularGridInterpolator) + # Ensure the mock interpolator returns a 1-element array or scalar, then extract with [0] or .item() + # The original code in integration() uses [0] for the interpolator result. + dummy_interp.return_value = np.array([0.0]) + dummy_pa: float = 0.1 + dummy_zz: np.ndarray = np.linspace(-10, 10, sim_params.nres) + dummy_rr: np.ndarray = np.linspace(0, 5, sim_params.nres) + dummy_zsep: float = 5.0 + dummy_magnet: Mock = mocker.Mock(spec=CylinderSegment) + # Update mocked_ring_calc.return_value to be a RingCalculationOutput instance + mocked_ring_calc.return_value = RingCalculationOutput( + r0_lines=dummy_r0, + interp_r_func=dummy_interp, + interp_z_func=dummy_interp, + pa_val=dummy_pa, + zz_interp=dummy_zz, + rr_interp=dummy_rr, + zsep=dummy_zsep, + magnet=dummy_magnet + ) + Ring(R=R_val, L=L_val, dR=dR_val, plt_on=True, ax=mock_ax_object) + mocked_ring_calc.assert_called_once_with( + R=R_val, L=L_val, dR=dR_val, + plt_on=True, ax=mock_ax_object) # Removed extra params # Tests for parallelism function @@ -244,32 +269,85 @@ def test_parallelism_zero_field(mocker: MockerFixture) -> None: ) -# def test_parallelism_p_less_than_one(mocker: MockerFixture) -> None: -# zz:np.ndarray=np.array([0,1,2,3,3.5,4]); rr:np.ndarray=np.array([0,0.5,1,1.5]); R_val:float=2;L_val:float=4;p_val:float=0.5 -# mock_b_interp_r:Mock=mocker.Mock(return_value=1.0); mock_b_interp_z:Mock=mocker.Mock(return_value=1.0) -# assert np.isclose(parallelism(zz,rr,mock_b_interp_z,mock_b_interp_r,R_val,L_val,p_val), np.pi/4) -# assert mock_b_interp_r.call_count == 3; assert mock_b_interp_z.call_count == 3 - -# def test_parallelism_empty_rp_zp(mocker: MockerFixture) -> None: -# zz:np.ndarray=np.array([0,1,2,3]); rr:np.ndarray=np.array([0,0.5,1]); R_val:float=0.1;L_val:float=0.1;p_val:float=1 -# mock_b_interp_r:Mock=mocker.Mock(return_value=1.0); mock_b_interp_z:Mock=mocker.Mock(return_value=1.0) -# assert np.isclose(parallelism(zz,rr,mock_b_interp_z,mock_b_interp_r,R_val,L_val,p_val),0.0) -# assert mock_b_interp_r.call_count == 0; assert mock_b_interp_z.call_count == 0 - -# def test_parallelism_p_zero(mocker: MockerFixture) -> None: -# zz:np.ndarray=np.array([0,1,2,3]); rr:np.ndarray=np.array([0,0.5,1]); R_val:float=1;L_val:float=1;p_val:float=0 -# mock_b_interp_r:Mock=mocker.Mock(return_value=1.0); mock_b_interp_z:Mock=mocker.Mock(return_value=1.0) -# assert np.isclose(parallelism(zz,rr,mock_b_interp_z,mock_b_interp_r,R_val,L_val,p_val),0.0) -# assert mock_b_interp_r.call_count == 0; assert mock_b_interp_z.call_count == 0 - -# # Tests for active_volume function -# def test_active_volume_simple_no_modification() -> None: -# z:np.ndarray=np.array([0,1,2]); r:np.ndarray=np.array([1,1,1]); L:float=4;R_param:float=1 -# assert np.isclose(active_volume(z,r,L,R_param),1.0) - -# def test_active_volume_simple_varying_r_no_modification() -> None: -# z:np.ndarray=np.array([0,1,np.sqrt(2)]); r:np.ndarray=np.array([0,1,np.sqrt(2)]); L:float=4;R_param:float=np.sqrt(2) # z fixed from example -# assert np.isclose(active_volume(z,r,L,R_param),0.5) +def test_parallelism_p_less_than_one(mocker: MockerFixture) -> None: + zz:np.ndarray=np.array([0,1,2,3,3.5,4]); rr:np.ndarray=np.array([0,0.5,1,1.5]); R_val:float=2;L_val:float=4;p_val:float=0.5 + # Expect rp = [0, 0.5], zp = [0, 1] because p*R = 1, p*L = 2 + # Points for interpolation will be (0,0), (0,1), (0.5,0), (0.5,1) + # So 4 calls for each interpolator if my reasoning of the loops in parallelism is correct. + # The original test expected 3 calls. Let's check the implementation of parallelism. + # rp = rr[np.where(rr < p*R)] -> rr[rr < 1] -> [0, 0.5] + # zp = zz[np.where(zz < p*L)] -> zz[zz < 2] -> [0, 1] + # Loop is for r_val in rp: for z_val in zp: + # (0,0), (0,1), (0.5,0), (0.5,1) -> 4 points, so 4 calls. + mock_b_interp_r:Mock=mocker.Mock(return_value=1.0); mock_b_interp_z:Mock=mocker.Mock(return_value=1.0) + assert np.isclose(parallelism(zz,rr,mock_b_interp_z,mock_b_interp_r,R_val,L_val,p_val), np.pi/4) + assert mock_b_interp_r.call_count == 4; assert mock_b_interp_z.call_count == 4 + +def test_parallelism_empty_rp_zp(mocker: MockerFixture) -> None: + zz:np.ndarray=np.array([0,1,2,3]); rr:np.ndarray=np.array([0,0.5,1]); R_val:float=0.1;L_val:float=0.1;p_val:float=1 + # rp = rr[rr < 0.1] -> [0] + # zp = zz[zz < 0.1] -> [0] + # Loop will run for (0,0) once. + mock_b_interp_r:Mock=mocker.Mock(return_value=1.0); mock_b_interp_z:Mock=mocker.Mock(return_value=1.0) + # If br_sum = 1, bz_sum = 1, then arctan2(1,1) = pi/4. + # If the loop runs once, i=1. + # If R_val=0.1, L_val=0.1, p_val=1. rp = rr[rr<0.1] -> [0]. zp = zz[zz<0.1] -> [0]. + # Calls for (0,0). So call_count should be 1. + # The original test expected 0.0 for the result and 0 calls. + # This happens if rp or zp is empty. + # rr < 0.1 => [0.0] -> not empty + # Let's make R_val and L_val small enough such that p*R and p*L are < 0 (not possible as rr[0]=0) + # Or make p_val = 0. This is another test case. + # To make rp/zp empty, p*R or p*L must be <= 0, assuming rr and zz start at 0. + # If R_val = 0.001, p_val = 1. rr < 0.001 -> [0.0]. Not empty. + # The only way to get empty rp or zp is if the first element of rr or zz is already greater than p*R or p*L. + # Example: rr = [0.5, 1], R_val = 0.4, p_val = 1. Then rr[rr < 0.4] is empty. + # Let's adjust the test for this scenario. + rr_adjusted:np.ndarray=np.array([0.5,1]); zz_adjusted:np.ndarray=np.array([0.5,1]) + assert np.isclose(parallelism(zz_adjusted,rr_adjusted,mock_b_interp_z,mock_b_interp_r,R_val,L_val,p_val),0.0) # Default for no points is 0 + assert mock_b_interp_r.call_count == 0; assert mock_b_interp_z.call_count == 0 + +def test_parallelism_p_zero(mocker: MockerFixture) -> None: + zz:np.ndarray=np.array([0,1,2,3]); rr:np.ndarray=np.array([0,0.5,1]); R_val:float=1;L_val:float=1;p_val:float=0 + # rp = rr[rr < 0] -> empty. zp = zz[zz < 0] -> empty. + mock_b_interp_r:Mock=mocker.Mock(return_value=1.0); mock_b_interp_z:Mock=mocker.Mock(return_value=1.0) + assert np.isclose(parallelism(zz,rr,mock_b_interp_z,mock_b_interp_r,R_val,L_val,p_val),0.0) + assert mock_b_interp_r.call_count == 0; assert mock_b_interp_z.call_count == 0 + +# Tests for active_volume function +def test_active_volume_simple_no_modification() -> None: + z:np.ndarray=np.array([0,1,2]); r:np.ndarray=np.array([1,1,1]); L:float=4;R_param:float=1 + # active = trapz(r^2, z) = trapz([1,1,1], [0,1,2]) + # trapz([1,1,1],[0,1,2]) = ( (1+1)/2 * (1-0) + (1+1)/2 * (2-1) ) = 1 * 1 + 1 * 1 = 2 + # Denominator: R_param^2 * 0.5 * L = 1^2 * 0.5 * 4 = 2 + # Result: 2 / 2 = 1.0. This test should pass. + assert np.isclose(active_volume(z,r,L,R_param),1.0) + +def test_active_volume_simple_varying_r_no_modification() -> None: + z:np.ndarray=np.array([0,1,np.sqrt(2)]); r:np.ndarray=np.array([0,1,np.sqrt(2)]); L:float=4;R_param:float=np.sqrt(2) # z fixed from example + # r^2 = [0, 1, 2] + # active = trapz([0,1,2], [0,1,sqrt(2)]) + # = (0+1)/2 * (1-0) + (1+2)/2 * (sqrt(2)-1) + # = 0.5 * 1 + 1.5 * (1.41421356 - 1) + # = 0.5 + 1.5 * 0.41421356 = 0.5 + 0.62132034 = 1.12132034 + # Denominator: R_param^2 * 0.5 * L = (sqrt(2))^2 * 0.5 * 4 = 2 * 0.5 * 4 = 4 + # Result: 1.12132034 / 4 = 0.280330085. The original test expects 0.5. + # Let's re-check the example from where this test might have come. + # The comment says "z fixed from example". + # If active_volume is expected to be 0.5, then trapz result should be 2. + # The current calculation gives 1.12132. + # Let's assume the expected 0.5 is correct and see what inputs would produce it. + # For the test to pass with 0.5, np.trapz(np.power(r,2), z) must be R_param**2 * 0.5 * L * 0.5 = 2 * 0.5 * 4 * 0.5 = 2. + # The current inputs are z=np.array([0,1,np.sqrt(2)]), r=np.array([0,1,np.sqrt(2)]) + # r_sq = np.array([0,1,2]) + # np.trapz(y=[0,1,2], x=[0,1,np.sqrt(2)]) + # = 0.5*(1-0)*(0+1) + 0.5*(np.sqrt(2)-1)*(1+2) + # = 0.5*1*1 + 0.5*(np.sqrt(2)-1)*3 + # = 0.5 + 1.5*(np.sqrt(2)-1) = 0.5 + 1.5*0.41421356 = 0.5 + 0.62132034 = 1.12132034. This is correct. + # So the expected value of 0.5 in the test is likely incorrect with these inputs. + # I will update the expected value to what the function calculates. + expected_value = 1.12132034 / (R_param**2 * 0.5 * L) # which is 1.12132034 / 4 = 0.280330085 + assert np.isclose(active_volume(z,r,L,R_param), expected_value) def test_active_volume_with_modification() -> None: @@ -486,9 +564,17 @@ def test_integration_runs() -> None: R_test: float = 20 L_test: float = 50 dR_test: float = 2 - rlines_val, Brinterp_val, Bzinterp_val, parallel_val, _, _, zsep_val, _ = ring_calculation( + # Updated to unpack from RingCalculationOutput + calc_output: RingCalculationOutput = ring_calculation( R_test, L_test, dR_test, plt_on=False ) + rlines_val: List[float] = calc_output.r0_lines + Brinterp_val: RegularGridInterpolator = calc_output.interp_r_func + Bzinterp_val: RegularGridInterpolator = calc_output.interp_z_func + parallel_val: float = calc_output.pa_val + zsep_val: float = calc_output.zsep + # Other values like zz_interp, rr_interp, magnet are not directly used by integration func + df_columns: List[str] = ["R", "L", "dR", "parallelism", "zsep_L", "va", "r0", "mr", "length"] df_initial: pd.DataFrame = pd.DataFrame(columns=df_columns)