master-thesis/python/richard_hops/energy_flow_thermal.org

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

* Setup
** Jupyter
#+begin_src jupyter-python
  %load_ext autoreload
  %autoreload 2
  %load_ext jupyter_spaces
#+end_src

#+RESULTS:

** Matplotlib
#+begin_src jupyter-python
  import matplotlib
  import matplotlib.pyplot as plt

  #matplotlib.use("TkCairo", force=True)
  %gui tk
  %matplotlib inline
  plt.style.use('ggplot')
#+end_src

#+RESULTS:

** Richard (old) HOPS
#+begin_src jupyter-python
  import hierarchyLib
  import hierarchyData
  import numpy as np

  from stocproc.stocproc import StocProc_FFT
  import bcf
  from dataclasses import dataclass, field
  import scipy
  import scipy.misc
  import scipy.signal
  import pickle
  from scipy.special import gamma as gamma_func
  from scipy.optimize import curve_fit
#+end_src

#+RESULTS:

** Auxiliary Definitions
#+begin_src jupyter-python
  σ1 = np.matrix([[0,1],[1,0]])
  σ2 = np.matrix([[0,-1j],[1j,0]])
  σ3 = np.matrix([[1,0],[0,-1]])
#+end_src

#+RESULTS:

* Model Setup
Basic parameters.
#+begin_src jupyter-python
  class params:
      T = 2.09

      t_max = 10
      t_steps = int(t_max * 1/.01)
      k_max = 3
      N = 800

      seed = 100
      dim = 2
      H_s = σ3 + np.eye(dim)
      L = 1 / 2 * (σ1 - 1j * σ2)
      ψ_0 = np.array([1, 0])

      s = 1
      num_exp_t = 5

      wc = 1

      with open("good_fit_data_abs_brute_force", "rb") as f:
          good_fit_data_abs = pickle.load(f)

      alpha = 0.8
      # _, g_tilde, w_tilde = good_fit_data_abs[(numExpFit, s)]
      # g_tilde = np.array(g_tilde)
      # w_tilde = np.array(w_tilde)
      # g = 1 / np.pi * gamma_func(s + 1) * wc ** (s + 1) * np.asarray(g_tilde)
      # w = wc * np.asarray(w_tilde)

      bcf_scale = np.pi / 2 * alpha * wc ** (1 - s)
#+end_src

#+RESULTS:

** BCF and Thermal BCF
#+begin_src jupyter-python
  @dataclass
  class CauchyBCF:
      δ: float
      wc: float

      def I(self, ω):
          return np.sqrt(self.δ) / (self.δ + (ω - self.wc) ** 2 / self.δ)

      def __call__(self, τ):
          return np.sqrt(self.δ) * np.exp(-1j * self.wc * τ - np.abs(τ) * self.δ)

      def __bfkey__(self):
          return self.δ, self.wc

  α = bcf.OBCF(s=params.s, eta=1, gamma=params.wc)
  I = bcf.OSD(s=params.s, eta=1, gamma=params.wc)
#+end_src

#+RESULTS:

*** Fit
We now fit a sum of exponentials against the BCF.
#+begin_src jupyter-python
  from lmfit import minimize, Parameters


  def α_apprx(τ, g, w):
      return np.sum(
          g[np.newaxis, :] * np.exp(-w[np.newaxis, :] * (τ[:, np.newaxis])), axis=1
      )


  def _fit():
      def residual(fit_params, x, data):
          resid = 0
          w = np.array([fit_params[f"w{i}"] for i in range(params.num_exp_t)]) + 1j * np.array(
              [fit_params[f"wi{i}"] for i in range(params.num_exp_t)]
          )
          g = np.array([fit_params[f"g{i}"] for i in range(params.num_exp_t)]) + 1j * np.array(
              [fit_params[f"gi{i}"] for i in range(params.num_exp_t)]
          )
          resid = data - α_apprx(x, g, w)

          return resid.view(float)

      fit_params = Parameters()
      for i in range(params.num_exp_t):
          fit_params.add(f"g{i}", value=.1)
          fit_params.add(f"gi{i}", value=.1)
          fit_params.add(f"w{i}", value=.1)
          fit_params.add(f"wi{i}", value=.1)

      ts = np.linspace(0, params.t_max, 1000)
      out = minimize(residual, fit_params, args=(ts, α(ts)))

      w = np.array([out.params[f"w{i}"] for i in range(params.num_exp_t)]) + 1j * np.array(
          [out.params[f"wi{i}"] for i in range(params.num_exp_t)]
      )
      g = np.array([out.params[f"g{i}"] for i in range(params.num_exp_t)]) + 1j * np.array(
          [out.params[f"gi{i}"] for i in range(params.num_exp_t)]
      )

      return w, g


  w, g = _fit()
#+end_src

#+RESULTS:

*** Plot
Let's look a the result.
#+begin_src jupyter-python
  class bcfplt:
      t = np.linspace(0, params.t_max, 1000)
      ω = np.linspace(params.wc - 10, params.wc + 10, 1000)
      fig, axs = plt.subplots(2)
      axs[0].plot(t, np.real(α(t)))
      axs[0].plot(t, np.imag(α(t)))
      axs[0].plot(t, np.real(α_apprx(t, g, w)))
      axs[0].plot(t, np.imag(α_apprx(t, g, w)))
      axs[1].plot(ω, I(ω).real)
      axs[1].plot(ω, I(ω).imag)
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/4477b58eeacab016f06c0cd5e97fbca4343ad8b7.png]]

Seems ok for now.
** Hops setup
#+begin_src jupyter-python
  HierachyParam = hierarchyData.HiP(
      k_max=params.k_max,
      # g_scale=None,
      # sample_method='random',
      seed=params.seed,
      nonlinear=True,
      normalized=False,
      # terminator=False,
      result_type=hierarchyData.RESULT_TYPE_ALL,
      # accum_only=None,
      # rand_skip=None
  )
#+end_src

#+RESULTS:

Integration.
#+begin_src jupyter-python
  IntegrationParam = hierarchyData.IntP(
      t_max=params.t_max,
      t_steps=params.t_steps,
      # integrator_name='zvode',
      # atol=1e-8,
      # rtol=1e-8,
      # order=5,
      # nsteps=5000,
      # method='bdf',
      # t_steps_skip=1
  )
#+end_src

#+RESULTS:

And now the system.
#+begin_src jupyter-python
  SystemParam = hierarchyData.SysP(
      H_sys=params.H_s,
      L=params.L,
      psi0=params.ψ_0,  # excited qubit
      g=np.array(g),
      w=np.array(w),
      H_dynamic=[],
      bcf_scale=params.bcf_scale,  # some coupling strength (scaling of the fit parameters 'g_i')
      gw_hash=None,  # this is used to load g,w from some database
      len_gw=len(g),
  )
#+end_src

#+RESULTS:

The quantum noise.
#+begin_src jupyter-python
  Eta = StocProc_FFT(
      I,
      params.t_max,
      α,
      negative_frequencies=False,
      seed=params.seed,
      intgr_tol=1e-2,
      intpl_tol=1e-2,
      scale=params.bcf_scale,
  )
#+end_src

#+RESULTS:
#+begin_example
  stocproc.stocproc - INFO - non neg freq only
  stocproc.method_ft - INFO - get_dt_for_accurate_interpolation, please wait ...
  stocproc.method_ft - INFO - acc interp N 33 dt 6.22e-01 -> diff 1.50e-01
  stocproc.method_ft - INFO - acc interp N 65 dt 3.11e-01 -> diff 3.37e-02
  stocproc.method_ft - INFO - acc interp N 129 dt 1.55e-01 -> diff 7.11e-03
  stocproc.method_ft - INFO - requires dt < 1.555e-01
  stocproc.method_ft - INFO - get_N_a_b_for_accurate_fourier_integral, please wait ...
  stocproc.method_ft - INFO - J_w_min:1.00e-02 N 32 yields: interval [0.00e+00,6.47e+00] diff 9.83e-03
  stocproc.method_ft - INFO - return, cause tol of 0.01 was reached
  stocproc.method_ft - INFO - requires dx < 2.023e-01
  stocproc.stocproc - INFO - Fourier Integral Boundaries: [0.000e+00, 4.575e+01]
  stocproc.stocproc - INFO - Number of Nodes            : 256
  stocproc.stocproc - INFO - yields dx                  : 1.787e-01
  stocproc.stocproc - INFO - yields dt                  : 1.373e-01
  stocproc.stocproc - INFO - yields t_max               : 3.502e+01
#+end_example

The sample trajectories are smooth.
#+begin_src jupyter-python
  %%space plot
  ts = np.linspace(0, params.t_max, 1000)
  Eta.new_process()
  plt.plot(ts, Eta(ts).real)
#+end_src

#+RESULTS:
:RESULTS:
| <matplotlib.lines.Line2D | at | 0x7fc2bb9e11f0> |
[[file:./.ob-jupyter/5142d39238c0ae10debc12c41c8a3257029e427b.png]]
:END:

And now the thermal noise.
#+begin_src jupyter-python
  EtaTherm = StocProc_FFT(
      spectral_density=bcf.OFTDens(s=params.s, eta=1, gamma=params.wc, beta=1 / params.T),
      t_max=params.t_max,
      alpha=bcf.OFTCorr(s=params.s, eta=1, gamma=params.wc, beta=1 / params.T),
      intgr_tol=1e-3,
      intpl_tol=1e-3,
      seed=params.seed,
      negative_frequencies=False,
      scale=params.bcf_scale,
  )
#+end_src

#+RESULTS:
#+begin_example
  stocproc.stocproc - INFO - non neg freq only
  stocproc.method_ft - INFO - get_dt_for_accurate_interpolation, please wait ...
  stocproc.method_ft - INFO - acc interp N 33 dt 6.25e-01 -> diff 2.88e-02
  stocproc.method_ft - INFO - acc interp N 65 dt 3.12e-01 -> diff 6.07e-03
  stocproc.method_ft - INFO - acc interp N 129 dt 1.56e-01 -> diff 1.41e-03
  stocproc.method_ft - INFO - acc interp N 257 dt 7.81e-02 -> diff 3.46e-04
  stocproc.method_ft - INFO - requires dt < 7.812e-02
  stocproc.method_ft - INFO - get_N_a_b_for_accurate_fourier_integral, please wait ...
  stocproc.method_ft - INFO - J_w_min:1.00e-02 N 32 yields: interval [0.00e+00,4.18e+00] diff 1.30e-02
  stocproc.method_ft - INFO - J_w_min:1.00e-03 N 32 yields: interval [0.00e+00,5.92e+00] diff 3.14e-02
  stocproc.method_ft - INFO - J_w_min:1.00e-02 N 64 yields: interval [0.00e+00,4.18e+00] diff 8.00e-03
  stocproc.method_ft - INFO - J_w_min:1.00e-04 N 32 yields: interval [0.00e+00,7.62e+00] diff 5.84e-02
  stocproc.method_ft - INFO - J_w_min:1.00e-03 N 64 yields: interval [0.00e+00,5.92e+00] diff 7.23e-03
  stocproc.method_ft - INFO - J_w_min:1.00e-02 N 128 yields: interval [0.00e+00,4.18e+00] diff 7.66e-03
  stocproc.method_ft - INFO - increasing N while shrinking the interval does lower the error -> try next level
  stocproc.method_ft - INFO - J_w_min:1.00e-05 N 32 yields: interval [0.00e+00,9.30e+00] diff 8.89e-02
  stocproc.method_ft - INFO - J_w_min:1.00e-04 N 64 yields: interval [0.00e+00,7.62e+00] diff 1.28e-02
  stocproc.method_ft - INFO - J_w_min:1.00e-03 N 128 yields: interval [0.00e+00,5.92e+00] diff 1.70e-03
  stocproc.method_ft - INFO - J_w_min:1.00e-02 N 256 yields: interval [0.00e+00,4.18e+00] diff 7.57e-03
  stocproc.method_ft - INFO - increasing N while shrinking the interval does lower the error -> try next level
  stocproc.method_ft - INFO - J_w_min:1.00e-06 N 32 yields: interval [0.00e+00,1.10e+01] diff 1.30e-01
  stocproc.method_ft - INFO - J_w_min:1.00e-05 N 64 yields: interval [0.00e+00,9.30e+00] diff 1.85e-02
  stocproc.method_ft - INFO - J_w_min:1.00e-04 N 128 yields: interval [0.00e+00,7.62e+00] diff 3.10e-03
  stocproc.method_ft - INFO - J_w_min:1.00e-03 N 256 yields: interval [0.00e+00,5.92e+00] diff 7.99e-04
  stocproc.method_ft - INFO - return, cause tol of 0.001 was reached
  stocproc.method_ft - INFO - requires dx < 2.311e-02
  stocproc.stocproc - INFO - Fourier Integral Boundaries: [0.000e+00, 8.725e+01]
  stocproc.stocproc - INFO - Number of Nodes            : 4096
  stocproc.stocproc - INFO - yields dx                  : 2.130e-02
  stocproc.stocproc - INFO - yields dt                  : 7.201e-02
  stocproc.stocproc - INFO - yields t_max               : 2.949e+02
#+end_example

The sample trajectories are smooth too.
#+begin_src jupyter-python
  %%space plot
  ts = np.linspace(0, params.t_max, 1000)
  EtaTherm.new_process()
  plt.plot(ts, EtaTherm(ts).real)
#+end_src

#+RESULTS:
:RESULTS:
| <matplotlib.lines.Line2D | at | 0x7fc2d1699820> |
[[file:./.ob-jupyter/16fa913e12828f1a33d5ad9e2768a514dc3182bd.png]]
:END:

* Actual Hops
Generate the key for binary caching.
#+begin_src jupyter-python
  hi_key = hierarchyData.HIMetaKey_type(
      HiP=HierachyParam,
      IntP=IntegrationParam,
      SysP=SystemParam,
      Eta=Eta,
      EtaTherm=None,
  )
#+end_src

#+RESULTS:

Initialize Hierarchy.
#+begin_src jupyter-python
  myHierarchy = hierarchyLib.HI(hi_key, number_of_samples=params.N, desc="calculate the heat flow")
#+end_src

#+RESULTS:
: init Hi class, use 112 equation
: /home/hiro/Documents/Projects/UNI/master/masterarb/python/richard_hops/hierarchyLib.py:1058: UserWarning: sum_k_max is not implemented! DO SO BEFORE NEXT USAGE (use simplex).HierarchyParametersType does not yet know about sum_k_max
:   warnings.warn(

Run the integration.
#+begin_src jupyter-python
  myHierarchy.integrate_simple(data_name="energy_flow_therm.data")
#+end_src

#+RESULTS:
#+begin_example
  (10, 1000, 56)
  samples     :0.0%
  integration :0.0%
  samples     :0.9%
  integration :32.9%
  samples     :2.0%
  integration :38.6%
  samples     :3.1%
  integration :89.1%
  samples     :4.4%
  integration :18.5%
  samples     :5.5%
  integration :26.5%
  samples     :6.6%
  integration :81.7%
  samples     :7.9%
  integration :0.0%
  samples     :9.0%
  integration :7.3%
  samples     :10.1%
  integration :31.6%
  samples     :11.2%
  integration :66.5%
  samples     :12.5%
  integration :8.3%
  samples     :13.6%
  integration :50.9%
  samples     :14.9%
  integration :2.8%
  samples     :16.0%
  integration :45.5%
  samples     :17.1%
  integration :79.7%
  samples     :18.4%
  integration :22.0%
  samples     :19.5%
  integration :63.2%
  samples     :20.6%
  integration :92.2%
  samples     :21.9%
  integration :11.1%
  samples     :23.0%
  integration :53.2%
  samples     :24.1%
  integration :0.5%
  samples     :25.2%
  integration :22.7%
  samples     :26.4%
  integration :80.2%
  samples     :27.5%
  integration :70.5%
  samples     :28.7%
  integration :2.6%
  samples     :29.9%
  integration :71.2%
  samples     :31.0%
  integration :21.9%
  samples     :32.1%
  integration :61.9%
  samples     :33.2%
  integration :79.9%
  samples     :34.4%
  integration :65.9%
  samples     :35.4%
  integration :88.1%
  samples     :36.6%
  integration :8.6%
  samples     :37.6%
  integration :96.3%
  samples     :38.9%
  integration :44.0%
  samples     :40.0%
  integration :91.7%
  samples     :41.2%
  integration :18.6%
  samples     :42.4%
  integration :66.3%
  samples     :43.6%
  integration :43.4%
  samples     :44.8%
  integration :96.9%
  samples     :46.0%
  integration :0.0%
  samples     :47.1%
  integration :53.5%
  samples     :48.4%
  integration :13.9%
  samples     :49.4%
  integration :50.8%
  samples     :50.6%
  integration :1.5%
  samples     :51.7%
  integration :55.7%
  samples     :52.9%
  integration :37.7%
  samples     :54.0%
  integration :88.0%
  samples     :55.1%
  integration :90.5%
  samples     :56.4%
  integration :10.5%
  samples     :57.5%
  integration :58.9%
  samples     :58.6%
  integration :64.9%
  samples     :59.8%
  integration :51.9%
  samples     :60.9%
  integration :84.1%
  samples     :62.1%
  integration :19.6%
  samples     :63.2%
  integration :49.9%
  samples     :64.5%
  integration :0.0%
  samples     :65.6%
  integration :26.1%
  samples     :66.8%
  integration :49.4%
  samples     :67.9%
  integration :99.9%
  samples     :69.1%
  integration :30.0%
  samples     :70.2%
  integration :11.1%
  samples     :71.4%
  integration :33.0%
  samples     :72.5%
  integration :86.8%
  samples     :73.8%
  integration :22.5%
  samples     :74.9%
  integration :2.0%
  samples     :75.9%
  integration :90.1%
  samples     :77.0%
  integration :20.2%
  samples     :78.0%
  integration :94.1%
  samples     :79.2%
  integration :7.8%
  samples     :80.4%
  integration :26.5%
  samples     :81.5%
  integration :36.5%
  samples     :82.6%
  integration :82.8%
  samples     :83.9%
  integration :0.2%
  samples     :85.0%
  integration :24.0%
  samples     :86.1%
  integration :77.7%
  samples     :87.4%
  integration :24.7%
  samples     :88.5%
  integration :1.1%
  samples     :89.6%
  integration :29.2%
  samples     :90.8%
  integration :65.4%
  samples     :92.0%
  integration :22.5%
  samples     :93.0%
  integration :90.6%
  samples     :94.2%
  integration :8.2%
  samples     :95.4%
  integration :59.5%
  samples     :96.5%
  integration :99.4%
  samples     :97.6%
  integration :39.5%
  samples     :98.8%
  integration :92.1%
  samples     : 100%
  integration :0.0%
  
#+end_example
Get the samples.
#+BEGIN_SRC jupyter-python
  # to access the data the 'hi_key' is used to find the data in the hdf5 file
  class int_result:
      with hierarchyData.HIMetaData(
          hid_name="energy_flow_therm.data", hid_path="."
      ) as metaData:
          with metaData.get_HIData(hi_key, read_only=True) as data:
              smp = data.get_samples()
              print("{} samples found in database".format(smp))
              τ = data.get_time()
              rho_τ = data.get_rho_t()
              s_proc = np.array(data.stoc_proc)
              states = np.array(data.aux_states).copy()
              ψ_1 = np.array(data.aux_states)[:, :, 0 : params.num_exp_t * params.dim]
              ψ_0 = np.array(data.stoc_traj)
              y = np.array(data.y)
              η = np.array(data.stoc_proc)
#+end_src

#+RESULTS:
: 800 samples found in database

Calculate energy.
#+begin_src jupyter-python
  %matplotlib inline
  energy = np.array([np.trace(ρ @ params.H_s).real for ρ in int_result.rho_τ])
  plt.plot(int_result.τ, energy)
#+end_src

#+RESULTS:
:RESULTS:
| <matplotlib.lines.Line2D | at | 0x7fc2d14a16d0> |
[[file:./.ob-jupyter/f5da2ef7d690140aef6072b282eb0aa0f398d34a.png]]
:END:

* Energy Flow
:PROPERTIES:
:ID:       eefb1594-e399-4d24-9dd7-a57addd42e65
:END:
#+begin_src jupyter-python
  int_result.ψ_1.shape
#+end_src

#+RESULTS:
| 1280 | 1000 | 10 |

Let's look at the norm.
#+begin_src jupyter-python
  plt.plot(int_result.τ, (int_result.ψ_0[0].conj() * int_result.ψ_0[0]).sum(axis=1).real)
#+end_src

#+RESULTS:
:RESULTS:
| <matplotlib.lines.Line2D | at | 0x7fc2d1410a60> |
[[file:./.ob-jupyter/1cf2d718580a0474d3e7b9521550e6f1ccdcf21d.png]]
:END:

And try to calculate the energy flow.
#+begin_src jupyter-python
  def flow_for_traj(ψ_0, ψ_1):
      a = np.array((params.L @ ψ_0.T).T)

      ψ_1 = (-w * g * params.bcf_scale)[None, :, None] * ψ_1.reshape(
          params.t_steps, params.num_exp_t, params.dim
      )

      # return np.array(np.sum(ψ_0.conj() * ψ_0, axis=1)).flatten().real
      return np.array(
          2
          ,* (
              1j
              ,* np.sum(a.conj()[:, None, :] * ψ_1, axis=(1, 2))
              / np.sum(ψ_0.conj() * ψ_0, axis=1)
          ).real
      ).flatten()


  def flow_for_traj_alt(ψ_0, y):
      Eta.new_process(y)
      eta_dot = scipy.misc.derivative(Eta, τ, dx=1e-8)
      a = np.array((L @ ψ_0.T).T)

      return -(
          2j * eta_dot.conj() * np.array((np.sum(ψ_0.conj() * a, axis=1))).flatten()
      ).real
#+end_src

#+RESULTS:

Now we calculate the average over all trajectories.
#+begin_src jupyter-python
  class Flow:
      j = np.zeros_like(int_result.τ)
      for i in range(0, params.N):
          j += flow_for_traj(int_result.ψ_0[i], int_result.ψ_1[i])
      j /= params.N
#+end_src

#+RESULTS:

And do the same with the alternative implementation.
#+begin_src jupyter-python
  ja = np.zeros_like(params.τ)
  for i in range(0, N):
      ja += flow_for_traj_alt(ψ_0[i], y[i])
  ja /= N
#+end_src

#+RESULTS:
:RESULTS:
# [goto error]
:
: AttributeErrorTraceback (most recent call last)
: <ipython-input-25-1eb991d00182> in <module>
: ----> 1 ja = np.zeros_like(params.τ)
:       2 for i in range(0, N):
:       3     ja += flow_for_traj_alt(ψ_0[i], y[i])
:       4 ja /= N
:
: AttributeError: type object 'params' has no attribute 'τ'
:END:

And plot it :)
#+begin_src jupyter-python
  %matplotlib inline
  plt.plot(int_result.τ, Flow.j)
  #plt.plot(τ, ja)
  plt.show()
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/d36fa5aee1635793a8667941513e2c0454f3be8f.png]]

Let's calculate the integrated energy.
#+begin_src jupyter-python
  E_t = np.array(
      [0]
      + [
          scipy.integrate.simpson(Flow.j[0:n], int_result.τ[0:n])
          for n in range(1, len(int_result.τ))
      ]
  )
  print(E_t[-1])
  E_t = E_t / E_t[-1] * 2
#+end_src

#+RESULTS:
: 1.9718688725570916

With this we can retrieve the energy of the interaction Hamiltonian.
#+begin_src jupyter-python
  E_I = 2 - energy - E_t
#+end_src

#+RESULTS:

#+begin_src jupyter-python
  %%space plot
  plt.rcParams['figure.figsize'] = [15, 10]
  #plt.plot(τ, j, label="$J$", linestyle='--')
  plt.plot(int_result.τ, E_t, label=r"$\langle H_{\mathrm{B}}\rangle$")
  plt.plot(int_result.τ, E_I, label=r"$\langle H_{\mathrm{I}}\rangle$")
  plt.plot(int_result.τ, energy, label=r"$\langle H_{\mathrm{S}}\rangle$")

  plt.xlabel("τ")
  plt.legend()
  plt.show()
#+end_src

#+RESULTS:
:RESULTS:
| <matplotlib.lines.Line2D | at | 0x7fc2d128f2b0> |
| <matplotlib.lines.Line2D | at | 0x7fc2d128f6a0> |
| <matplotlib.lines.Line2D | at | 0x7fc2d128fa30> |
: Text(0.5, 0, 'τ')
: <matplotlib.legend.Legend at 0x7fc2d128f820>
[[file:./.ob-jupyter/3d9e49cd424f5b7dc3de2a397f7d596f525aff74.png]]
:END:

* System + Interaction Energy
#+begin_src jupyter-python
  def h_si_for_traj(ψ_0, ψ_1):
      a = np.array((params.L @ ψ_0.T).T)
      b = np.array((params.H_s @ ψ_0.T).T)
      ψ_1 = (g*params.bcf_scale)[None, :, None] * ψ_1.reshape(
          params.t_steps, params.num_exp_t, params.dim
      )
      E_i = np.array(
          2
          ,* (
              -1j
              ,* np.sum(
                  a.conj()[:, None, :]
                  ,* ψ_1,
                  axis=(1, 2),
              )
          ).real
      ).flatten()
      E_s = np.array(np.sum(b.conj() * ψ_0, axis=1)).flatten().real

      return (E_i + E_s) / np.sum(ψ_0.conj() * ψ_0, axis=1).real


  def h_si_for_traj_alt(ψ_0, y):
      Eta.new_process(y)

      a = np.array((L.conj().T @ ψ_0.T).T)
      b = np.array((H_s @ ψ_0.T).T)
      E_i = np.array(2 * (Eta(τ) * 1j * np.sum(a.conj() * ψ_0, axis=1)).real).flatten()
      E_s = np.array(np.sum(b.conj() * ψ_0, axis=1)).flatten().real

      return E_i + E_s
#+end_src

#+RESULTS:

#+begin_src jupyter-python
  e_si = np.zeros_like(int_result.τ)
  for i in range(0, params.N):
      e_si += h_si_for_traj(int_result.ψ_0[i], int_result.ψ_1[i])
  e_si /= params.N
#+end_src

#+RESULTS:

Not too bad...
#+begin_src jupyter-python
  plt.plot(int_result.τ, e_si, label="real")
  plt.plot(int_result.τ, E_I + energy)
  plt.legend()
#+end_src


#+RESULTS:
:RESULTS:
: <matplotlib.legend.Legend at 0x7fc2d113ec70>
[[file:./.ob-jupyter/5563b4ea985bd0cfeaee0ad67412152851290286.png]]
:END:

* Rivas Internal energy

*