diff --git a/src/grogu_magn/core.py b/src/grogu_magn/core.py
index c294e06..8e3ab4e 100644
--- a/src/grogu_magn/core.py
+++ b/src/grogu_magn/core.py
@@ -142,8 +142,8 @@ def setup_pairs_and_magnetic_entities(
     # for every site we have to store 3 Greens function (and the associated _tmp-s) in the 3 reference directions
     for mag_ent in magnetic_entities:
         parsed = parse_magnetic_entity(dh, **mag_ent)  # parse orbital indexes
-        mag_ent["orbital_indeces"] = parsed
-        mag_ent["spin_box_indeces"] = blow_up_orbindx(
+        mag_ent["orbital_indices"] = parsed
+        mag_ent["spin_box_indices"] = blow_up_orbindx(
             parsed
         )  # calculate spin box indexes
         # if orbital is not set use all
@@ -165,7 +165,7 @@ def setup_pairs_and_magnetic_entities(
                 mag_ent["xyz"].append(dh.xyz[atom_idx])
 
         # calculate size for Greens function generation
-        spin_box_shape = len(mag_ent["spin_box_indeces"])
+        spin_box_shape = len(mag_ent["spin_box_indices"])
 
         mag_ent["energies"] = (
             []
@@ -205,8 +205,8 @@ def setup_pairs_and_magnetic_entities(
         pair["dist"] = np.linalg.norm(xyz_ai - xyz_aj)
 
         # calculate size for Greens function generation
-        spin_box_shape_i = len(magnetic_entities[pair["ai"]]["spin_box_indeces"])
-        spin_box_shape_j = len(magnetic_entities[pair["aj"]]["spin_box_indeces"])
+        spin_box_shape_i = len(magnetic_entities[pair["ai"]]["spin_box_indices"])
+        spin_box_shape_j = len(magnetic_entities[pair["aj"]]["spin_box_indices"])
         pair["tags"] = []
         for mag_ent in [magnetic_entities[pair["ai"]], magnetic_entities[pair["aj"]]]:
             tag = ""
@@ -261,3 +261,31 @@ def setup_pairs_and_magnetic_entities(
             )
 
     return pairs, magnetic_entities
+
+
+def onsite_projection(matrix, idx1, idx2):
+    """_summary_
+
+    Args:
+        matrix (_type_): _description_
+        idx (_type_): _description_
+
+    Returns:
+        _type_: _description_
+    """
+
+    return matrix[..., idx1, :][..., idx2]
+
+
+def onsite_projection(matrix, idx):
+    """_summary_
+
+    Args:
+        matrix (_type_): _description_
+        idx (_type_): _description_
+
+    Returns:
+        _type_: _description_
+    """
+
+    return matrix[..., idx, :][..., idx]
diff --git a/src/grogu_magn/io.py b/src/grogu_magn/io.py
index f9a5b47..f7c2f19 100644
--- a/src/grogu_magn/io.py
+++ b/src/grogu_magn/io.py
@@ -17,6 +17,9 @@ default_args = dict(
     ebot=None,
     eset=42,
     esetp=1000,
+    calculate_charge=True,
+    charges=[],
+    parallel_solver_for_Gk=True,
 )
 
 # parser = ArgumentParser()
@@ -101,6 +104,9 @@ def print_parameters(simulation_parameters):
     print(
         "================================================================================================================================================================"
     )
+    if simulation_parameters["calculate_charge"]:
+        print("The calculated charge of the Hamiltonian in the quantization axes: ")
+        print(simulation_parameters["charges"])
 
 
 def print_atoms_and_pairs(magnetic_entities, pairs):
diff --git a/test.ipynb b/test.ipynb
index 65186e2..62819e8 100644
--- a/test.ipynb
+++ b/test.ipynb
@@ -2,25 +2,16 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [
     {
-     "name": "stderr",
+     "name": "stdout",
      "output_type": "stream",
      "text": [
-      "[Daniels-Air:78400] shmem: mmap: an error occurred while determining whether or not /var/folders/yh/dx7xl94n3g52ts3td8qcxjcc0000gn/T//ompi.Daniels-Air.501/jf.0/412614656/sm_segment.Daniels-Air.501.18980000.0 could be created.\n"
+      "0.14.3\n",
+      "1.24.4\n"
      ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "'0.14.3'"
-      ]
-     },
-     "execution_count": 1,
-     "metadata": {},
-     "output_type": "execute_result"
     }
    ],
    "source": [
@@ -54,7 +45,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 22,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -63,7 +54,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 23,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -116,7 +107,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 24,
    "metadata": {},
    "outputs": [
     {
@@ -149,7 +140,7 @@
       "Eset:  300\n",
       "Esetp:  1000\n",
       "================================================================================================================================================================\n",
-      "Setup done. Elapsed time: 2.303232 s\n",
+      "Setup done. Elapsed time: 2435.031293583 s\n",
       "================================================================================================================================================================\n"
      ]
     }
@@ -207,14 +198,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 25,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Hamiltonian and exchange field rotated. Elapsed time: 2.914966208 s\n",
+      "Hamiltonian and exchange field rotated. Elapsed time: 2435.900105791 s\n",
       "================================================================================================================================================================\n"
      ]
     }
@@ -268,14 +259,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 26,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Site and pair dictionaries created. Elapsed time: 2.944875791 s\n",
+      "Site and pair dictionaries created. Elapsed time: 2438.737295 s\n",
       "================================================================================================================================================================\n"
      ]
     }
@@ -297,7 +288,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 27,
    "metadata": {},
    "outputs": [
     {
@@ -311,7 +302,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "k set created. Elapsed time: 2.974842875 s\n",
+      "k set created. Elapsed time: 2441.389173 s\n",
       "================================================================================================================================================================\n"
      ]
     }
@@ -336,14 +327,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 29,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Rotations done perpendicular to quantization axis. Elapsed time: 3.253283416 s\n",
+      "Rotations done perpendicular to quantization axis. Elapsed time: 2460.81759 s\n",
       "================================================================================================================================================================\n"
      ]
     }
diff --git a/test.py b/test.py
index 4e5aa93..a83cf79 100644
--- a/test.py
+++ b/test.py
@@ -178,9 +178,14 @@ if rank == root_node:
 kset = make_kset(
     dirs=simulation_parameters["kdirs"], NUMK=simulation_parameters["kset"]
 )
-wkset = np.ones(len(kset)) / len(kset)  # generate weights for k points
-kpcs = np.array_split(kset, size)  # split the k points based on MPI size
 
+# generate weights for k points
+wkset = np.ones(len(kset)) / len(kset)
+
+# split the k points based on MPI size
+kpcs = np.array_split(kset, size)
+
+# use progress bar if available
 if tqdm_imported:
     kpcs[root_node] = tqdm(kpcs[root_node], desc="k loop")
 
@@ -190,3 +195,237 @@ if rank == root_node:
     print(
         "================================================================================================================================================================"
     )
+
+# this will contain the three Hamiltonian in the
+# reference directions needed to calculate the energy
+# variations upon rotation
+hamiltonians = []
+
+# iterate over the reference directions (quantization axes)
+for i, orient in enumerate(simulation_parameters["ref_xcf_orientations"]):
+    # obtain rotated exchange field and Hamiltonian
+    R = RotMa2b(simulation_parameters["scf_xcf_orientation"], orient["o"])
+    rot_XCF = np.einsum("ij,jklm->iklm", R, XCF)
+    rot_H_XCF = sum(
+        [np.kron(rot_XCF[i], tau) for i, tau in enumerate([tau_x, tau_y, tau_z])]
+    )
+    rot_H_XCF_uc = rot_H_XCF[uc_in_sc_idx]
+
+    # obtain total Hamiltonian with the rotated exchange field
+    rot_H = hTRS + rot_H_XCF  # equation 76
+
+    # store the relevant information of the Hamiltonian
+    hamiltonians.append(dict(orient=orient["o"], H=rot_H))
+
+    if simulation_parameters["calculate_charge"]:
+        hamiltonians[-1]["GS"] = np.zeros(
+            (simulation_parameters["eset"], rot_H.shape[1], rot_H.shape[2]),
+            dtype="complex128",
+        )
+        hamiltonians[-1]["GS_tmp"] = np.zeros(
+            (simulation_parameters["eset"], rot_H.shape[1], rot_H.shape[2]),
+            dtype="complex128",
+        )
+
+    # these are the rotations (for now) perpendicular to the quantization axis
+    for u in orient["vw"]:
+        # section 2.H
+        Tu = np.kron(np.eye(NO, dtype=int), tau_u(u))
+        Vu1, Vu2 = calc_Vu(rot_H_XCF_uc, Tu)
+
+        for mag_ent in magnetic_entities:
+            idx = mag_ent["spin_box_indices"]
+            # fill up the perturbed potentials (for now) based on the on-site projections
+            mag_ent["Vu1"][i].append(onsite_projection(Vu1, idx))
+            mag_ent["Vu2"][i].append(onsite_projection(Vu2, idx))
+
+if rank == root_node:
+    times["reference_rotations_time"] = timer()
+    print(
+        f"Rotations done perpendicular to quantization axis. Elapsed time: {times['reference_rotations_time']} s"
+    )
+    print(
+        "================================================================================================================================================================"
+    )
+
+# provide helpful information to estimate the runtime and memory
+# requirements of the Greens function calculations
+if rank == root_node:
+    print("Starting matrix inversions.")
+    print(f"Total number of k points: {kset.shape[0]}")
+    print(f"Number of energy samples per k point: {simulation_parameters['eset']}")
+    print(f"Total number of directions: {len(hamiltonians)}")
+    print(
+        f"Total number of matrix inversions: {kset.shape[0] * len(hamiltonians) * simulation_parameters['eset']}"
+    )
+    print(f"The shape of the Hamiltonian and the Greens function is {NO}x{NO}={NO*NO}")
+    # https://stackoverflow.com/questions/70746660/how-to-predict-memory-requirement-for-np-linalg-inv
+    # memory is O(64 n**2) for complex matrices
+    memory_size = getsizeof(hamiltonians[0]["H"].base) / 1024
+    print(
+        f"Memory taken by a single Hamiltonian is: {getsizeof(hamiltonians[0]['H'].base) / 1024} KB"
+    )
+    print(f"Expected memory usage per matrix inversion: {memory_size * 32} KB")
+    print(
+        f"Expected memory usage per k point for parallel inversion: {memory_size * len(hamiltonians) * simulation_parameters['eset'] * 32} KB"
+    )
+    print(
+        f"Expected memory usage on root node: {len(np.array_split(kset, size)[0]) * memory_size * len(hamiltonians) * simulation_parameters['eset'] * 32 / 1024} MB"
+    )
+    print(
+        "================================================================================================================================================================"
+    )
+
+# MPI barrier
+comm.Barrier()
+
+# make energy contour
+cont = make_contour(
+    emin=simulation_parameters["ebot"],
+    enum=simulation_parameters["eset"],
+    p=simulation_parameters["esetp"],
+)
+eran = cont.ze
+
+# sampling the integrand on the contour and the BZ
+for k in kpcs[rank]:
+    # weight of k point in BZ integral
+    wk = wkset[rank]
+
+    # iterate over reference directions
+    for i, hamiltonian_orientation in enumerate(hamiltonians):
+        # calculate Hamiltonian and Overlap matrix in a given k point
+        H = hamiltonian_orientation["H"]
+        HK, SK = hsk(H, ss, dh.sc_off, k)
+
+        if simulation_parameters["parallel_solver_for_Gk"]:
+            Gk = parallel_Gk(HK, SK, eran, simulation_parameters["eset"])
+        else:
+            # solve Greens function sequentially for the energies, because of memory bound
+            Gk = sequential_GK(HK, SK, eran, simulation_parameters["eset"])
+
+        # saving this for total charge
+        if simulation_parameters["calculate_charge"]:
+            hamiltonian_orientation["GS_tmp"] += Gk @ SK * wk
+
+        # store the Greens function slice of the magnetic entities
+        for mag_ent in magnetic_entities:
+            idx = mag_ent["spin_box_indices"]
+            mag_ent["Gii_tmp"][i] += onsite_projection(Gk, idx) * wk
+
+        for pair in pairs:
+            # add phase shift based on the cell difference
+            phase = np.exp(1j * 2 * np.pi * k @ pair["Ruc"].T)
+
+            # get the pair orbital sizes from the magnetic entities
+            ai = magnetic_entities[pair["ai"]]["spin_box_indices"]
+            aj = magnetic_entities[pair["aj"]]["spin_box_indices"]
+
+            # store the Greens function slice of the magnetic entities
+            pair["Gij_tmp"][i] += onsite_projection(Gk, ai, aj) * phase * wk
+            pair["Gji_tmp"][i] += onsite_projection(Gk, aj, ai) / phase * wk
+
+# summ reduce partial results of mpi nodes
+for i in range(len(hamiltonians)):
+    # for total charge calculation
+    if simulation_parameters["calculate_charge"]:
+        comm.Reduce(hamiltonians[i]["GS_tmp"], hamiltonians[i]["GS"], root=root_node)
+
+    for mag_ent in magnetic_entities:
+        comm.Reduce(mag_ent["Gii_tmp"][i], mag_ent["Gii"][i], root=root_node)
+
+    for pair in pairs:
+        comm.Reduce(pair["Gij_tmp"][i], pair["Gij"][i], root=root_node)
+        comm.Reduce(pair["Gji_tmp"][i], pair["Gji"][i], root=root_node)
+
+if rank == root_node:
+    times["green_function_inversion_time"] = timer()
+    print(
+        f"Calculated Greens functions. Elapsed time: {times['green_function_inversion_time']} s"
+    )
+    print(
+        "================================================================================================================================================================"
+    )
+
+
+if rank == root_node:
+    # Calculate total charge
+    if simulation_parameters["calculate_charge"]:
+        for hamiltonian in hamiltonians:
+            GS = hamiltonian["GS"]
+            traced = np.trace((GS), axis1=1, axis2=2)
+            simulation_parameters["charges"].append(int_de_ke(traced, cont.we))
+
+    # iterate over the magnetic entities
+    for tracker, mag_ent in enumerate(magnetic_entities):
+        # iterate over the quantization axes
+        for i, Gii in enumerate(mag_ent["Gii"]):
+            storage = []
+            # iterate over the first and second order local perturbations
+            for Vu1, Vu2 in zip(mag_ent["Vu1"][i], mag_ent["Vu2"][i]):
+                # The Szunyogh-Lichtenstein formula
+                traced = np.trace(
+                    (Vu2 @ Gii + 0.5 * Gii @ Vu1 @ Gii), axis1=1, axis2=2
+                )  # this is the on site projection
+                # evaluation of the contour integral
+                storage.append(int_de_ke(traced, cont.we))
+            # fill up the magnetic entities dictionary with the energies
+            magnetic_entities[tracker]["energies"].append(storage)
+        # convert to np array
+        magnetic_entities[tracker]["energies"] = np.array(
+            magnetic_entities[tracker]["energies"]
+        )
+
+    # iterate over the pairs
+    for tracker, pair in enumerate(pairs):
+        # iterate over the quantization axes
+        for i, (Gij, Gji) in enumerate(zip(pair["Gij"], pair["Gji"])):
+            site_i = magnetic_entities[pair["ai"]]
+            site_j = magnetic_entities[pair["aj"]]
+            storage = []
+            # iterate over the first order local perturbations in all possible orientations for the two sites
+            for Vui in site_i["Vu1"][i]:
+                for Vuj in site_j["Vu1"][i]:
+                    # The Szunyogh-Lichtenstein formula
+                    traced = np.trace(
+                        (Vui @ Gij @ Vuj @ Gji), axis1=1, axis2=2
+                    )  # this is the on site projection
+                    # evaluation of the contour integral
+                    storage.append(int_de_ke(traced, cont.we))
+            # fill up the pairs dictionary with the energies
+            pairs[tracker]["energies"].append(storage)
+        # convert to np array
+        pairs[tracker]["energies"] = np.array(pairs[tracker]["energies"])
+
+    # calculate magnetic parameters
+    for mag_ent in magnetic_entities:
+        Kxx, Kyy, Kzz, consistency = calculate_anisotropy_tensor(mag_ent)
+        mag_ent["K"] = np.array([Kxx, Kyy, Kzz]) * sisl.unit_convert("eV", "meV")
+        mag_ent["K_consistency"] = consistency
+
+    for pair in pairs:
+        J_iso, J_S, D, J = calculate_exchange_tensor(pair)
+        pair["J_iso"] = J_iso * sisl.unit_convert("eV", "meV")
+        pair["J_S"] = J_S * sisl.unit_convert("eV", "meV")
+        pair["D"] = D * sisl.unit_convert("eV", "meV")
+        pair["J"] = J * sisl.unit_convert("eV", "meV")
+
+    times["end_time"] = timer()
+    print(
+        "##################################################################### GROGU OUTPUT #############################################################################"
+    )
+
+    print_parameters(simulation_parameters)
+    print_atoms_and_pairs(magnetic_entities, pairs)
+    print_runtime_information(times)
+
+    pairs, magnetic_entities = remove_clutter_for_save(pairs, magnetic_entities)
+    # create output dictionary with all the relevant data
+    results = dict(
+        parameters=simulation_parameters,
+        magnetic_entities=magnetic_entities,
+        pairs=pairs,
+        runtime=times,
+    )
+
+    save_pickle(simulation_parameters["outfile"], results)