Move notebook, add fitting dev

docs/examples/fields/solenoid_modeling.ipynb

@@ -0,0 +1,275 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0fe9a7b5-2538-4fb1-93f5-c3dd9e4b1023",
+   "metadata": {},
+   "source": [
+    "# Solenoid Modeling"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7cc6ce6d-99ba-4254-8c49-55cb1a3297c9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pmd_beamphysics.fields.solenoid import make_solenoid_fieldmesh\n",
+    "\n",
+    "from pmd_beamphysics.fields.analysis import check_static_div_equation\n",
+    "from pmd_beamphysics.units import mu_0\n",
+    "\n",
+    "from pmd_beamphysics import FieldMesh\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cfd33797-b867-4353-9626-cf6ef049937f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "FM = make_solenoid_fieldmesh(\n",
+    "    radius=0.05,\n",
+    "    L=0.2,\n",
+    "    rmax=0.1,\n",
+    "    zmin=-0.4,\n",
+    "    zmax=0.4,\n",
+    "    nr=101,\n",
+    "    nz=200,\n",
+    "    nI=1,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3bb318b7-c742-4296-9e22-36128eb7b90f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "FM.plot()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fb8043b4-8030-464a-99a2-fb2888d48eb2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "FM.plot_onaxis()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8a317af4-d705-4481-bc3a-0ea9d767cbb2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "check_static_div_equation(FM, rtol=1e-2, plot=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "40936679-14e0-439d-bd12-462c908312a6",
+   "metadata": {},
+   "source": [
+    "## Hard edge\n",
+    "\n",
+    "Making the radius very small approximates a hard-edge model.\n",
+    "\n",
+    "Here we expext that $B_z = \\mu_0 n I$"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2616d48f-47f6-4b28-b637-8347e93f15a0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "FM_hard = make_solenoid_fieldmesh(\n",
+    "    radius=1e-9,\n",
+    "    L=0.2,\n",
+    "    rmax=0.1,\n",
+    "    zmin=-0.4,\n",
+    "    zmax=0.4,\n",
+    "    nr=101,\n",
+    "    nz=200,\n",
+    "    nI=1,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f21a3461-b9c5-4dfa-9cac-e82cca09eead",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "FM_hard.Bz[0, 0, :].max() == mu_0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "886f334c-3024-4650-8081-ae68ebdcaaaa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "FM_hard.plot_onaxis()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bfb0168b-cf80-47d4-a77d-23391391d112",
+   "metadata": {},
+   "source": [
+    "## Compare with a real solenoid"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8148485f-09e5-44dd-96a7-c9cc506f35ff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "FM2 = FieldMesh(\"../data/solenoid.h5\")\n",
+    "FM2.plot_onaxis()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c8ce66f4-1b14-45f5-b0d7-6e402868cf05",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "FM3 = make_solenoid_fieldmesh(\n",
+    "    radius=0.0191,\n",
+    "    L=0.02936 * 2,\n",
+    "    rmax=0.1,\n",
+    "    zmin=-0.1,\n",
+    "    zmax=0.1,\n",
+    "    nr=100,\n",
+    "    nz=40,\n",
+    "    nI=1,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1313ec9a-2338-46ae-a6b0-03da50510c7a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots()\n",
+    "z0, Bz0 = FM2.axis_values(\"z\", \"Bz\")\n",
+    "Bz0 = np.real(Bz0 / Bz0.max())\n",
+    "ax.plot(z0, Bz0, label=\"Superfish\")\n",
+    "\n",
+    "z, Bz = FM3.axis_values(\"z\", \"Bz\")\n",
+    "Bz = np.real(Bz / Bz.max())\n",
+    "ax.plot(z, Bz, label=\"Ideal\")\n",
+    "plt.legend()\n",
+    "ax.set_ylim(0, None)\n",
+    "ax.set_xlabel(r\"$z$ (m)\")\n",
+    "ax.set_ylabel(r\"B_z (T)\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "78ca976c-b608-461d-95ce-d7b147bba952",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "FM2.plot()\n",
+    "FM3.plot()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e61c96b2-2ad9-4920-bb80-6a4d54549ab3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "from scipy.optimize import curve_fit\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "\n",
+    "# Define the Bz model from the given formula\n",
+    "def Bz_model(z, a, b):\n",
+    "    \"\"\"\n",
+    "    Normalized ideal on-axis model\n",
+    "    \"\"\"\n",
+    "    A = np.hypot(a, b) / (2 * b)\n",
+    "    term1 = (z + b) / np.hypot(z + b, a)\n",
+    "    term2 = (z - b) / np.hypot(z - b, a)\n",
+    "    return A * (term1 - term2)\n",
+    "\n",
+    "\n",
+    "z_data = z0\n",
+    "Bz_data = Bz0\n",
+    "\n",
+    "# Normalize the data\n",
+    "Bz_data = Bz_data / np.max(Bz_data)  # Normalize Bz to a maximum of 1\n",
+    "z_data = z_data - z_data[np.argmax(Bz_data)]  # Shift z so maximum is at z=0\n",
+    "\n",
+    "# Fit the data to the model\n",
+    "popt, pcov = curve_fit(Bz_model, z_data, Bz_data, p0=[0.1, 0.1])\n",
+    "\n",
+    "# Extract fitted parameters\n",
+    "a_fit, b_fit = popt\n",
+    "a_fit = abs(a_fit)\n",
+    "b_fit = abs(b_fit)\n",
+    "print(f\"Fitted parameters: a = {a_fit}, b = {b_fit}\")\n",
+    "\n",
+    "# Plot the data and the fitted curve\n",
+    "z_fit = np.linspace(np.min(z_data), np.max(z_data), 500)\n",
+    "Bz_fit = Bz_model(z_fit, a_fit, b_fit)\n",
+    "\n",
+    "plt.figure(figsize=(8, 6))\n",
+    "plt.scatter(z_data, Bz_data, label=\"Normalized Data\", color=\"blue\", alpha=0.7)\n",
+    "plt.plot(z_fit, Bz_fit, label=\"Fitted Model\", color=\"red\", linewidth=2)\n",
+    "plt.xlabel(\"z (normalized)\")\n",
+    "plt.ylabel(\"Bz (normalized)\")\n",
+    "plt.title(\"Fitting Bz Model to Data\")\n",
+    "plt.legend()\n",
+    "plt.grid()\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.13.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}