diff --git a/src/config.json b/config.json
similarity index 69%
rename from src/config.json
rename to config.json
index 82f4daf..925d7d5 100644
--- a/src/config.json
+++ b/config.json
@@ -1,6 +1,6 @@
 {
   "simulation": {
-    "num_walker": 20000,
+    "num_walker": 30000,
     "seed": null
   },
   "molecule": {
@@ -9,15 +9,10 @@
   },
   "correlation_times": {
     "distribution": "DeltaDistribution",
-    "tau": {
-      "start": 1e-3,
-      "stop": 1e0,
-      "steps": 1,
-      "is_log": true
-    }
+    "tau": 1e-3
   },
   "motion": {
-    "model": "RandomJump"
+    "model": "TetrahedralJump"
   },
   "spectrum": {
     "dwell_time": 1e-6,
@@ -37,16 +32,16 @@
   },
   "stimulated_echo": {
     "t_evo": {
-      "start": 10e-6,
+      "start": 1e-6,
       "stop": 100e-6,
-      "steps": 98
+      "steps": 99
     },
     "t_mix": {
-      "start": 1e-7,
-      "stop": 1e-1,
-      "steps": 31,
+      "start": 1e-6,
+      "stop": 1e0,
+      "steps": 19,
       "is_log": true
     },
-    "t_echo": 15e-6
+    "t_echo": 0e-6
   }
 }
\ No newline at end of file
diff --git a/main.py b/main.py
new file mode 100644
index 0000000..8874982
--- /dev/null
+++ b/main.py
@@ -0,0 +1,22 @@
+from numpy.random import default_rng
+import matplotlib.pyplot as plt
+
+from rwsims.sims import run_ste_sim, run_spectrum_sim
+from rwsims.motions import TetrahedralJump
+
+
+# run_ste_sim('config.json')
+# run_spectrum_sim('config.json')
+
+rng = default_rng()
+
+tetra = TetrahedralJump(1, 0, rng)
+
+for _ in range(100):
+    tetra.start()
+    omegas = tetra.jump(100)
+
+    plt.plot(omegas, '.')
+    break
+
+plt.show()
diff --git a/src/rwsims/__init__.py b/rwsims/__init__.py
similarity index 100%
rename from src/rwsims/__init__.py
rename to rwsims/__init__.py
diff --git a/src/rwsims/distributions.py b/rwsims/distributions.py
similarity index 93%
rename from src/rwsims/distributions.py
rename to rwsims/distributions.py
index 8752d78..87a83af 100644
--- a/src/rwsims/distributions.py
+++ b/rwsims/distributions.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 
 import numpy as np
-from numpy.typing import ArrayLike
+# from numpy.typing import ArrayLike
 from numpy.random import Generator
 
 
diff --git a/src/rwsims/functions.py b/rwsims/functions.py
similarity index 66%
rename from src/rwsims/functions.py
rename to rwsims/functions.py
index 0d33bbf..64d1982 100644
--- a/src/rwsims/functions.py
+++ b/rwsims/functions.py
@@ -1,17 +1,17 @@
 from __future__ import annotations
 
 import numpy as np
-from numpy.typing import ArrayLike
+# from numpy.typing import ArrayLike
 
-from distributions import BaseDistribution
-from motions import BaseMotion
+from .distributions import BaseDistribution
+from .motions import BaseMotion
 
 
 def ste(x, a, f_infty, tau, beta):
     return a*((1-f_infty) * np.exp(-(x/tau)**beta) + f_infty)
 
 
-def pulse_attn(freq: ArrayLike, t_pulse: float) -> ArrayLike:
+def pulse_attn(freq, t_pulse: float):
     # cf. Schmitt-Rohr/Spieß eq. 2.126; omega_1 * t_p = pi/2
     pi_half_squared = np.pi**2 / 4
     omega = 2 * np.pi * freq
@@ -19,4 +19,4 @@ def pulse_attn(freq: ArrayLike, t_pulse: float) -> ArrayLike:
     numerator = np.sin(np.sqrt(pi_half_squared + omega**2 * t_pulse**2 / 2))
     denominator = np.sqrt(pi_half_squared + omega**2 * t_pulse**2 / 4)
 
-    return np.pi * numerator/denominator / 2
\ No newline at end of file
+    return np.pi * numerator/denominator / 2
diff --git a/src/rwsims/motions.py b/rwsims/motions.py
similarity index 86%
rename from src/rwsims/motions.py
rename to rwsims/motions.py
index de2644d..c18633a 100644
--- a/src/rwsims/motions.py
+++ b/rwsims/motions.py
@@ -4,7 +4,6 @@ from abc import ABC, abstractmethod
 
 import numpy as np
 from numpy.random import Generator
-from numpy.typing import ArrayLike
 
 
 class BaseMotion(ABC):
@@ -26,7 +25,7 @@ class BaseMotion(ABC):
         pass
 
     @abstractmethod
-    def jump(self, size: int = 1) -> ArrayLike:
+    def jump(self, size: int = 1) -> 'ArrayLike':
         pass
 
 
@@ -35,7 +34,7 @@ class RandomJump(BaseMotion):
     def __repr__(self):
         return 'Random Jump'
 
-    def jump(self, size: int = 1) -> ArrayLike:
+    def jump(self, size: int = 1) -> 'ArrayLike':
         return omega_q(self._delta, self._eta, *draw_orientation(self._rng, size=size))
 
 
@@ -74,15 +73,28 @@ class TetrahedralJump(BaseMotion):
             corners[0],
             np.array(spherical_to_xyz(1., np.arccos(cos_theta0), phi0)),
         )
+
         orientations = np.zeros(4)
         for i in range(4):
-            corner_lab = np.dot(rot, corners[i])
+            corner_lab = rot @ corners[i]
             _, theta_i, phi_i = xyz_to_spherical(*corner_lab)
             orientations[i] = omega_q(self._delta, self._eta, theta_i, phi_i)
 
+        # print(orientations)
+        #
+        # theta0 = np.arccos(cos_theta0)
+        # v0 = np.array([np.sin(theta0) * np.cos(phi0), np.sin(theta0)*np.sin(theta0), cos_theta0])
+        # norm = np.linalg.norm(v0)
+        # print(norm)
+        #
+        #
+        # corners = np.zeros((4, 3))
+        # corners[0] = v0
+
+
         return orientations
 
-    def jump(self, size: int = 1) -> ArrayLike:
+    def jump(self, size: int = 1) -> 'ArrayLike':
         jumps = self._rng.choice([1, 2, 3], size=size)
         jumps = np.cumsum(jumps) + self._start
         jumps %= 4
@@ -94,10 +106,10 @@ class TetrahedralJump(BaseMotion):
 
 # Helper functions
 
-def xyz_to_spherical(x_in: float, y_in: float, z_in: float) -> tuple[float, float, float]:
+def xyz_to_spherical(x_in: float, y_in: float, z_in: float) -> tuple[np.floating, float, float]:
     r = np.linalg.norm([x_in, y_in, z_in])
-    theta = np.arccos(z_in)
-    phi = np.arctan2(y_in, x_in)
+    theta: float = np.arccos(z_in)
+    phi: float = np.arctan2(y_in, x_in)
 
     return r, theta, phi
 
diff --git a/src/rwsims/parameter.py b/rwsims/parameter.py
similarity index 91%
rename from src/rwsims/parameter.py
rename to rwsims/parameter.py
index 762f6de..d4888fe 100644
--- a/src/rwsims/parameter.py
+++ b/rwsims/parameter.py
@@ -6,9 +6,9 @@ from math import prod
 from typing import Any
 
 import numpy as np
-from numpy._typing import ArrayLike
+# from numpy.typing import ArrayLike
 
-from functions import pulse_attn
+from .functions import pulse_attn
 from .distributions import BaseDistribution
 from .motions import BaseMotion
 
@@ -41,8 +41,8 @@ class MoleculeParameter:
 
 @dataclass
 class StimEchoParameter:
-    t_evo: ArrayLike
-    t_mix: ArrayLike
+    t_evo: 'ArrayLike'
+    t_mix: 'ArrayLike'
     t_echo: float
     t_max: float = field(init=False)
 
@@ -54,14 +54,14 @@ class StimEchoParameter:
 class SpectrumParameter:
     dwell_time: float
     num_points: int
-    t_echo: ArrayLike
+    t_echo: 'ArrayLike'
     lb: float
     t_pulse: float
-    t_acq: ArrayLike = field(init=False)
-    freq: ArrayLike = field(init=False)
+    t_acq: 'ArrayLike' = field(init=False)
+    freq: 'ArrayLike' = field(init=False)
     t_max: float = field(init=False)
-    dampening: ArrayLike = field(init=False)
-    pulse_attn: ArrayLike = field(init=False)
+    dampening: 'ArrayLike' = field(init=False)
+    pulse_attn: 'ArrayLike' = field(init=False)
 
     def __post_init__(self):
         self.t_acq = np.arange(self.num_points) * self.dwell_time
diff --git a/src/rwsims/parsing.py b/rwsims/parsing.py
similarity index 100%
rename from src/rwsims/parsing.py
rename to rwsims/parsing.py
diff --git a/src/rwsims/sims.py b/rwsims/sims.py
similarity index 96%
rename from src/rwsims/sims.py
rename to rwsims/sims.py
index 75083db..5e4e3a4 100644
--- a/src/rwsims/sims.py
+++ b/rwsims/sims.py
@@ -54,6 +54,7 @@ def run_ste_sim(config_file: str):
                     dephased = phase(t_evo_k)
                     t0 = t_mix + t_evo_k
                     rephased = phase(t0 + t_evo_k + 2*t_echo) + phase(t0) - 2 * phase(t0+t_echo)
+                    # print(t_evo_k, t0 + t_evo_k + 2*t_echo, t0)
                     cc[:, k] += np.cos(dephased)*np.cos(rephased)
                     ss[:, k] += np.sin(dephased)*np.sin(rephased)
 
@@ -174,17 +175,20 @@ def make_trajectory(
 
     # number of jumps that are simulated at once
     chunks = min(int(0.51 * t_max / dist.tau_jump), 100_000) + 1
+    # print(chunks)
 
     t = [np.array([t_passed])]
     phase = [np.array([init_phase])]
+    # omega = [np.array([0])]
     while t_passed < t_max:
         # frequencies between jumps
         current_omega = motion.jump(size=chunks)
         # times at a particular position
         t_wait = dist.wait(size=chunks)
 
-        accumulated_phase = np.cumsum(t_wait * current_omega) + phase[-1]
+        accumulated_phase = np.cumsum(t_wait * current_omega) + phase[-1][-1]
         phase.append(accumulated_phase)
+        # omega.append(current_omega)
 
         t_wait = np.cumsum(t_wait) + t_passed
         t_passed = t_wait[-1]
@@ -192,6 +196,12 @@ def make_trajectory(
 
     t = np.concatenate(t)
     phase = np.concatenate(phase)
+    # omega = np.concatenate(omega)
+
+    # fig_test, ax_test = plt.subplots()
+    # ax_test.plot(t, phase, 'x-')
+
+    # np.savetxt('trajectory.dat', np.c_[t, phase, omega])
 
     # convenient interpolation to get phase at arbitrary times
     phase_interpol = interp1d(t, phase)
diff --git a/src/rwsims/spectrum.py b/rwsims/spectrum.py
similarity index 100%
rename from src/rwsims/spectrum.py
rename to rwsims/spectrum.py
diff --git a/src/rwsims/ste.py b/rwsims/ste.py
similarity index 100%
rename from src/rwsims/ste.py
rename to rwsims/ste.py
diff --git a/src/rwsims/tetrahedral_spectrum.py b/rwsims/tetrahedral_spectrum.py
similarity index 100%
rename from src/rwsims/tetrahedral_spectrum.py
rename to rwsims/tetrahedral_spectrum.py
diff --git a/src/rwsims/tetrahedral_spectrum_chunk.py b/rwsims/tetrahedral_spectrum_chunk.py
similarity index 100%
rename from src/rwsims/tetrahedral_spectrum_chunk.py
rename to rwsims/tetrahedral_spectrum_chunk.py
diff --git a/src/main.py b/src/main.py
deleted file mode 100644
index d75e40e..0000000
--- a/src/main.py
+++ /dev/null
@@ -1,5 +0,0 @@
-from rwsims.sims import run_ste_sim, run_spectrum_sim
-
-# run_ste_sim('config.json')
-
-run_spectrum_sim('config.json')