---
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()
```

# Simulations

+++

Here we discuss strategies for managing simulations performed locally, and on other servers like CoCalc or on clusters like Kamiak.

+++

## Organizing Experiments

+++

In order to reliably keep track of results, I recommend organizing everything in terms of "experiments".  The original idea here is to allow easy comparison with real experimental results, but the method can be extended to purely numerical work.  To do this, we define two classes:

`Experiment`: This class should provide an interface to the problem.  It's main role should be to accept experimentally relevant parameters, and then produce an appropriate initial state that representing the experimental protocol.  I delegate from the `State` class to the `Experiment` to determine this like time-dependence of parameters (e.g. because these may be most naturally expressed in experimental units.)  It should provide the following methods:

* `Experiment.__init__(**kw)`: The constructor should only accept keyword arguments.  These arguments will be stored with their values and used to generate a unique directory name for checkpoints.  For this to work, there needs to be a one-to-one mapping between simulation runs and the parameters here.

* `Experiment.get_state()`: Should quickly return a valid `State` object, but quickly.  The data may be initialized using something like the Thomas Fermi approximation, but expensive minimization should not be performed as this will also be used by the `Simulation` class which will insert checkpoint data into this object.

* `Experiment.get_initial_state()`: This should perform the appropriate experimental protocol for initializing the experimental state - i.e. the correct state at time $t=0$.  Thus, appropriate minimizations should be performed.

* `Experiment.dir_name`: This should be a property that returns a valid directory name (without the base path) where the experimental data should be stored.  This name should encode all of the information passed to the constructor.  I organize my work first into a directory name with the same name as the class, then into parameter-specific subfolders.

Note that only the constructor should take arguments (unless arguments are used for testing or debugging).  This will be the reproducible path to generating data.  To support this, you should only use `Experiment.get_state()` or `Experiment.get_initial_state()` to obtain a state object.  Finally, the `Experiment` objects should be able to be made persistent and will generally be stored for reference along with the checkpoint data.

`State`: The state represents all of the actual physical quantities of interest, but should express these in the most natural way for the purpose of computation.  It is the role of the `Experiment` class to convert from experimental units to the units that `State` uses.  For my work, I may overload the default states slightly if modifications are needed, and generally pass the `Experiment` object to `State.experiment` through the `State` constructor so that I can delegate properties such as time-dependence of coupling constants or trapping potentials.

* `State.E_max`: Maximum energy representable in the basis.  (Used to determine the evolver timestep).
* `State.hbar`: Converts `E_max` to a time for use by the evolver.


In designing your problem, you should consider what parameters make sense to vary in a single `Experiment` object, and when it is better to define a new `Experiment` class altogether.

```{code-cell}
from gpe import utils, soc_soliton

reload(utils)
reload(soc_soliton)

experiment = soc_soliton.Experiment1()
sim = utils.Simulation(experiment=experiment, dt_t_unit=0.5)
sim.initialize()
sim.run()
```

```{code-cell}
%pylab inline --no-import-all
from gpe import utils, soc_soliton

reload(utils)
reload(soc_soliton)

experiment = soc_soliton.Experiment1()
sim = utils.Simulation(experiment=experiment, dt_t_unit=0.5)
plt.figure(figsize=(20, 5))
sim.view()
```

```{code-cell}
sim = utils.Simulation(
    dt_t_unit=0.5,
    dir_name="_data/Experiment1/Lxyz=(1000.0,)/Nxyz=(16384,)/Rx_TF=234.6/detuning_E_R=0.1/magnetic_field_gradient=1e-06/mu_Bs_mu_B=(0.5,0.0)/rabi_frequency_E_R=1.9/recoil_frequency_Hz=3683.8/states=((1,-1),(1,0))/t_SOC_on_ms=10/t_expand_ms=7/t_wait_ms=100/trapping_frequencies_Hz=(3,273,277)",
)
sim.initialize()

# state = expt.get_initial_state()
# a = Archive()
# a.insert(expt=expt)
# d = {}
# exec str(a) in d
# d['expt']
```

```{code-cell}
persist.archive.Archive?
```

# Visualization

+++

Here we consider visualizing a simulation.  First we generate some data:

```{code-cell}
sim.get_state(1.0)
```

```{code-cell}
import signal
import ipyparallel as ipp

rc = ipp.Client()
engine_pids = rc[:].apply(os.getpid).get_dict()


def signal_engine(engine_id, sig=signal.SIGINT):
    """send a signal to a local engine"""
    pid = engine_pids[engine_id]
    os.kill(pid, sig)
```

```{code-cell}
!rm -rf _data/ExperimentLoading/
```

```{code-cell}
%%px --noblock
from gpe import utils, soc_soliton

reload(utils)
reload(soc_soliton)

experiment = soc_soliton.ExperimentLoading(cells_x=100)
sim = utils.Simulation(dt_=1.0, experiment=experiment, dt_t_scale=0.1)
sim.run()
```

```{code-cell}
from gpe import utils, soc_soliton
```

```{code-cell}
reload(utils)
reload(soc_soliton)
experiment = soc_soliton.ExperimentLoading(cells_x=100)
sim = utils.Simulation(dt_=1.0, experiment=experiment, dt_t_scale=0.1)
```

```{code-cell}
def plot_state(state, fig=None):
    if fig is None:
        fig = plt.figure()
    plt.plot(state.xyz[0], state.get_density()[0])
    return fig
```

```{code-cell}
sim.view(plot_state=plot_state)
```

```{code-cell}
map(signal_engine, engine_pids)
```

```{code-cell}
s = sim.get_state(0)
s.experiment.trapping_frequencies_Hz
```

```{code-cell}
plt.plot(s.get_density()[1])
```

```{code-cell}
s.experiment.get_initial_state(psi_tol=1e-16).plot()
```

```{code-cell}
%pylab inline --no-import-all
from IPython.display import clear_output
import logging

from gpe import utils, soc_soliton

reload(utils)
reload(soc_soliton)

experiment = soc_soliton.ExperimentLoading()
sim = utils.Simulation(
    dt_=1.0, experiment=experiment, dt_t_scale=0.1, logging_level=logging.WARNING
)
```

```{code-cell}
import holoviews as hv
```

```{code-cell}
hv.extension("bokeh")
```

```{code-cell}
import attr

from collections import namedtuple


@attr.s
class Sim:
    sim = attr.ib()

    @property
    def ts(self):
        return self.sim.saved_ts_

    @property
    def t(self):
        return hv.Dimension("t", values=self.ts)

    @property
    def extents(self):
        xmin = 0
        nmin = np.inf
        xmax = nmax = -np.inf
        for t_ in self.ts:
            state = self.sim.get_state(t_)
            xmin = min(xmin, state.xyz[0].min())
            xmax = max(xmax, state.xyz[0].max())
            ns = state.get_density()
            nmin = min(nmin, ns.min())
            nmax = min(nmax, ns.max())
        return (xmin, xmax, nmin, nmax)

    def densities(self, t_):
        state = self.sim.get_state(t_)
        x = hv.Dimension("x", values=state.xyz[0])
        na, nb = state.get_density()
        return (
            hv.Curve(na, kdims=[x], vdims=["na"])
            * hv.Curve(nb, kdims=[x], vdims=["nb"])
            * hv.Curve(na + nb, kdims=[x], vdims=["n"])
        )

    def plots(self, t_):
        state = self.sim.get_state(t_)
        state.plot()
        fig = plt.gcf()
        return Figure(fig)


class Figure(hv.Element2D):
    pass


class FigurePlot(hv.plotting.mpl.ElementPlot):
    def initialize_plot(self, ranges=None):
        element = self.hmap.last
        self.handles["fig"] = element.data
        return self.handles["fig"]

    def update_handles(self, key, axis, element, ranges, style):
        self.handles["fig"] = element.data

    def update_frame(self, key, ranges=None, element=None):
        element = self._get_frame(key)
        self.handles["fig"] = element.data
        return self.handles["fig"]


hv.Store.register({Figure: FigurePlot}, "matplotlib")
```

```{code-cell}
hv.plotting.mpl.ElementPlot.update_handles?
```

```{code-cell}
#%%output size=150
#%%opts Curve  [height=100 width=400]

s = Sim(sim)
# hv.DynamicMap(s.densities, kdims=[s.t])
hv.DynamicMap(s.plots, kdims=[s.t])
```

```{code-cell}
def plot_state(state, fig=None):
    if fig is None:
        fig = plt.figure()
    state.plot_k()
    return fig


plot_state = None
sim.view(plot_state=plot_state)
plt.close("all")
```

```{code-cell}
import numpy as np

frequencies = [0.5, 0.75, 1.0, 1.25]


def sine_curve(phase, freq):
    xvals = [0.1 * i for i in range(100)]
    return hv.Curve((xvals, [np.sin(phase + freq * x) for x in xvals]))


# When run live, this cell's output should match the behavior of the GIF below
dmap = hv.DynamicMap(sine_curve, kdims=["phase", "frequency"])
dmap.redim.range(phase=(0.5, 1)).redim.range(frequency=(0.5, 1.25))
```