From 3fc2589be0b7683d49de36f8a8b39b5ad7b838ab Mon Sep 17 00:00:00 2001 From: Pedro Juan Royo Date: Wed, 4 Oct 2023 15:18:03 +0100 Subject: [PATCH] feat: buck potentials look weird? --- pu_pjr/force_fields_plotting/__main__.py | 72 ++++++++++-- pu_pjr/force_fields_plotting/force_fields.py | 38 ++++-- pu_pjr/force_fields_plotting/utils.py | 117 ++++++++++++++++++- 3 files changed, 210 insertions(+), 17 deletions(-) diff --git a/pu_pjr/force_fields_plotting/__main__.py b/pu_pjr/force_fields_plotting/__main__.py index 019832b..b436041 100644 --- a/pu_pjr/force_fields_plotting/__main__.py +++ b/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() diff --git a/pu_pjr/force_fields_plotting/force_fields.py b/pu_pjr/force_fields_plotting/force_fields.py index 802a055..080154c 100644 --- a/pu_pjr/force_fields_plotting/force_fields.py +++ b/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) \ No newline at end of file diff --git a/pu_pjr/force_fields_plotting/utils.py b/pu_pjr/force_fields_plotting/utils.py index ef76801..af3109b 100644 --- a/pu_pjr/force_fields_plotting/utils.py +++ b/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 \ No newline at end of file