simulate uniform electromagnetic field

.github/workflows/release-docs.yml

@@ -0,0 +1,33 @@
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - uniform-em
+
+env:
+  PACK_DIR: /root/.pack
+
+jobs:
+  docs:
+    runs-on: ubuntu-latest
+    container: ghcr.io/stefan-hoeck/idris2-pack
+    steps:
+    - uses: actions/checkout@v4
+    - name: Build docs
+      run: |
+        apt-get update && apt-get install -y curl python3-venv
+        pack --no-prompt run plasma.ipkg
+        python3 -m venv .venv
+        . .venv/bin/activate
+        pip install -r plot/requirements.txt -c plot/constraints.txt
+        python3 plot/plot.py
+        mkdir docs
+        mv index.html docs/
+    - name: Upload docs
+      run: |
+        git config --global --add safe.directory "$GITHUB_WORKSPACE"
+        git add -f docs/
+        git config user.email "none"
+        git config user.name "none"
+        git commit -m "build documentation"
+        git push -f origin HEAD:gh-pages

.gitignore

@@ -0,0 +1,171 @@
+.idea/
+.venv/
+*.csv
+*.html
+
+# Idris 2
+*.ttc
+*.ttm
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

README.md

@@ -1,2 +1,5 @@
 # plasma
-Plasma physics simulations
+
+_Plasma physics simulations_
+
+See the [website](https://joelberkeley.github.io/plasma/) for progress so far.

pack.toml

@@ -0,0 +1,16 @@
+[custom.all.plasma]
+type = "local"
+path = "."
+ipkg = "plasma.ipkg"
+test = "test/test.ipkg"
+
+[custom.all.plasma-test]
+type = "local"
+path = "test"
+ipkg = "test.ipkg"
+
+[custom.all.gnuplot]
+type   = "git"
+url    = "https://github.com/stefan-hoeck/idris2-gnuplot"
+commit = "latest:main"
+ipkg   = "gnuplot.ipkg"

plasma.ipkg

@@ -0,0 +1,16 @@
+package plasma
+version = 0.1.0
+
+depends =
+    spidr
+  , pjrt-plugin-xla-cpu
+  , gnuplot
+
+modules =
+    EM
+  , Sampling
+
+main = Main
+
+executable = "plasma"
+sourcedir = "src"

plot/constraints.txt

@@ -0,0 +1,9 @@
+numpy==2.0.0
+packaging==24.1
+pandas==2.2.2
+plotly==5.22.0
+python-dateutil==2.9.0.post0
+pytz==2024.1
+six==1.16.0
+tenacity==8.5.0
+tzdata==2024.1

plot/plot.py

@@ -0,0 +1,26 @@
+import pandas as pd
+import numpy as np
+import plotly.graph_objs as go
+
+p = pd.read_csv("ions.csv")
+e = pd.read_csv("electrons.csv")
+
+p = [x.drop("p", axis="columns") for _, x in p.groupby("p")]
+e = [x.drop("p", axis="columns") for _, x in e.groupby("p")]
+
+fig = go.Figure()
+
+for colour, trace in [("blue", p), ("purple", e)]:
+    for particle, t in enumerate(trace):
+        t = np.array(t)
+        fig.add_trace(
+            go.Scatter3d(
+                x=t[:, 0],
+                y=t[:, 1],
+                z=t[:, 2],
+                mode="lines",
+                marker={"color": colour},
+            )
+        )
+
+fig.write_html("index.html")

plot/requirements.txt

@@ -0,0 +1,3 @@
+numpy
+plotly
+pandas

src/EM.idr

@@ -0,0 +1,38 @@
+module EM
+
+import Tensor
+import Literal
+
+||| The trajectory of `n` charged particles through a uniform electromagnetic field.
+|||
+||| @ex The x-component of the electric field.
+||| @ey The y-component of the electric field.
+||| @bz The z-component of the magnetic field.
+||| @r0 The starting coordinate of each particle.
+||| @v0 The starting velocity of each particle.
+||| @q The charge, in units of e, of each particle.
+||| @m The mass, in eV, of each particle.
+||| @interval The time interval between steps.
+||| @steps The number of time steps.
+export
+uniformEM :
+  (ex, ez, bz : Tensor [] F64) ->
+  {n : _} ->
+  (r0 : Tensor [n, 3] F64) ->
+  (v0 : Tensor [n, 3] F64) ->
+  (q, m : Tensor [] F64) ->
+  (interval : Tensor [] F64) ->
+  (steps : Nat) ->
+  Tag $ Tensor [n, steps, 3] F64
+uniformEM ex ez bz r0 v0 q m dt steps = do
+  let v0x = slice [all, 0.to 1] v0
+      v0y = slice [all, 1.to 2] v0
+      v0orth = broadcast {to = [1, n, steps]} $ sqrt (v0x ^ fill 2.0 + v0y ^ fill 2.0)
+      w = abs q * bz / m
+      sign = select (q > 0.0) 1.0 (-1.0)
+  t <- tag $ dt * iota {shape = [1, n, steps]} 2
+  let x : Tensor [1, n, steps] F64 = reshape $ v0orth / w * sin (w * t)
+      y : Tensor [1, n, steps] F64 = reshape $ sign * v0orth / w * cos (w * t) - t * broadcast (ex / bz)
+      z : Tensor [1, n, steps] F64 = reshape $
+        q * ez * (t ^ fill 2.0) / (2.0 * m) + t * broadcast (slice [all, 2.to 3] v0)
+  pure (broadcast (expand 1 r0) + transpose [1, 2, 0] (concat 0 x $ concat 0 y z))

src/Main.idr

@@ -0,0 +1,68 @@
+module Main
+
+import Data.String
+import Data.Vect
+import System
+import System.File
+
+import Tensor
+import Literal
+import PjrtPluginXlaCpu
+
+import Sampling
+import EM
+
+toCSV : {m, n : _} -> Vect n String -> Literal [m, n] Double -> String
+toCSV headers xs =
+  let strs = map tostr xs
+      lines = map (joinBy "," . toList) $ cast @{toArray} strs
+      headers = joinBy "," $ toList headers
+   in joinBy "\n" (headers :: toList lines)
+
+  where
+  tostr : Double -> String
+  tostr x =
+    let inf = 1.0 / 0.0 in
+    if x == inf then "inf" else
+    if x == -inf then "-inf" else
+    if x /= x then "nan" else show x
+
+product2 : (a, b : Nat) -> (rest : List Nat) ->
+           product (the (List _) (a :: b :: rest)) === product (the (List _) (a * b :: rest))
+product2 Z _ _ = Refl
+product2 (S aa) _ _ = ?product2Succ
+
+main : IO ()
+main = do
+  let key = tensor {dtype = U64} 1
+      temp : Tensor [] F64 = 1.0e4  -- eV
+      me : Tensor [] F64 = 0.511e6  -- eV
+      mp : Tensor [] F64 = 938.272e6  -- eV
+      e : Tensor [] F64 = 1.0
+      steps = 1000
+      particleCount := 20
+      paths : Tag $ TensorList [[particleCount * steps, 4], [particleCount * steps, 4]] _ = do
+        let rand = Sampling.normal {n = particleCount} key
+        (r0e, r0p, v0e, v0p) <- evalStateT (tensor {dtype = U64} [Scalar 1]) $ do
+          let r0e = Tensor.(*) (broadcast (tensor [1.0e-2, 1.0e-2, 1.0])) !rand
+              r0p = Tensor.(*) (broadcast (tensor [1.0e-2, 1.0e-2, 1.0])) !rand
+              v0e = sqrt (3.0 * temp / me) * !rand
+              v0p = sqrt (3.0 * temp / mp) * !rand
+          pure (r0e, r0p, v0e, v0p)
+        -- these are unrealistic field strengths, but they show paths more clearly
+        -- https://en.wikipedia.org/wiki/Orders_of_magnitude_(magnetic_field)
+        let applyEM = uniformEM {ex = 1.0e6, ez = 1.0e3, bz = 1.0e9}
+        ions <- applyEM {r0 = r0p, v0 = v0p, q = e, m = mp, interval = 0.01, steps = steps}
+        electrons <- applyEM {r0 = r0e, v0 = v0e, q = -e, m = me, interval = 0.01, steps = steps}
+        let count = reshape {sizesEqual = product2 particleCount steps [1]}
+              {to = [particleCount * steps, 1]} $ iota {shape = [particleCount, steps, 1]} 0
+            ions = reshape {sizesEqual = product2 particleCount steps [3]}
+              {to = [particleCount * steps, 3]} ions
+            electrons = reshape {sizesEqual = product2 particleCount steps [3]}
+              {to = [particleCount * steps, 3]} electrons
+        pure [concat 1 count ions, concat 1 count electrons]
+
+  device <- eitherT (die . show) pure device
+  [p, e] <- TensorList.Tag.eval device paths
+  either (const $ die {io = IO} "file error") pure =<< writeFile "ions.csv" (toCSV ["p", "x", "y", "z"] p)
+  either (const $ die {io = IO} "file error") pure =<< writeFile "electrons.csv" (toCSV ["p", "x", "y", "z"] e)