From 3419258f0e5815ea108e6f162eb6b904bf1d1dac Mon Sep 17 00:00:00 2001
From: Chris Jones <50423140+chrisjonesBSU@users.noreply.github.com>
Date: Tue, 5 Mar 2024 13:19:42 -0700
Subject: [PATCH] Add concentration profile function for plotting diffusion
 across an interface (#84)

* add diffusion/conc profile function

* fix x_range returned

* add test for concentration profile

* rename box_edge param to box_axis

* add Note and example

* better param clarification in doc strings
---
 cmeutils/structure.py            | 69 ++++++++++++++++++++++++++++++++
 cmeutils/tests/base_test.py      |  9 +++++
 cmeutils/tests/test_structure.py | 15 +++++++
 3 files changed, 93 insertions(+)

diff --git a/cmeutils/structure.py b/cmeutils/structure.py
index 3ac6766..83fdb6f 100644
--- a/cmeutils/structure.py
+++ b/cmeutils/structure.py
@@ -726,6 +726,75 @@ def order_parameter(aa_gsd, cg_gsd, mapping, r_max, a_max, large=6, start=-10):
     return order, cl_idx
 
 
+def concentration_profile(snap, A_indices, B_indices, n_bins=70, box_axis=0):
+    """Calculate the concentration profile for two species
+    along a spatial dimension.
+
+    Parameters
+    ----------
+    snap : gsd.hoomd.Frame
+        A snapshot object containing particle information, including positions.
+    A_indices : numpy array
+        Indices of particles belonging to species A.
+    B_indices : numpy array
+        Indices of particles belonging to species B.
+    n_bins : int, optional, default 70
+        Number of bins for the concentration profile.
+    box_axis : int, optional, default 0
+        Index of the box edge that the concentration profile is calculated.
+        Options are 0, 1, or 2 which correspond to [x, y, z].
+
+    Returns
+    -------
+    d_profile : numpy array
+        Positions corresponding to the center of each bin
+        in the concentration profile.
+    A_count : numpy array
+        Particle count for species A in each bin.
+    B_count : numpy array
+        Particle count for species B in each bin.
+    total_count : numpy array
+        Total particle count in each bin.
+
+    Notes
+    -----
+    Use this to create a concentration profile plot of "left" species
+    and "right" species in the simulation's volume.
+
+    Example::
+        # Plot the concentration profile for a snapshot with 200 particles
+        # "left" species are particles 0-99 and "right" species are 100-199
+
+        from cmeutils.structure import concentration_profile
+        import matplotlib.pyplot as plt
+
+        x_range, left, right, total = concentration_profile(
+            snap=snapshot,
+            A_indices=range(0, 100),
+            B_indices=range(100, 200),
+            n_bins=50,
+            box_axis=0
+        )
+
+        plt.plot(x_range, left/total)
+        plt.plot(x_range, right/total)
+
+    """
+
+    L = snap.configuration.box[box_axis]
+    dl = L / n_bins
+    d_profile = np.linspace(-L / 2 + dl, L / 2, n_bins)
+
+    A_pos = snap.particles.position[A_indices, box_axis]
+    B_pos = snap.particles.position[B_indices, box_axis]
+    A_count, _ = np.histogram(A_pos, bins=d_profile, density=False)
+    B_count, _ = np.histogram(B_pos, bins=d_profile, density=False)
+
+    total_count = A_count + B_count
+
+    return d_profile[:-1], A_count, B_count, total_count
+
+
 def all_atom_rdf(
     gsdfile,
     start=0,
diff --git a/cmeutils/tests/base_test.py b/cmeutils/tests/base_test.py
index a75210b..fa9a3af 100644
--- a/cmeutils/tests/base_test.py
+++ b/cmeutils/tests/base_test.py
@@ -11,6 +11,15 @@
 
 
 class BaseTest:
+    @pytest.fixture
+    def slab_snapshot(self):
+        frame = gsd.hoomd.Frame()
+        A_positions = np.random.uniform(-5, 0, size=(20, 3))
+        B_positions = np.random.uniform(0, 5, size=(20, 3))
+        frame.particles.position = np.concatenate([A_positions, B_positions])
+        frame.configuration.box = np.array([10.0, 10.0, 10.0, 0, 0, 0])
+        return frame
+
     @pytest.fixture
     def gsdfile(self, tmp_path):
         filename = tmp_path / "test.gsd"
diff --git a/cmeutils/tests/test_structure.py b/cmeutils/tests/test_structure.py
index d57eb40..075b1a9 100644
--- a/cmeutils/tests/test_structure.py
+++ b/cmeutils/tests/test_structure.py
@@ -8,6 +8,7 @@
     all_atom_rdf,
     angle_distribution,
     bond_distribution,
+    concentration_profile,
     diffraction_pattern,
     dihedral_distribution,
     get_centers,
@@ -258,3 +259,17 @@ def test_get_centers(self, gsdfile):
         new_gsdfile = "centers.gsd"
         centers = get_centers(gsdfile, new_gsdfile)
         assert isinstance(centers, type(None))
+
+    def test_conc_profiel(self, slab_snapshot):
+        A_indices = np.arange(20)
+        B_indices = np.arange(20, 40)
+        d_profile, A_count, B_count, total_count = concentration_profile(
+            slab_snapshot, A_indices, B_indices, n_bins=5, box_axis=0
+        )
+        assert (
+            len(d_profile) == len(A_count) == len(B_count) == len(total_count)
+        )
+        assert (A_count / total_count)[0] == 1
+        assert (A_count / total_count)[-1] == 0
+        assert (B_count / total_count)[-1] == 1
+        assert (B_count / total_count)[0] == 0