#+PROPERTY: header-args :session finite_temp_new :kernel python :pandoc yes :async yes

* Configuration and Setup
This will be tangled into the [[file:config.py][config file]] that can be used with the HOPS cli.

#+begin_src jupyter-python :results none :tangle config.py
  from hops.core.hierarchy_parameters import HIParams, HiP, IntP, SysP, ResultType
  from hops.core.hierarchyLib import HI
  from hops.util.bcf_fits import get_ohm_g_w
  from hops.util.truncation_schemes import TruncationScheme_Power_multi
  import hops.util.bcf
  import numpy as np
  import hops.util.matrixLib as ml
  from stocproc import StocProc_FFT

  wc = 2
  s = 1

  # The BCF fit
  bcf_terms = 5
  g, w = get_ohm_g_w(bcf_terms, s, wc)

  integration = IntP(t_max=30, t_steps=int(20 // 0.01))
  system = SysP(
      H_sys=0.5 * np.array([[-1, 0], [0, 1]]),
      L=0.5 * np.array([[0, 1], [1, 0]]),
      psi0=np.array([0, 1]),
      g=g,
      w=w,
      bcf_scale=0.5,
      T=1.5,
  )

  params = HIParams(
      SysP=system,
      IntP=integration,
      HiP=HiP(
          nonlinear=True,
          normalized_by_hand=True,
          result_type=ResultType.ZEROTH_AND_FIRST_ORDER,
          truncation_scheme=TruncationScheme_Power_multi.from_g_w(
              g=.5 * g, w=w, p=1, q=0.5, kfac=1.7
          ),
          save_therm_rng_seed=True,
      ),
      Eta=StocProc_FFT(
          spectral_density=hops.util.bcf.OhmicSD_zeroTemp(
              s,
              1,
              wc,
          ),
          alpha=hops.util.bcf.OhmicBCF_zeroTemp(
              s,
              1,
              wc,
          ),
          t_max=integration.t_max,
          intgr_tol=1e-5,
          intpl_tol=1e-5,
          negative_frequencies=False,
      ),
      EtaTherm=StocProc_FFT(
        spectral_density=hops.util.bcf.Ohmic_StochasticPotentialDensity(
            s, 1, wc, beta=1 / system.__non_key__["T"]
        ),
        alpha=hops.util.bcf.Ohmic_StochasticPotentialCorrelations(
            s, 1, wc, beta=1 / system.__non_key__["T"]
        ),
        t_max=integration.t_max,
        intgr_tol=1e-5,
        intpl_tol=1e-5,
        negative_frequencies=False,
    ),
  )
#+end_src

* Using the Data
** Jupyter Setup
#+begin_src jupyter-python :results none
  import numpy as np
  import matplotlib.pyplot as plt
  import utilities as ut
  import figsaver as fs
#+end_src


Let's export some infos about the model to TeX.
#+begin_src jupyter-python
  fs.tex_value(system.bcf_scale, prec=1, save="bcf_scale", prefix="η="), fs.tex_value(
      wc, prec=0, save="cutoff_freq", prefix="ω_c="
  ), fs.tex_value(system.__non_key__["T"], prec=1, save="temp", prefix="T=")
#+end_src

#+RESULTS:
| \(η=0.5\) | \(ω_c=2\) | \(T=1.5\) |

** Load the Data
#+begin_src jupyter-python :results none
  from hopsflow import hopsflow, util
  from hops.core.hierarchyData import HIMetaData
#+end_src

Now we read the trajectory data.
#+begin_src jupyter-python
  class result:
      hd = HIMetaData("data", ".").get_HIData(params, read_only=True)
      N = hd.samples
      τ = hd.get_time()
      ψ_1 = hd.aux_states
      ψ = hd.stoc_traj
      seeds = hd.rng_seed

  result.N
  fs.tex_value(result.N, prefix="N=", save="samples")
#+end_src

#+RESULTS:
: \(N=10000\)

** Calculate System Energy
Simple sanity check.
#+begin_src jupyter-python
  _, e_sys, σ_e_sys = util.operator_expectation_ensemble(
      iter(result.ψ),
      system.H_sys,
      result.N,
      params.HiP.nonlinear,
      save="./results/new_energy1.npy",
  )

  with fs.hiro_style():
      plt.gcf().set_size_inches(fs.get_figsize(239, 1, .8))
      plt.errorbar(result.τ[::100], e_sys.real[::100], yerr=σ_e_sys.real[::100], ecolor="yellow")
      plt.ylabel(r"$\langle H_S\rangle$")
      plt.xlabel(r"$τ$")
      fs.export_fig("system_energy")
#+end_src

#+RESULTS:
:RESULTS:
: 100% 9999/9999 [00:09<00:00, 1009.10it/s]
[[file:./.ob-jupyter/6ca9a74ecaeb4a206f93f97921ee4a4bb212254f.svg]]
:END:

** Calculate the Heat Flow
Now let's calculate the heatflow. In this simple case it is engouh to
know the first hierarchy states.

First we set up some parameter objects for the alogrithm.
#+begin_src jupyter-python :results none
  hf_system = hopsflow.SystemParams(
      system.L, system.g, system.w, system.bcf_scale, params.HiP.nonlinear
  )

  η = params.Eta
  ξ = params.EtaTherm
  ξ.calc_deriv = True
  ξ.set_scale(params.SysP.bcf_scale)

  hf_therm = hopsflow.ThermalParams(ξ=ξ, τ=result.τ, num_deriv=False)
#+end_src

Now we can apply our tooling to one trajectory for testing.
#+begin_src jupyter-python
  hf_sample_run = hopsflow.HOPSRun(result.ψ[0], result.ψ_1[0], hf_system)
  hf_sample_run_therm = hopsflow.ThermalRunParams(hf_therm, result.seeds[0])
  first_flow = hopsflow.flow_trajectory_coupling(hf_sample_run, hf_system)
  first_flow_therm = hopsflow.flow_trajectory_therm(hf_sample_run, hf_sample_run_therm)
  plt.plot(result.τ, first_flow)
  plt.plot(result.τ, first_flow_therm)
#+end_src

#+RESULTS:
:RESULTS:
| <matplotlib.lines.Line2D | at | 0x7f8956fadfd0> |
[[file:./.ob-jupyter/0c5b760af32ab0310202f083aaa6c54692a8f8dc.svg]]
:END:

And now for all trajectories.
#+begin_src jupyter-python :results none
  full_flow = hopsflow.heat_flow_ensemble(
      iter(result.ψ),
      iter(result.ψ_1),
      hf_system,
      result.N,
      (iter(result.seeds), hf_therm),
      every=result.N // 4,
      save="results/flow_more.npy",
      n_proc=4
  )

  with fs.hiro_style():
      fig, ax = fs.plot_convergence(result.τ, full_flow, transform=lambda y: -y)
      fig.set_size_inches(fs.get_figsize(239, 1, .8))
      ax.legend()
      ax.set_xlabel("$τ$")
      ax.set_ylabel("$-J$")
      #fs.export_fig("flow", fig)
#+end_src

#+RESULTS:
:RESULTS:
:   1% 90/9999 [00:21<39:31,  4.18it/s]Process ForkPoolWorker-173:
: Process ForkPoolWorker-172:
# [goto error]
#+begin_example
  ---------------------------------------------------------------------------
  IndexError                                Traceback (most recent call last)
  /nix/store/dn4fwp0yx6nsa85cr20cwvdmg64xwmcy-python3-3.9.9/lib/python3.9/multiprocessing/pool.py in next(self, timeout)
      852             try:
  --> 853                 item = self._items.popleft()
      854             except IndexError:

  IndexError: pop from an empty deque

  During handling of the above exception, another exception occurred:

  KeyboardInterrupt                         Traceback (most recent call last)
  /tmp/ipykernel_94001/1430602402.py in <module>
  ----> 1 full_flow = hopsflow.heat_flow_ensemble(
        2     iter(result.ψ),
        3     iter(result.ψ_1),
        4     hf_system,
        5     result.N,

  /nix/store/c9msmd5k6clygvhbasl6871wb2ldg58c-python3-3.9.9-env/lib/python3.9/site-packages/hopsflow/hopsflow.py in heat_flow_ensemble(ψ_0s, ψ_1s, params, N, therm_args, only_therm, **kwargs)
      357         raise ValueError("Can't calculate only thermal part if therm_args are None.")
      358 
  --> 359     return util.ensemble_mean(
      360         iter(zip(ψ_0s, ψ_1s, therm_args[0]))
      361         if therm_args

  /nix/store/c9msmd5k6clygvhbasl6871wb2ldg58c-python3-3.9.9-env/lib/python3.9/site-packages/hopsflow/util.py in ensemble_mean(arg_iter, function, N, const_args, const_kwargs, n_proc, every, save)
      221         )
      222 
  --> 223         for res in tqdm(result_iter, total=(N - 1) if N else None):
      224             aggregate.update(res)
      225 

  /nix/store/c9msmd5k6clygvhbasl6871wb2ldg58c-python3-3.9.9-env/lib/python3.9/site-packages/tqdm/std.py in __iter__(self)
     1178 
     1179         try:
  -> 1180             for obj in iterable:
     1181                 yield obj
     1182                 # Update and possibly print the progressbar.

  /nix/store/dn4fwp0yx6nsa85cr20cwvdmg64xwmcy-python3-3.9.9/lib/python3.9/multiprocessing/pool.py in <genexpr>(.0)
      446                     result._set_length
      447                 ))
  --> 448             return (item for chunk in result for item in chunk)
      449 
      450     def apply_async(self, func, args=(), kwds={}, callback=None,

  /nix/store/dn4fwp0yx6nsa85cr20cwvdmg64xwmcy-python3-3.9.9/lib/python3.9/multiprocessing/pool.py in next(self, timeout)
      856                     self._pool = None
      857                     raise StopIteration from None
  --> 858                 self._cond.wait(timeout)
      859                 try:
      860                     item = self._items.popleft()

  /nix/store/dn4fwp0yx6nsa85cr20cwvdmg64xwmcy-python3-3.9.9/lib/python3.9/threading.py in wait(self, timeout)
      310         try:    # restore state no matter what (e.g., KeyboardInterrupt)
      311             if timeout is None:
  --> 312                 waiter.acquire()
      313                 gotit = True
      314             else:

  KeyboardInterrupt:
#+end_example
:END:


We can integrate the energy change in the bath:
#+begin_src jupyter-python
  import scipy.integrate

  e_bath = np.array([0] + [
      scipy.integrate.simpson(-full_flow[-1][1][:i], result.τ[:i])
      for i in range(1, len(result.τ))
  ])
  plt.plot(result.τ, e_bath)
  σ_e_bath = np.sqrt(np.array([0] + [
      scipy.integrate.simpson(full_flow[-1][2][:i]**2, result.τ[:i])
      for i in range(1, len(result.τ))
  ])).real
  plt.errorbar(result.τ, e_bath, yerr=σ_e_bath, ecolor="yellow")
#+end_src

#+RESULTS:
:RESULTS:
: <ErrorbarContainer object of 3 artists>
[[file:./.ob-jupyter/a65a13edc562621973be18c8fbd4271bf3afa733.svg]]
:END:

** Calculate the Interaction Energy
First we calculate it from energy conservation.
#+begin_src jupyter-python
  e_int = (1/2 - e_sys - e_bath).real
  σ_e_int = np.sqrt(σ_e_sys ** 2 + σ_e_bath ** 2).real
  plt.errorbar(result.τ, e_int, yerr=σ_e_int, ecolor="yellow")
#+end_src

#+RESULTS:
:RESULTS:
: <ErrorbarContainer object of 3 artists>
[[file:./.ob-jupyter/5950975b59e103671be041250987da84adfdca83.svg]]
:END:

And then from first principles:
#+begin_src jupyter-python
  _, e_int_ex, σ_e_int_ex = hopsflow.interaction_energy_ensemble(
      result.ψ,
      result.ψ_1,
      hf_system,
      result.N,
      (result.seeds, hf_therm),
      save="results/interaction_energy_new.npy",
  )
#+end_src

#+RESULTS:
: 100% 499/499 [02:09<00:00,  3.86it/s]

And both together:
#+begin_src jupyter-python
  with fs.hiro_style():
      plt.errorbar(result.τ, e_int, yerr=σ_e_int, label="from energy conservation", ecolor="yellow")
      plt.errorbar(result.τ, e_int_ex, yerr=σ_e_int_ex, label="direct", ecolor="pink")
      plt.gcf().set_size_inches(fs.get_figsize(239, 1, .8))
      plt.legend()
      plt.ylabel(r"$\langle H_I\rangle$")
      plt.xlabel(r"$τ$")
      #fs.export_fig("interaction")
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/ecbcd0497b02327dd0a678f193edc9c3d9e7c553.svg]]

Seems to work :P.
#+begin_src jupyter-python
  plt.plot(result.τ, 1/2 - np.array(e_bath) - e_int_ex )
  plt.plot(result.τ, e_sys)
#+end_src

#+RESULTS:
:RESULTS:
: /nix/store/maj240ris2x1a0r1wz26m0m785bc0xn9-python3-3.9.9-env/lib/python3.9/site-packages/matplotlib/cbook/__init__.py:1298: ComplexWarning: Casting complex values to real discards the imaginary part
:   return np.asarray(x, float)
| <matplotlib.lines.Line2D | at | 0x7f15e986fbb0> |
[[file:./.ob-jupyter/8161e70494da51689f01e45a2beb0d57798a7436.svg]]
:END:

* Close the Data File
We need to release the hold on the file.
#+begin_src jupyter-python :results none
 result.hd.close()
#+end_src