# emout

A Python library for parsing and visualizing output files generated by [EMSES](https://github.com/Nkzono99/MPIEMSES3D) simulations.

- **Documentation:** [https://nkzono99.github.io/emout/](https://nkzono99.github.io/emout/)

## Installation

```bash
pip install emout
```

## Example

- [Visualization of simulation results for lunar surface charging](https://nbviewer.org/github/Nkzono99/examples/blob/main/examples/emout/example.ipynb)

## Overview

When you run EMSES simulations, the results (e.g., potentials, densities, currents) are output in `.h5` files, and a parameter file (`plasma.inp`) contains the simulation settings. **emout** helps you:

1. [**Load and handle** these outputs easily in Python.](#loading-data)
2. [**Access simulation parameters** from `plasma.inp`](#retrieving-the-parameter-file-plasmainp).
3. [**Plot scalar/vector fields** for quick data visualization.](#plotting-data)
4. [**Convert units** between EMSES internal units and SI units.](#working-with-units)
5. [**Combine multi-run outputs** for continued or appended simulations.](#handling-appended-simulation-outputs)
6. [**Apply data masking** to filter data regions of interest.](#data-masking)
7. [**Create animations** to see how data evolve over time.](#creating-animations)
8. [**(Experimental) Solve Poisson’s equation** directly from charge distributions for 3D potentials.](#solving-poissons-equation-experimental)
9. [**(Experimental) Compute backtraces and probabilities** for particle trajectories (via data.backtrace)](#backtrace-usage-examples-experimental).

Below, you will find usage examples that assume the following directory structure:

```
.
└── output_dir
    ├── plasma.inp
    ├── phisp00_0000.h5
    ├── nd1p00_0000.h5
    ├── nd2p00_0000.h5
    ├── j1x00_0000.h5
    ├── j1y00_0000.h5
    ...
    └── bz00_0000.h5
```

---

## Usage

### Loading Data

```python
import emout

# Initialize Emout with the path to the output directory
data = emout.Emout('output_dir')

# Access arrays by their variable names (derived from EMSES filename)
data.phisp   # Data from phisp00_0000.h5
len(data.phisp)   # Number of time steps
data.phisp[0].shape

data.j1x
data.bz
data.j1xy  # Vector data from "j1x00_0000.h5" + "j1y00_0000.h5"

# Data created by "relocating" ex00_0000.h5
data.rex

# Access data as a pandas DataFrame
data.icur
data.pbody
```

---

### Retrieving the Parameter File (plasma.inp)

```python
# The namelist is parsed into a dictionary-like structure
data.inp
data.inp['tmgrid']['nx']  # Access via group name and parameter name
data.inp['nx']            # Group name can be omitted if not ambiguous
data.inp.tmgrid.nx        # Access like object attributes
data.inp.nx               # Still valid
```

---

### Plotting Data

```python
x, y, z = 32, 32, 100

# Basic 2D plot (xy-plane at z=100 for the last timestep)
data.phisp[-1, z, :, :].plot()

# Line plot along z-axis at x=32, y=32
data.phisp[-1, :, y, x].plot()

# Plot using SI units
data.phisp[-1, z, :, :].plot(use_si=True)  # Default is True

# Show or save plot
data.phisp[-1, z, :, :].plot(show=True)
data.phisp[-1, z, :, :].plot(savefilename='phisp.png')

# Plot vector field as a streamline
data.j1xy[-1, z, :, :].plot()
```

---

### Working with Units

> **Note**  
> EMSES → SI conversion is supported only when the first line of `plasma.inp` includes something like:
> ```text
> !!key dx=[0.5],to_c=[10000.0]
> ```
> where `dx` is the grid spacing [m], and `to_c` is the (normalized) speed of light used internally by EMSES.

```python
# Converting between physical (SI) and EMSES units
data.unit.v.trans(1)    # Convert 1 m/s to EMSES velocity unit
data.unit.v.reverse(1)  # Convert 1 EMSES velocity unit to m/s

# Access converted data in SI units
phisp_volt = data.phisp[-1, :, :, :].val_si       # Potential in volts [V]
j1z_A_per_m2 = data.j1z[-1, :, :, :].val_si       # Current density [A/m^2]
nd1p_per_cc = data.nd1p[-1, :, :, :].val_si       # Number density [1/cm^3]
```


<details>
    
<summary>Unit Name List</summary>

```
B = Magnetic flux density [T]
C = Capacitance [F]
E = Electric field [V/m]
F = Force [N]
G = Conductance [S]
J = Current density [A/m^2]
L = Inductance [H]
N = Flux [/m^2s]
P = Power [W]
T = Temperature [K]
W = Energy [J]
a = Acceleration [m/s^2]
c = Light Speed [m/s]
e = Napiers constant []
e0 = FS-Permttivity [F/m]
eps = Permittivity  [F/m]
f = Frequency [Hz]
i = Current [A]
kB = Boltzmann constant [J/K]
length = Sim-to-Real length ratio [m]
m = Mass [kg]
m0 = FS-Permeablity [N/A^2]
mu = Permiability [H/m]
n = Number density [/m^3]
phi = Potential [V]
pi = Circular constant []
q = Charge [C]
q_m = Charge-to-mass ratio [C/kg]
qe = Elementary charge [C]
qe_me = Electron charge-to-mass ratio [C/kg]
rho = Charge density [C/m^3]
t = Time [s]
v = Velocity [m/s]
w = Energy density [J/m^3]
```

</details>

---

### Handling Appended Simulation Outputs

<details>

<summary>Examples</summary>

If your simulation continues and creates new directories:

```python
import emout

# Merge multiple output directories into one Emout object
data = emout.Emout('output_dir', append_directories=['output_dir_2', 'output_dir_3'])

# Same as above if 'ad="auto"' is specified (detects appended outputs automatically)
data = emout.Emout('output_dir', ad='auto')
```

</details>

---

### Data Masking

<details>

<summary>Examples</summary>

```python
# Mask values below the average
data.phisp[1].masked(lambda phi: phi < phi.mean())

# Equivalent manual approach
phi = data.phisp[1].copy()
phi[phi < phi.mean()] = float('nan')
```
    
</details>

---

### Creating Animations

<details>

<summary>Examples</summary>

```python
# Create a time-series animation along the first axis (time = 0)
x, y, z = 32, 32, 100
data.phisp[:, z, :, :].gifplot()

# Specify a different axis (default is axis=0)
data.phisp[:, z, :, :].gifplot(axis=0)

# Save animation as a GIF
data.phisp[:, z, :, :].gifplot(action='save', filename='phisp.gif')

# Display the animation inline in a Jupyter notebook
data.phisp[:, z, :, :].gifplot(action='to_html')

# Combining multiple frames for a single animation
updater0 = data.phisp[:, z, :, :].gifplot(action='frames', mode='cmap')
updater1 = data.phisp[:, z, :, :].build_frame_updater(mode='cont')
updater2 = data.nd1p[:, z, :, :].build_frame_updater(mode='cmap', vmin=1e-3, vmax=20, norm='log')
updater3 = data.nd2p[:, z, :, :].build_frame_updater(mode='cmap', vmin=1e-3, vmax=20, norm='log')
updater4 = data.j2xy[:, z, :, :].build_frame_updater(mode='stream')

layout = [
    [
        [updater0, updater1],
        [updater2],
        [updater3, updater4]
    ]
]
animator = updater0.to_animator(layout=layout)
animator.plot(action='to_html')  # or 'save', 'show', etc.
```

</details>

---

### Solving Poisson’s Equation (Experimental)

You can solve Poisson’s equation from 3D charge distributions using `emout.poisson` (depends on `scipy`):

<details>

<summary>Examples</summary>
    
```python
import numpy as np
import scipy.constants as cn
from emout import Emout, poisson

data = Emout('output_dir')
dx = data.inp.dx  # [m] Grid spacing
rho = data.rho[-1].val_si  # [C/m^3] Charge distribution
btypes = ["pdn"[i] for i in data.inp.mtd_vbnd]  # Boundary conditions

# Solve Poisson’s equation for potential
phisp = poisson(rho, dx=dx, btypes=btypes, epsilon_0=cn.epsilon_0)

# Compare with EMSES potential
np.allclose(phisp, data.phisp[-1])  # Should be True (within numerical tolerance)
```

</details>

---

### Backtrace Usage Examples (Experimental)

<details>
    
<summary>Examples</summary>

#### Running backtrace on HPC computational nodes with Dask (>= Python3.10)

If you’ve set up a Dask cluster via `emout.distributed`, all of the `data.backtrace` calls below will actually run on your computational nodes instead of your login node.

```python
from emout.distributed import start_cluster, stop_cluster
import emout

# ① Dask クラスタを起動（SLURM ジョブを一時的に作成して Worker を常駐させる）
client = start_cluster(
    partition="gr20001a",   # 使用するキュー
    processes=1,            # プロセス数
    cores=112,              # コア数
    memory="60G",           # メモリ
    walltime="03:00:00",    # 最大実行時間
    scheduler_ip=None,  # ログインノード上の Scheduler IP (e.g. "10.10.64.1", Noneで自動検索)
    scheduler_port=32332,       # Scheduler ポート
)

# ② 通常の data.backtrace API を呼び出すだけで、
#    図のようにバックトレース関数群が計算ノード上で実行されます
data = emout.Emout("output_dir")
result = data.backtrace.get_probabilities(
    128, 128, 200,
    (-data.inp.path[0]*3, data.inp.path[0]*3, 500),
    1,
    (-data.inp.path[0]*4, data.inp.path[0]*3, 500),
    ispec=0,
    istep=-1,
    dt=data.inp.dt,
    max_step=100000,
    n_threads=112,
)
result.vxvz.plot()

# ③ 終了後はクライアントを閉じて Scheduler を停止
stop_cluster()
````


Below are three example workflows demonstrating how to use the `data.backtrace` interface. All examples assume you have already created an `Emout` object named `data`.


#### 1. Perform a batch backtrace and plot sampled trajectories

In this example, we generate 10 particles with varying initial positions and velocities, compute their backtraces, randomly sample 5 backtrace trajectories, and then plot the $x–y$ projection of each sampled trajectory.

```python
import numpy as np
import matplotlib.pyplot as plt

# 1) Prepare 10 initial positions and velocities
#    Positions: (128, 128, 60 + i*2) for i = 0, 1, …, 9
#    Velocities: (0, 0, -10*i) for i = 0, 1, …, 9
positions = np.array([[128, 128, 60 + 2 * i] for i in range(10)])
velocities = np.array([[0, 0, -10 * i] for i in range(10)])

# 2) Compute backtraces for all 10 particles
#    This returns a MultiBacktraceResult object
backtrace_result = data.backtrace.get_backtraces(positions, velocities)

# 3) Sample 5 trajectories. Then plot the x vs y projection.
sampled_result = backtrace_result.sample(5)
ax = sampled_result.xy.plot()
ax.set_title("Backtrace: Sampled 10 Trajectories (x vs y)")
plt.show()
```


#### 2. Compute a velocity–space probability distribution and plot $v_x$ vs $v_z$

Here, we call `get_probabilities(...)` to generate a 6D phase grid $(x, y, z, vx, vy, vz)$, invoke the backend solver to compute the probability for each $(x,y,z,v_x,v_y,v_z)$ cell, and then plot the 2D slice in $(v_x, v_z)$ space.

```python
# 1) Call get_probabilities on a 1×1×1 (x,y,z) cell at (128,128,60),
#    with vx in [-3*dx, +3*dx] sampled over 10 points,
#    vy fixed at 0, and vz in [-3*dx, 0] sampled over 10 points.
#    ispec=0 indicates electron species.
probability_result = data.backtrace.get_probabilities(
    128,                       # x = constant 128 (single cell)
    128,                       # y = constant 128
    60,                        # z = constant 60
    (-data.inp.path[0] * 3,    # vx from –3*dx
     +data.inp.path[0] * 3, 10),  #    to +3*dx in 10 steps
    0,                         # vy = 0
    (-data.inp.path[0] * 3,    # vz from –3*dx
     0, 10),                   #    to 0 in 10 steps
    ispec=0,                   # electron species
)

# 2) Plot the resulting probability on the vx–vz plane
ax = probability_result.vxvz.plot(shading="auto", cmap="plasma")
ax.set_title("Probability Distribution: vx vs vz (ispec=0)")
plt.show()
```


#### 3. Backtrace using the particles from a previous probability calculation, then plot $x$–$z$ trajectories with probability as transparency

In this final example, we take the `particles` array produced internally by `get_probabilities(...)`, run backtraces on each of those particles, and then plot the $x$–$z$ projections of all backtraced trajectories.

 We normalize each trajectory’s probability to the maximum probability across all phase‐grid cells, and pass that normalized array to `alpha`, so that high‐probability trajectories appear more opaque and low‐probability trajectories more transparent.

```python
# 1) Reuse the ProbabilityResult from Example 2:
probability_result = data.backtrace.get_probabilities(
    128,
    128,
    60,
    (-data.inp.path[0] * 3, data.inp.path[0] * 3, 10),
    0,
    (-data.inp.path[0] * 3, 0, 10),
    ispec=0,
)

# 2) Extract the `particles` array and their associated `probabilities`
particles = probability_result.particles        # Sequence of Particle objects
prob_flat  = probability_result.probabilities    # 2D array of shape (nvz, nvx)

# 3) Flatten the 2D probability grid back into the 1D array matching `particles` order
prob_1d = prob_flat.ravel()

# 4) Normalize probabilities to [0,1] by dividing by the global maximum
alpha_values = prob_1d / prob_1d.max()

# 5) Compute backtraces for all particles
backtrace_result = data.backtrace.get_backtraces_from_particles(
    particles,
    ispec=0,
)

# 6) Plot x vs z for every trajectory, using the normalized probabilities as alpha
ax = backtrace_result.xz.plot(alpha=alpha_values)
ax.set_title("Backtrace Trajectories (x vs z) with Probability Transparency")
plt.show()
```


**Notes on the above examples:**

* In Example 1, `get_backtraces(positions, velocities)` returns a `MultiBacktraceResult` whose `xy` property is a `MultiXYData` object. You can sample, reorder, or subset the trajectories and then call `.plot()` on `.xy`, `.vxvy`, `.xz`, etc.

* In Example 2, `get_probabilities(...)` returns a `ProbabilityResult` whose `.vxvz`, `.xy`, `.xz`, etc. are all `HeatmapData` objects. Calling `.plot()` on any of these displays a 2D probability heatmap for the chosen pair of axes.

* In Example 3, `probability_result.particles` is the list of `Particle` objects used internally to compute the 6D probability grid. We pass that list to `get_backtraces_from_particles(...)` to compute backtraced trajectories for exactly those same particles. Normalizing their probabilities to `[0,1]` and passing that array into `alpha` makes high‐probability trajectories draw more opaquely.

These three patterns demonstrate the flexibility of the `data.backtrace` facade for:

1. **Direct backtracing** from arbitrary $(\mathbf{r}, \mathbf{v})$ arrays,
2. **Probability‐space calculations** on a structured phase grid, and
3. **Combining the two** so that you can visualize backtraced trajectories with opacity weighted by their computed probabilities.

</details>

---

Feel free to explore the documentation for more details and advanced usage:

[emout Documentation](https://nkzono99.github.io/emout/)

Happy analyzing!
