---
jupytext:
  formats: md:myst,ipynb
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.16.4
kernelspec:
  display_name: Python 3 (ipykernel)
  language: python
  name: python3
---

```{code-cell}
:init_cell: true

import mmf_setup

mmf_setup.nbinit()
from gpe.imports import *
```

# Plotting

+++

Here we demonstrate some tools for plotting, animating results, and making movies.

+++

## [MPLGrid](../gpe/plot_utils.py)

+++

Here as a simple example, we start with a harmonically trapped gas, phase imprint a dark soliton, and then watch the "evolution" as we cool into the soliton state.

```{code-cell}
from gpe.bec import State, u

s = State(Nxyz=(128,), Lxyz=(20.0,))
s[:] *= np.sign(s.xyz[0])
s.cooling_phase = 1j
e = EvolverABM(s, dt=0.1 * s.t_scale)


def plot(state):
    plt.plot(state.xyz[0] / u.micron, state.get_density())


states = [s]
with NoInterrupt(ignore=True) as interrupted:
    for n in range(6):
        e.evolve(200)
        states.append(e.get_y())
        plt.clf()
        plot(e.y)
        display(plt.gcf())
        clear_output(wait=True)
```

Our first demonstration is to stack these frames using MPLGrid.  We can stack going down or going right.

+++

### `direction='down'`

```{code-cell}
import gpe.plot_utils

reload(gpe.plot_utils)
from gpe.plot_utils import MPLGrid

plt.figure(figsize=(10, 5))
grid = MPLGrid(direction="down", share=True)
for s in states:
    ax = grid.next()
    ax.plot(s.xyz[0] / u.micron, s.get_density())
plt.tight_layout()
```

### `direction='right'`

+++

The plots can also be stacked to the right:

```{code-cell}
plt.figure(figsize=(15, 2))
grid = MPLGrid(direction="right", share=True)
for s in states:
    ax = grid.next()
    ax.plot(s.xyz[0] / u.micron, s.get_density())
plt.tight_layout()
```

### `right=True`

```{code-cell}
plt.figure(figsize=(15, 2))
grid = MPLGrid(direction="right", share=True, right=True)
for s in states:
    ax = grid.next()
    ax.plot(s.xyz[0] / u.micron, s.get_density())
plt.tight_layout()
```

### `grid()`

+++

More complex layouts can be created using subgrids:

```{code-cell}
plt.figure(figsize=(10, 3))
grid = MPLGrid(direction="down", share=True)
for s in states[::3]:
    ax = grid.next(size=1.0 / 3)  # Make these 1/3 of the size
    ax.plot(s.xyz[0] / u.micron, s.get_density())
subgrid = grid.grid(size=1, direction="right", share=True)
for s in states:
    ax = subgrid.next(2)
    ax.plot(s.xyz[0] / u.micron, s.get_density())

plt.tight_layout()
```

Note in this case that we obscure the x-axis of the upper plots.  This is a limitation of the current MPLGrid class - all axes must either be shared or not shared.  To add a space, the first three panes should be in their own subgrid:

```{code-cell}
reload(gpe.plot_utils)
from gpe.plot_utils import MPLGrid
from matplotlib.gridspec import GridSpecFromSubplotSpec

plt.figure(figsize=(10, 3))
grid = MPLGrid(direction="down", share=False, space=0.5)
subgrid = grid.grid(direction="down", share=True)
for s in states[::3]:
    ax = subgrid.next(1.0 / 2)  # Make these 1/2 of the size
    ax.plot(s.xyz[0] / u.micron, s.get_density())
ax.set_xlabel("x")

subgrid = grid.grid(
    direction="right", share=True, space=0.1
)  # Some space might help lables not clash
for s in states:
    ax = subgrid.next(2)
    ax.plot(s.xyz[0] / u.micron, s.get_density())

plt.tight_layout()
```

Note that `tight_layout()` does not work properly here.  Not sure why.  Probably the solution is to use the new [constrained layout](https://matplotlib.org/devdocs/tutorials/intermediate/constrainedlayout_guide.html) feature of matplotlib once it is settled.