#+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 = 5
  s = 1

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

  integration = IntP(t_max=30, t_steps=int(30 // 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.8,
      T=0.1,
  )

  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=g, w=w, p=1, q=0.5, kfac=1.4
          ),
          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-4,
          intpl_tol=1e-4,
          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-4,
        intpl_tol=1e-4,
        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
#+end_src

** 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
#+end_src

#+RESULTS:
: 2000

** 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/energy.npy",
  )

  plt.errorbar(result.τ, e_sys.real, yerr=σ_e_sys.real, ecolor="yellow")
#+end_src

#+RESULTS:
:RESULTS:
: <ErrorbarContainer object of 3 artists>
[[file:./.ob-jupyter/4380c3cff03d9ecce002e7639da95347ebeb7e4b.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

  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 | 0x7fb51d6ebeb0> |
[[file:./.ob-jupyter/784bb9515cdf0d7723991b501cf27baed17cadc6.svg]]
:END:

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

  with ut.hiro_style():
      _, ax = ut.plot_convergence(result.τ, full_flow, transform=lambda y: -y)
      ax.legend()
#+end_src

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


We can integrate the energy change in the bath:
#+begin_src jupyter-python
  for i in range(len(full_flow)):
      e_bath = util.integrate_array(-full_flow[i][1], result.τ)
      plt.plot(result.τ, e_bath)
#+end_src

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

** Calculate the Interaction Energy
First we calculate it from energy conservation.
#+begin_src jupyter-python
  e_int = (1/2 - e_sys - e_bath).real
  plt.plot(result.τ, e_int)
#+end_src

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

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

#+RESULTS:

And both together:
#+begin_src jupyter-python
  plt.plot(result.τ, e_int, label="integrated")
  plt.plot(result.τ, e_int_ex, label="exact")
  plt.legend()
#+end_src

#+RESULTS:
:RESULTS:
: <matplotlib.legend.Legend at 0x7fb51c4d0b50>
[[file:./.ob-jupyter/31340f5271d2e0755bc77033d88f6a385465a5dc.svg]]
:END:

Seems to work :P.

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