move diagnostics to config; implement callable configs

.gitignore

@@ -27,3 +27,4 @@
 /transverse*/
 __pycache__/
 /docs/index.rst
+/docs/_build

config_example.py

@@ -4,7 +4,6 @@
 xi_step_size = .005  #: Step size in time-space coordinate xi
 xi_steps = int(3000 // xi_step_size)  #: Amount of xi steps
 
-diagnostics_each_N_steps = int(1 / xi_step_size)
 
 field_solver_subtraction_trick = 1  #: 0 for Laplace eqn., Helmholtz otherwise
 field_solver_variant_A = True  #: Use Variant A or Variant B for Ex, Ey, Bx, By
@@ -28,3 +27,49 @@ def beam(xi_i, x, y):
             (1 - cos(xi * COMPRESS * sqrt(pi / 2))))
 
 gpu_index = 0  #: Index of the GPU that should perform the calculations
+
+diagnostics_each_N_steps = int(1 / xi_step_size)  #: Diagnostics calls rarity
+diagnostics_wants_numpy_arrays = True  #: Wrap GPUArrays into GPUArraysView
+
+# Some sloppy example diagnostics
+
+Ez_00_history = []
+max_zn = 0
+def diagnostics(config, xi_i, state):
+    global Ez_00_history, max_zn
+
+    import lcode
+    import scipy.ndimage, scipy.signal
+    import numpy as np
+
+    xi = -xi_i * config.xi_step_size
+
+    ro = state.ro  # Unpack once to prevent copying ro several times
+
+    lcode.savepng('transverse', 'ro', xi, ro, vmin=-0.1, vmax=0.1, cmap='bwr')
+
+    # Calculate some high-frequency noise metric
+    bl = scipy.ndimage.gaussian_filter(ro, sigma=(.25 / config.grid_step_size))
+    zn = np.abs(ro - bl).mean() / 4.23045376e-04
+    max_zn = max(max_zn, zn)
+
+    # Copy the entire Ez array, and only use the center cell value
+    Ez_00 = state.Ez[config.grid_steps // 2, config.grid_steps // 2]
+    Ez_00_history.append(Ez_00)
+
+    # Display some information on relative on-axis Ez peak heights
+    Ez_00_array = np.array(Ez_00_history)
+    peak_indices = scipy.signal.argrelmax(Ez_00_array)[0]
+    if peak_indices.size:
+        peak_values = Ez_00_array[peak_indices]
+        rel_deviations_perc = 100 * (peak_values / peak_values[0] - 1)
+        peaks_info = f'{peak_values[-1]:0.4e} {rel_deviations_perc[-1]:+0.2f}%'
+    else:
+        peaks_info = '...'
+
+    print(f'xi={xi:+.4f} {Ez_00:+.4e}|{peaks_info}|zn={max_zn:.3f}')
+
+
+if __name__ == '__main__':  # If executed directly, run LCODE with this config
+    import sys, lcode; lcode.main(sys.modules[__name__])
+    # Or, if you want to, write out all the main loop in here

docs/technicalities/gpu.rst

@@ -116,7 +116,14 @@ but beware!
    so one should bind it to the variable name (``x_init = const_view.x_init``) once
    and reuse the resulting ``numpy`` array.
 
-.. todo:: CODE: wrap the returned array with GPUArraysView by default
+.. autodata:: config_example.diagnostics_wants_numpy_arrays
+
+   When this option is on, func:`config_example.diagnostics` will receive GPUArraysView
+   that will perform the conversion to numpy arrays transparently.
+   Disable this if
+   you want to run your diagnostics on GPU as well,
+   you know what you are doing
+   and your diagnostics function is GPU-aware.
 
 
 Selecting GPU

docs/tour/xi_step.rst

@@ -91,3 +91,18 @@ The result, or the 'new prev'
 The fields from 7., coordinates and momenta from 8., and densities from 9.
 make up the new ``GPUArrays`` collection that would be passed as ``prev``
 to the next iteration of :func:`step()`.
+
+
+Looping it
+----------
+.. autofunction:: lcode.main
+
+   The main loop doesn't do much fancy stuff.
+   It initializes the constant arrays and the initial state,
+   feeds it into :func:`lcode.step` repeatedly and
+   calls diagnostics every :data:`config_example.diagnostics_each_N_steps` steps
+   and after the last step.
+
+   Feel free to rewrite/replace it for embedding purposes.
+
+.. autodata:: config_example.diagnostics_each_N_steps

docs/usage/running_and_embedding.rst

@@ -10,9 +10,10 @@ Installing LCODE into ``PYTHONPATH`` with the likes of ``pip install .`` is poss
 but is not officially supported.
 
 
-Configuration
--------------
-LCODE 3D is configured by placing a file ``config.py`` into the current working directory.
+Configuration with config.py
+----------------------------
+The simplest way to configure LCODE 3D is
+to place a file ``config.py`` into the current working directory.
 An example is provided as ``config_example.py``.
 
 The file gets imported by the standard Python importing mechanism,
@@ -21,10 +22,36 @@ the resulting module is passed around internally as ``config``.
 One can use all the features of Python inside the configuration file,
 from arithmetic expressions and functions to other modules and metaprogramming.
 
+For executing, try
+``python3 lcode.py``, ``python lcode.py`` or ``./lcode.py``.
 
-Execution
----------
-``python3 lcode.py``, ``python lcode.py`` or ``./lcode.py``
 
+Configuration with executable config
+------------------------------------
+The second way is to create a config file,
+make sure that it ends with
 
-.. todo:: CODE: embedding
+.. code-block:: python
+
+  if __name__ == '__main__':
+      import sys, lcode; lcode.main(sys.modules[__name__])
+
+and execute the configuration file directly.
+
+``lcode.py`` should be either located inside the working directory
+or installed to ``PYTHONPATH``.
+
+This way may be preferrable if you want to manage several configuration files.
+
+For executing, try
+``python3 some_config.py``, ``python some_config.py`` or ``./some_config.py``.
+
+
+Configuration with a Python object
+----------------------------------
+The third way is to create a Python object with all the required attributes
+and to pass it to
+:func:`lcode.main`, :func:`lcode.init` + :func:`lcode.step`,
+or even smaller bits and pieces from your code.
+
+This way may come in handy if you are calling LCODE 3D from other Python programs.

lcode.py

@@ -18,22 +18,13 @@
 
 from math import sqrt, floor
 
-import os
-import sys
-
-import matplotlib.pyplot as plt
-
 import numpy as np
 
 import numba
 import numba.cuda
 
 import cupy as cp
 
-import scipy.ndimage
-import scipy.signal
-
-
 # Prevent all CPU cores waiting for the GPU at 100% utilization (under conda).
 # os.environ['OMP_NUM_THREADS'] = '1'
 
@@ -975,80 +966,39 @@ def zeros():
     return xs, ys, const, virt_params, state
 
 
-# Some really sloppy diagnostics #
-
-max_zn = 0
-def diags_ro_zn(config, ro):
-    global max_zn
-
-    sigma = 0.25 / config.grid_step_size
-    blurred = scipy.ndimage.gaussian_filter(ro, sigma=sigma)
-    hf = ro - blurred
-    zn = np.abs(hf).mean() / 4.23045376e-04
-    max_zn = max(max_zn, zn)
-    return max_zn
-
-
-def diags_peak_msg(Ez_00_history):
-    Ez_00_array = np.array(Ez_00_history)
-    peak_indices = scipy.signal.argrelmax(Ez_00_array)[0]
-
-    if peak_indices.size:
-        peak_values = Ez_00_array[peak_indices]
-        rel_deviations_perc = 100 * (peak_values / peak_values[0] - 1)
-        return (f'{peak_values[-1]:0.4e} '
-                f'{rel_deviations_perc[-1]:+0.2f}%')
-    else:
-        return '...'
+# Diagnostics helper #
 
+def savepng(dirname, prefix, xi, arr, **kwargs):
+    import matplotlib.pyplot as plt
+    import os
 
-def diags_ro_slice(config, xi_i, xi, ro):
-    if xi_i % int(1 / config.xi_step_size):
-        return
-    if not os.path.isdir('transverse'):
-        os.mkdir('transverse')
-
-    fname = f'ro_{xi:+09.2f}.png' if xi else 'ro_-00000.00.png'
-    plt.imsave(os.path.join('transverse', fname), ro.T,
-               origin='lower', vmin=-0.1, vmax=0.1, cmap='bwr')
-
-
-def diagnostics(view_state, config, xi_i, Ez_00_history):
-    xi = -xi_i * config.xi_step_size
-
-    Ez_00 = Ez_00_history[-1]
-    peak_report = diags_peak_msg(Ez_00_history)
+    if not os.path.isdir(dirname):
+        os.mkdir(dirname)
 
-    ro = view_state.ro
-    max_zn = diags_ro_zn(config, ro)
-    diags_ro_slice(config, xi_i, xi, ro)
-
-    print(f'xi={xi:+.4f} {Ez_00:+.4e}|{peak_report}|zn={max_zn:.3f}')
-    sys.stdout.flush()
+    fname = prefix + (f'_{xi:+09.2f}.png' if xi else '_-00000.00.png')
+    plt.imsave(os.path.join(dirname, fname), arr.T, origin='lower', **kwargs)
 
 
 # Main loop #
 
-def main():
-    import config
-    with cp.cuda.Device(config.gpu_index):
+def main(config=None):
+    if config is None:
+        import config
 
+    with cp.cuda.Device(config.gpu_index):
         xs, ys, const, virt_params, state = init(config)
-        Ez_00_history = []
 
         for xi_i in range(config.xi_steps):
             beam_ro = config.beam(xi_i, xs, ys)
 
             state = step(config, const, virt_params, state, beam_ro)
-            view_state = GPUArraysView(state)
-
-            ez = view_state.Ez[config.grid_steps // 2, config.grid_steps // 2]
-            Ez_00_history.append(ez)
 
             time_for_diags = xi_i % config.diagnostics_each_N_steps == 0
             last_step = xi_i == config.xi_steps - 1
             if time_for_diags or last_step:
-                diagnostics(view_state, config, xi_i, Ez_00_history)
+                wrap = config.diagnostics_wants_numpy_arrays
+                config.diagnostics(config, xi_i,
+                                   GPUArraysView(state) if wrap else state)
 
 
 if __name__ == '__main__':