Added new function for calculating the occupation matrix

src/mdevaluate/free_energy_landscape.py

@@ -9,10 +9,61 @@ import cmath
 import pandas as pd
 import multiprocessing as mp
 
-
 VALID_GEOMETRY = {"cylindrical", "slab"}
 
 
+def occupation_matrix(trajectory, edge_length=0.05, segments=1000, skip=0.1, nodes=8):
+    frame_indices = np.unique(
+        np.int_(np.linspace(len(trajectory) * skip, len(trajectory) - 1, num=segments))
+    )
+
+    box = trajectory[0].box
+    x_bins = np.arange(0, box[0][0] + edge_length, edge_length)
+    y_bins = np.arange(0, box[1][1] + edge_length, edge_length)
+    z_bins = np.arange(0, box[2][2] + edge_length, edge_length)
+    bins = [x_bins, y_bins, z_bins]
+    # Trajectory is split for parallel computing
+    size = math.ceil(len(frame_indices) / nodes)
+    indices = [
+        np.arange(len(frame_indices))[i : i + size]
+        for i in range(0, len(frame_indices), size)
+    ]
+    pool = mp.Pool(nodes)
+    results = pool.map(
+        partial(_calc_histogram, trajectory=trajectory, bins=bins), indices
+    )
+    pool.close()
+    matbin = np.sum(results, axis=0)
+    x = (x_bins[:-1] + x_bins[1:]) / 2
+    y = (y_bins[:-1] + y_bins[1:]) / 2
+    z = (z_bins[:-1] + z_bins[1:]) / 2
+
+    coords = np.array(np.meshgrid(x, y, z, indexing="ij"))
+    coords = np.array([x.flatten() for x in coords])
+    matbin_new = matbin.flatten()
+    occupation_df = pd.DataFrame(
+        {"x": coords[0], "y": coords[1], "z": coords[2], "occupation": matbin_new}
+    )
+    occupation_df = occupation_df.query("occupation != 0")
+    return occupation_df
+
+
+def _calc_histogram(numberlist, trajectory, bins):
+    matbin = None
+    for index in range(0, len(numberlist), 1000):
+        try:
+            indices = numberlist[index : index + 1000]
+        except IndexError:
+            indices = numberlist[index:]
+        frames = np.concatenate(np.array([trajectory.pbc[i] for i in indices]))
+        hist, _ = np.histogramdd(frames, bins=bins)
+        if matbin is None:
+            matbin = hist
+        else:
+            matbin += hist
+    return matbin
+
+
 def get_fel(
     traj,
     path,