From b981c9ee3a495f11c87f5aa08dde92dcaef50c68 Mon Sep 17 00:00:00 2001
From: Pedro Juan Royo <pedro.juan.royo@gmail.com>
Date: Tue, 3 Oct 2023 23:09:49 +0100
Subject: [PATCH] feat: plottings fields start

---
 pu_pjr/force_fields_plotting/__init__.py     |  0
 pu_pjr/force_fields_plotting/__main__.py     | 37 ++++++++++++++++++++
 pu_pjr/force_fields_plotting/force_fields.py | 28 +++++++++++++++
 pu_pjr/force_fields_plotting/utils.py        | 37 ++++++++++++++++++++
 pyproject.toml                               |  1 +
 tests/test_force_fields_plotting.py          | 14 ++++++++
 6 files changed, 117 insertions(+)
 create mode 100644 pu_pjr/force_fields_plotting/__init__.py
 create mode 100644 pu_pjr/force_fields_plotting/__main__.py
 create mode 100644 pu_pjr/force_fields_plotting/force_fields.py
 create mode 100644 pu_pjr/force_fields_plotting/utils.py
 create mode 100644 tests/test_force_fields_plotting.py

diff --git a/pu_pjr/force_fields_plotting/__init__.py b/pu_pjr/force_fields_plotting/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/pu_pjr/force_fields_plotting/__main__.py b/pu_pjr/force_fields_plotting/__main__.py
new file mode 100644
index 0000000..019832b
--- /dev/null
+++ b/pu_pjr/force_fields_plotting/__main__.py
@@ -0,0 +1,37 @@
+import argparse
+from rich.console import Console
+
+from . import utils
+from . import force_fields
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Plot potentials.",
+        prog="fields",
+        epilog="Created by Pedro Juan Royo @UnstrayCato"
+    )
+    parser.set_defaults(which="main")
+
+    parser.add_argument(
+        "--version", "-v", action="version", version="%(prog)s v0.17.0"
+    )
+
+    available_fields = [a.value for a in utils.Fields]
+    parser.add_argument(
+        "--field", "-f", nargs="*", choices=available_fields,
+        help="Force field to plot"
+    )
+
+    args = parser.parse_args()
+    print(args)
+
+    console = Console()
+    if args.which == "main":
+        with console.status("[bold green]Plot created"):
+            force_fields.plot_field(points = utils.create_range(0.9,3,60),
+                                    epsilon = 1, sigma = 1.0, cut = 3.5)
+    else:
+        parser.print_help()
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/pu_pjr/force_fields_plotting/force_fields.py b/pu_pjr/force_fields_plotting/force_fields.py
new file mode 100644
index 0000000..802a055
--- /dev/null
+++ b/pu_pjr/force_fields_plotting/force_fields.py
@@ -0,0 +1,28 @@
+from matplotlib import pyplot as plt
+
+from . import utils
+
+def get_field_data(
+    field: utils.Fields, points: list[float], **kwargs
+) -> list[float]:
+    """Get the values for some points in a field"""
+    if field == utils.Fields.LJ_CUT:
+        # print("lj/cut field")
+        return utils.lj_cut(points, **kwargs)
+    else:
+        raise NotImplementedError(f'Field {field} not implemented')
+    
+def plot_field(
+    field: utils.Fields = utils.Fields.LJ_CUT, 
+    points: list[float] = utils.create_range(), 
+    **kwargs
+) -> None:
+    """Plot a field"""
+
+    x = points
+    y = get_field_data(field, points, **kwargs)
+    # print(x, len(x))
+    # print(y)
+
+    plt.plot(x, y, 'o-')
+    plt.show(block=True)
\ No newline at end of file
diff --git a/pu_pjr/force_fields_plotting/utils.py b/pu_pjr/force_fields_plotting/utils.py
new file mode 100644
index 0000000..ef76801
--- /dev/null
+++ b/pu_pjr/force_fields_plotting/utils.py
@@ -0,0 +1,37 @@
+from enum import Enum
+
+class Fields(Enum):
+    LJ_CUT = 'lj/cut'
+
+def create_range(
+    start: float = 0.1, end: float = 3.0, items: int = 30
+) -> list[float]:
+    """Return a float range between 'start' and 'end'"""
+    # Check end is bigger than start
+    if start > end:
+        raise ValueError("start value is bigger than end value")
+    
+    # Check that items is bigger than 1
+    if items <= 1:
+        raise ValueError("items value must be bigger than 1")
+
+    step_val = (end - start) / (items - 1)
+    range_vals = []
+    for i in range(items):
+        range_vals.append(start + step_val*i)
+
+    return range_vals
+
+def lj_cut(
+    points: list[float], epsilon: float, sigma: float, cut: float
+) -> list[float | None]:
+    """Lennard-Jones potential with a cut-off at `cut`"""
+
+    values = []
+    for r in points:
+        if r < cut:
+            values.append(4 * epsilon * ((sigma / r) ** 12 - (sigma / r) ** 6))
+        else:
+            values.append(None)
+
+    return values
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index f635f80..82aecb3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,6 +20,7 @@ pytest = "^7.4.2"
 [tool.poetry.scripts]
 quickplot = "pu_pjr.plotting.__main__:main"
 dirstats = "pu_pjr.dir_stats.__main__:main"
+fields = "pu_pjr.force_fields_plotting.__main__:main"
 
 [build-system]
 requires = ["poetry-core"]
diff --git a/tests/test_force_fields_plotting.py b/tests/test_force_fields_plotting.py
new file mode 100644
index 0000000..efd3a77
--- /dev/null
+++ b/tests/test_force_fields_plotting.py
@@ -0,0 +1,14 @@
+from pu_pjr.force_fields_plotting import utils
+
+import pytest
+
+def test_create_range():
+    assert [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.7999999999999999, 0.8999999999999999, 
+            0.9999999999999999, 1.0999999999999999, 1.2, 1.3, 1.4, 1.5, 
+            1.5999999999999999, 1.7, 1.8, 1.9, 2.0, 2.0999999999999996, 
+            2.1999999999999997, 2.3, 2.4, 2.5, 2.6, 2.6999999999999997, 
+            2.8, 2.9, 3.0] == \
+            utils.create_range()
+    assert [0, 1.0] == utils.create_range(0.0, 1.0, 2)
+    with pytest.raises(ValueError):
+        utils.create_range(1, 0)
\ No newline at end of file