forked from psc-code/psc
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
radiation: add some basic dipole / quadrupole radiation
This really could use some more work in cleaning it up, etc, but it does pretty much do what it's supposed to do.
- Loading branch information
Showing
2 changed files
with
252 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,194 @@ | ||
| { | ||
| "cells": [ | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": 1, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "import matplotlib as mpl\n", | ||
| "import matplotlib.pyplot as plt\n", | ||
| "from matplotlib.animation import FuncAnimation\n", | ||
| "import xarray as xr\n", | ||
| "import numpy as np\n", | ||
| "#import pscpy\n", | ||
| "\n", | ||
| "%matplotlib ipympl \n", | ||
| "#plt.rcParams['figure.figsize'] = [16, 10]" | ||
| ] | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": 2, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "# dipole far field\n", | ||
| "dir = \"/Users/kai/src/psc/build-arm64\"\n", | ||
| "steps = range(0, 400, 5)\n", | ||
| "vmax = .0001\n", | ||
| "\n", | ||
| "# orig dipole near field\n", | ||
| "dir = \"/Users/kai/src/psc/build-arm64-2\"\n", | ||
| "steps = range(0, 800, 5)\n", | ||
| "vmax = .01\n", | ||
| "\n", | ||
| "# quadrupole\n", | ||
| "dir = \"/Users/kai/src/psc/build-arm64-2\"\n", | ||
| "steps = range(0, 1000, 20)\n", | ||
| "vmax = .00005\n", | ||
| "\n", | ||
| "def open_step(step):\n", | ||
| " return xr.open_dataset(f\"{dir}/pfd.{step:09d}.bp\", #engine='pscadios2',\n", | ||
| " species_names=['e', 'i'])\n", | ||
| "\n" | ||
| ] | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "datas = {}\n", | ||
| "for step in steps:\n", | ||
| " ds = open_step(step)\n", | ||
| " datas[step] = ds.ez_ec.sel(y=0.).T, float(ds.time)\n", | ||
| "#datas" | ||
| ] | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "def get_data(step):\n", | ||
| " return datas[step]\n", | ||
| "\n", | ||
| "steps = list(steps)\n", | ||
| "\n", | ||
| "fig, ax = plt.subplots()\n", | ||
| "step = steps[0]\n", | ||
| "fld, time = get_data(step)\n", | ||
| "cax = fld.plot(vmin=-vmax, vmax=vmax, cmap='coolwarm')\n", | ||
| "ax.set_title(f\"step {step} time {time:6.2f}\")\n", | ||
| "ax.set_aspect(1.)\n", | ||
| "\n", | ||
| "def animate(step):\n", | ||
| " fld, time = get_data(step)\n", | ||
| " cax.set_array(fld.values.flatten())\n", | ||
| " ax.set_title(f\"step {step} time {time:6.2f}\")\n", | ||
| "\n", | ||
| "ani = FuncAnimation(\n", | ||
| " fig, # figure\n", | ||
| " animate, # name of the function above\n", | ||
| " frames=steps[1:], # Could also be iterable or list\n", | ||
| " interval=200, # ms between frames\n", | ||
| " blit=False\n", | ||
| ")\n", | ||
| "\n", | ||
| "plt.show()\n", | ||
| "#ani" | ||
| ] | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": 5, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "def plot_fields(fldnames, fld_kwargs=None):\n", | ||
| " fig, axs = plt.subplots(1, len(fldnames))\n", | ||
| " if len(fldnames) == 1: axs = [axs]\n", | ||
| " for i, fldname in enumerate(fldnames):\n", | ||
| " fld = ds[fldname].sel(y=0)\n", | ||
| " if fld_kwargs:\n", | ||
| " kwargs = fld_kwargs[i]\n", | ||
| " else:\n", | ||
| " kwargs = {}\n", | ||
| " fld.plot(ax=axs[i], **kwargs)\n", | ||
| " axs[i].set_aspect('equal')\n", | ||
| " plt.tight_layout()" | ||
| ] | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "\n", | ||
| "plot_fields(['ex_ec', 'ey_ec', 'ez_ec'])\n", | ||
| "# fld_kwargs=[{\"vmin\": -.0065}, {\"vmin\": -.02}, {}])" | ||
| ] | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "def format_radians_label(float_in):\n", | ||
| " # Converts a float value in radians into a\n", | ||
| " # string representation of that float\n", | ||
| " string_out = str(float_in / (np.pi))+\"π\"\n", | ||
| " \n", | ||
| " return string_out\n", | ||
| "\n", | ||
| "def convert_polar_xticks_to_radians(ax):\n", | ||
| " # Converts x-tick labels from degrees to radians\n", | ||
| " \n", | ||
| " # Get the x-tick positions (returns in radians)\n", | ||
| " label_positions = ax.get_xticks()\n", | ||
| " \n", | ||
| " # Convert to a list since we want to change the type of the elements\n", | ||
| " labels = list(label_positions)\n", | ||
| " \n", | ||
| " # Format each label (edit this function however you'd like)\n", | ||
| " labels = [format_radians_label(label) for label in labels]\n", | ||
| " \n", | ||
| " ax.set_xticklabels(labels)\n", | ||
| " \n", | ||
| "theta = np.linspace(-np.pi, np.pi, 100)\n", | ||
| "\n", | ||
| "fig, ax = plt.subplots(subplot_kw={'projection': 'polar'})\n", | ||
| "ax.plot(theta, np.sin(theta)**2)\n", | ||
| "ax.set_rticks([0.25, 0.5, 0.75, 1])\n", | ||
| "ax.set_theta_zero_location(\"N\")\n", | ||
| "#convert_polar_xticks_to_radians(ax)\n", | ||
| "\n", | ||
| "ax.set_title(\"Radiated dipole power\")\n", | ||
| "plt.show()" | ||
| ] | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": null, | ||
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [] | ||
| } | ||
| ], | ||
| "metadata": { | ||
| "kernelspec": { | ||
| "display_name": ".venv", | ||
| "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.12.5" | ||
| } | ||
| }, | ||
| "nbformat": 4, | ||
| "nbformat_minor": 4 | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters