#+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.τ, e_sys.real, yerr=σ_e_sys.real, ecolor="yellow")
      plt.ylabel(r"$\langle H_S\rangle$")
      plt.xlabel(r"$τ$")
      fs.export_fig("system_energy")
#+end_src

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

** 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 | 0x7fee23b5ebb0> |
[[file:./.ob-jupyter/0775a369159dac940ed20b03310e57d003ec7c42.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 // 4,
      save="results/flow_more.npy",
      n_proc=4
  )

  with fs.hiro_style():
      fig, ax = fs.plot_convergence(result.τ, [full_flow[1], full_flow[-1]], 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:
[[file:./.ob-jupyter/4baa612c7e7201f51f74886c2f74caa2a56383a8.svg]]


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/fd2493973fc2770c764cc73be94f2513def80b8e.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/64d9fac11c4d3356642e6e585b0ceb4b22d7e8ad.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:

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/48b4df7cf090fe5c307194ad23cdb51da1538247.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/c9msmd5k6clygvhbasl6871wb2ldg58c-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 | 0x7fee23c017f0> |
[[file:./.ob-jupyter/2bb493ebb9eb6bc714253b115f92b26457542e60.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

* Observations
- convergence slower
- error estimate is of course too small
- energy change not as steep for smaller \(\omega_c\)

* Temperature
#+begin_src jupyter-python
  ρ = result.hd.rho_t_accum.mean[-1]
  ρ
#+end_src

#+RESULTS:
: array([[ 0.65441242+0.j        , -0.00319444+0.00213465j],
:        [-0.00319444-0.00213465j,  0.34558758+0.j        ]])

#+begin_src jupyter-python :results none
  v, w = np.linalg.eig(params.SysP.H_sys)
  v = np.sort(v.real)
  v
#+end_src

#+begin_src jupyter-python :results none
  def thermal_e(v: np.ndarray, β: float):
        return np.sum(v[None, :] * np.exp(-β * v[None, :]), axis=1) / np.sum(np.exp(-β * v[None,:]), axis=1)

  def trace_norm(ρ: np.ndarray):
        return np.trace(scipy.linalg.sqrtm(ρ @ ρ.conj().T)).real
#+end_src

#+begin_src jupyter-python :results none
    e_final = e_sys[-1]
#+end_src

#+begin_src jupyter-python :results none
  def therm_state(T):
      ρ_therm = scipy.linalg.expm(-params.SysP.H_sys/T)
      ρ_therm = ρ_therm / np.trace(ρ_therm)

      return ρ_therm
#+end_src

#+begin_src jupyter-python :results none
  def dist_to_therm(T):
      return trace_norm(ρ - therm_state(T))
#+end_src

#+begin_src jupyter-python
    import scipy.optimize
    closest_temp = scipy.optimize.minimize(dist_to_therm, (params.SysP.__non_key__["T"]))
    closest_temp.x, (params.SysP.__non_key__["T"]), closest_temp.fun
#+end_src

#+RESULTS:
| array | ((1.56619123)) | 1.5 | 0.007684057062556999 |


#+begin_src jupyter-python
  dist_to_therm(closest_temp.x), dist_to_therm(params.SysP.__non_key__["T"])
#+end_src

#+RESULTS:
| 0.007684057062556999 | 0.01483332959261369 |

#+begin_src jupyter-python
  ρ - therm_state(closest_temp.x)
#+end_src

#+RESULTS:
: array([[ 3.31136851e-09+0.j        , -3.19443907e-03+0.00213465j],
:        [-3.19443907e-03-0.00213465j, -3.31136663e-09+0.j        ]])

How does the effective Hamiltonian look like?
#+begin_src jupyter-python
  def effective_objective(H):
      H = H.reshape((2,2))
      H_eff = H + H.conj().T
      T = params.SysP.__non_key__["T"]
      Z = np.trace(np.expm(-H_eff/T))
      return scipy.linalg.logm(Z * ρ) + H_eff/T

  scipy.optimze.newton()
#+end_src