feat: buck potentials look weird?

pu_pjr/force_fields_plotting/__main__.py

@@ -1,5 +1,6 @@
 import argparse
 from rich.console import Console
+from rich.text import Text
 
 from . import utils
 from . import force_fields
@@ -16,20 +17,77 @@ def main():
         "--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"
+    subparsers = parser.add_subparsers(title="subcommands", dest="which")
+    list_parser = subparsers.add_parser(
+        "list", help="list available force fields"
+    )
+    list_parser.set_defaults(which="list")
+
+    plot_parser = subparsers.add_parser(
+        "show", help="plot a force field"
+    )
+    plot_parser.set_defaults(which="main")
+
+    plot_parser.add_argument(
+        "--y-range", "-y", type=float, nargs=2, default=[-5.0, 10.0],
+        help = "Min value for range in plot points, default 0.9"
+    )
+    plot_parser.add_argument(
+        "--range-min", "-r", type=float, default=0.9,
+        #metavar="kwargs",
+        help = "Min value for range in plot points, default 0.9"
+    )
+    plot_parser.add_argument(
+        "--range-max", "-R", type=float, default=3.0,
+        #metavar="kwargs",
+        help = "Min value for range in plot points, default 3.0"
+    )
+    plot_parser.add_argument(
+        "--range-points", "-p", type=int, default=30,
+        #metavar="kwargs",
+        help = "Number of points in range, default 30"
+    )
+    plot_parser.add_argument(
+        "--line-type", "-l", type=str, default="o-",
+        #metavar="kwargs",
+        help = "Line type for plot, default 'o-'"
+    )
+    plot_parser.add_argument(
+        "potential_data", nargs="+",
+        #metavar="kwargs",
+        help = "Force field data to plot POTENTIAL must be a valid field name and DATA"+
+            "must be valid keyword arguments for the potential."
     )
 
     args = parser.parse_args()
-    print(args)
+    # print(args)
 
     console = Console()
     if args.which == "main":
+        keyargpairs = utils.parse_keyargpairs(args.potential_data)
+        # print(keyargpairs)
+
         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)
+            # force_fields.plot_field(points = utils.create_range(0.9,3,60),
+            #                         epsilon = 1.0, sigma = 1.0, cut = 3.5)
+            for field in keyargpairs:
+                points = utils.create_range(args.range_min, 
+                                            args.range_max, 
+                                            args.range_points)
+                force_fields.plot_field(utils.Potentials[utils.remove_extras_for_class(field)], 
+                                        points, 
+                                        args.line_type,
+                                        args.y_range,
+                                        **keyargpairs[field])
+    elif args.which == "list":
+        descriptions, max_length = utils.field_descriptions()
+        for field in descriptions:
+            text = Text(f"{field.replace('/', '_'):>{max_length}}: "
+                        f"{descriptions[field]}")
+            text.stylize("bold green", 0, max_length)
+            text.stylize("bold blue", max_length + 2, len(text))
+            text.highlight_regex("ARGS:.*", "bold yellow")
+            console.print(text)
     else:
         parser.print_help()
 

pu_pjr/force_fields_plotting/force_fields.py

@@ -3,26 +3,50 @@
 from . import utils
 
 def get_field_data(
-    field: utils.Fields, points: list[float], **kwargs
+    field: utils.Potentials, 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")
+    """Get the values for some points in a potential"""
+    if field == utils.Potentials.lj_cut:
+        # print("lj/cut potential")
         return utils.lj_cut(points, **kwargs)
+    elif field == utils.Potentials.buck:
+        # print("buck potential")
+        return utils.buck(points, **kwargs)
+    elif field == utils.Potentials.buck_coul:
+        # print("buck_coul potential")
+        return utils.buck_coul(points, **kwargs)
     else:
         raise NotImplementedError(f'Field {field} not implemented')
 
 def plot_field(
-    field: utils.Fields = utils.Fields.LJ_CUT, 
+    field: utils.Potentials = utils.Potentials.lj_cut, 
     points: list[float] = utils.create_range(), 
+    line_type: str = 'o-',
+    y_range: list[float] = [-5.0, 10.0],
     **kwargs
 ) -> None:
-    """Plot a field"""
+    """Plot a potential"""
 
     x = points
     y = get_field_data(field, points, **kwargs)
     # print(x, len(x))
     # print(y)
 
-    plt.plot(x, y, 'o-')
+    # line at y = 0
+    plt.axhline(y=0, color='k')
+
+    # if 'cut' in kwargs add vertical line at cut with label
+    if 'cut' in kwargs:
+        plt.axvline(x=kwargs['cut'], color='k', linestyle='--', label='cut')
+        plt.legend()
+
+    # Set y range 
+    plt.ylim(y_range)
+
+    plt.title(f"{utils.remove_extras_from_field_name(field.name)}: {kwargs}")
+    plt.plot(x, y, line_type)
+
+    plt.xlabel("r")
+    plt.ylabel("V(r)")
+
     plt.show(block=True)

pu_pjr/force_fields_plotting/utils.py

@@ -1,7 +1,25 @@
 from enum import Enum
+import math
 
-class Fields(Enum):
-    LJ_CUT = 'lj/cut'
+### Constants ###
+EPSILON_0 = 55.26349406 # e^2*eV^-1*um^-1
+
+class Potentials(Enum):
+    lj_cut = 'Lennard-Jones potential with a cut-off. ARGS: epsilon, sigma, cut'
+    buck = 'Buckingham potential. ARGS: A, rho, C'
+    buck_coul = 'Buckingham potential with Coulombic term. ARGS: A, rho, C, q1, q2'
+
+def field_descriptions() -> (dict, int):
+    """Return a description for all potentials"""
+    descpritions = {}
+    max_length = 0
+    for field in Potentials:
+        if len(field.name) > max_length:
+            max_length = len(field.name)
+
+        descpritions[field.name] = field.value
+
+    return descpritions, max_length
 
 def create_range(
     start: float = 0.1, end: float = 3.0, items: int = 30
@@ -22,16 +40,109 @@ def create_range(
 
     return range_vals
 
+def remove_extras_from_field_name(
+    field_name: str
+) -> str:
+    """Remove extra characters from a potential name"""
+    return field_name.split("-")[0].replace("_", "/")
+
+def remove_extras_for_class(
+    field_name: str
+) -> str:
+    """Remove extra characters from a potential name"""
+    return field_name.split("-")[0]
+
+def parse_keyargpairs(
+    keyargpairs: list[str]
+) -> dict[str, str]:
+    """Parse a list of key-argument pairs"""
+    keyargdict = {}
+    current_field = None
+    for pos,keyargpair in enumerate(keyargpairs):
+        if "=" not in keyargpair:
+            current_field = keyargpair
+            if current_field in keyargdict:
+                current_field = current_field + "-" + str(pos)
+
+            keyargdict[current_field] = {}
+        else:
+            key, arg = keyargpair.split("=")
+            keyargdict[current_field][key] = float(arg)
+
+    return keyargdict
+
+def check_inputs_non_neg(
+    input: float,
+) -> None:
+    """Check that the input is a float and non-negative"""
+    if not isinstance(input, float):
+        raise ValueError(f"Input must be a float: {input}")
+
+    if input < 0:
+        raise ValueError(f"Input must be non-negative: {input}")
+
+def check_inputs(
+    input: float,
+) -> None:
+    """Check that the input is a float and non-negative"""
+    if not isinstance(input, float):
+        raise ValueError(f"Input must be a float: {input}")
+
+### Potentials ###
+
 def lj_cut(
     points: list[float], epsilon: float, sigma: float, cut: float
 ) -> list[float | None]:
     """Lennard-Jones potential with a cut-off at `cut`"""
 
+    # Check inputs
+    check_inputs_non_neg(epsilon)
+    check_inputs_non_neg(sigma)
+    check_inputs_non_neg(cut)
+
     values = []
     for r in points:
-        if r < cut:
+        if r <= cut:
             values.append(4 * epsilon * ((sigma / r) ** 12 - (sigma / r) ** 6))
         else:
             values.append(None)
 
+    return values
+
+def buck(
+    points: list[float], A: float, rho: float, C: float
+) -> list[float | None]:
+    """Buckingham potential"""
+
+    # Check inputs
+    check_inputs_non_neg(A)
+    check_inputs_non_neg(rho)
+    check_inputs_non_neg(C)
+
+    values = []
+    for r in points:
+        values.append(A * math.exp(-r/rho) - C / r**6)
+
+    return values
+
+def buck_coul(
+    points: list[float], A: float, rho: float, C: float, 
+    q1: float = 1.0, q2: float = -1.0
+) -> list[float | None]:
+    """Buckingham potential with Coulombic term"""
+
+    # Check inputs
+    check_inputs_non_neg(A)
+    check_inputs_non_neg(rho)
+    check_inputs_non_neg(C)
+    check_inputs(q1)
+    check_inputs(q2)
+
+    values = []
+    for r in points:
+        values.append(A * math.exp(-r/rho) - C / r**6 + 
+                      q1*q2/(4*math.pi*EPSILON_0*r))
+
+    print(min(values))
+
     return values