From b981c9ee3a495f11c87f5aa08dde92dcaef50c68 Mon Sep 17 00:00:00 2001 From: Pedro Juan Royo Date: Tue, 3 Oct 2023 23:09:49 +0100 Subject: [PATCH 1/4] 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 From 3fc2589be0b7683d49de36f8a8b39b5ad7b838ab Mon Sep 17 00:00:00 2001 From: Pedro Juan Royo Date: Wed, 4 Oct 2023 15:18:03 +0100 Subject: [PATCH 2/4] 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 From 8c4b8d3fe4799691bacfbe9af2a755f86dacfa7f Mon Sep 17 00:00:00 2001 From: Pedro Juan Royo Date: Thu, 5 Oct 2023 00:13:07 +0100 Subject: [PATCH 3/4] feat: fix coulomb term in buck_coul potential --- pu_pjr/force_fields_plotting/utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/pu_pjr/force_fields_plotting/utils.py b/pu_pjr/force_fields_plotting/utils.py index af3109b..2a5d491 100644 --- a/pu_pjr/force_fields_plotting/utils.py +++ b/pu_pjr/force_fields_plotting/utils.py @@ -2,7 +2,6 @@ import math ### 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' @@ -121,7 +120,7 @@ def buck( values = [] for r in points: - values.append(A * math.exp(-r/rho) - C / r**6) + values.append(A * math.exp(-r/rho) - (C / r**6)) return values @@ -140,9 +139,10 @@ def buck_coul( values = [] for r in points: - values.append(A * math.exp(-r/rho) - C / r**6 + - q1*q2/(4*math.pi*EPSILON_0*r)) + values.append(A * math.exp(-r/rho) - C / r**6 + + q1*q2/(r)) - print(min(values)) + # print(min(values)) + # print(max(values)) return values \ No newline at end of file From 2f004c6ddfea1b3d9bc5a717a31454c34987205f Mon Sep 17 00:00:00 2001 From: Pedro Juan Royo Date: Thu, 5 Oct 2023 19:42:38 +0100 Subject: [PATCH 4/4] feat: Can read potential from file and plt in one --- pu_pjr/dir_stats/__main__.py | 2 +- pu_pjr/force_fields_plotting/__main__.py | 122 ++++++++++++++++--- pu_pjr/force_fields_plotting/force_fields.py | 41 ++++++- pu_pjr/force_fields_plotting/utils.py | 4 +- pu_pjr/plotting/__main__.py | 2 +- pyproject.toml | 4 +- tests/test_files/potentials.pot | 5 + tests/test_force_fields_plotting.py | 35 +++++- 8 files changed, 189 insertions(+), 26 deletions(-) create mode 100644 tests/test_files/potentials.pot diff --git a/pu_pjr/dir_stats/__main__.py b/pu_pjr/dir_stats/__main__.py index 59c23ab..468908f 100644 --- a/pu_pjr/dir_stats/__main__.py +++ b/pu_pjr/dir_stats/__main__.py @@ -15,7 +15,7 @@ def main(): parser.set_defaults(which="main") parser.add_argument( - "--version", "-v", action="version", version="%(prog)s v0.17.0" + "--version", "-v", action="version", version="%(prog)s v0.18.0" ) subparser = parser.add_subparsers(required=True) diff --git a/pu_pjr/force_fields_plotting/__main__.py b/pu_pjr/force_fields_plotting/__main__.py index b436041..554acf9 100644 --- a/pu_pjr/force_fields_plotting/__main__.py +++ b/pu_pjr/force_fields_plotting/__main__.py @@ -8,26 +8,33 @@ def main(): parser = argparse.ArgumentParser( description="Plot potentials.", - prog="fields", + prog="potentials", 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" + "--version", "-v", action="version", version="%(prog)s v0.18.0" ) subparsers = parser.add_subparsers(title="subcommands", dest="which") list_parser = subparsers.add_parser( - "list", help="list available force fields" + "list", help="list available force fields", + epilog="Created by Pedro Juan Royo @UnstrayCato" ) list_parser.set_defaults(which="list") plot_parser = subparsers.add_parser( - "show", help="plot a force field" + "show", help="plot a force field", + epilog="Created by Pedro Juan Royo @UnstrayCato" ) plot_parser.set_defaults(which="main") + plot_parser.add_argument( + "--together", "-t", action="store_true", + help = "Plot all potentials together, if not the potentials will" + + " be plotted one by one default False" + ) 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" @@ -53,7 +60,17 @@ def main(): help = "Line type for plot, default 'o-'" ) plot_parser.add_argument( - "potential_data", nargs="+", + "--from-file", "-f", action="store_true", dest="read_from_file", + #metavar="kwargs", + help = "Read arguments from file, ignores other potential_data arguments" + ) + plot_parser.add_argument( + "--file-name", "-F", type=str, default="potentials.pot", + #metavar="kwargs", + help = "File name to read arguments from, default 'potentials.pot'" + ) + 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." @@ -64,21 +81,90 @@ def main(): console = Console() if args.which == "main": - keyargpairs = utils.parse_keyargpairs(args.potential_data) - # print(keyargpairs) + if args.read_from_file: + try: + with open(args.file_name, "r") as f: + # Remove newlines and split by spaces, eliminate multiple spaces + args.potential_data = " ".join(f.read().replace("\n", " ").split())\ + .split(" ") + except FileNotFoundError as e: + console.print(f"⛔️[bold red]FILE NOT FOUND ERROR: " + f"[bold yellow][underline]{e.filename}[/underline]" + f" is not a valid file") + exit(4) + except PermissionError as e: + console.print(f"⛔️[bold red]PERMISSION ERROR: " + f"[bold yellow][underline]{e.filename}[/underline]" + f" cannot be read") + exit(5) + + # print(args.potential_data) + + try: + keyargpairs = utils.parse_keyargpairs(args.potential_data) + # print(keyargpairs) + except ValueError as e: + console.print(f"⛔️[bold red]VALUE ERROR: [bold yellow]{e}") + exit(1) + except TypeError as e: + console.print(f"⛔️[bold red]TYPE ERROR: [bold yellow]{e}") + exit(3) + except KeyError as e: + console.print(f"⛔️[bold red]KEY ERROR: " + f"[bold yellow][underline]{e}[/underline]" + f" is not a valid potential name") + exit(2) + + try: + points = utils.create_range(args.range_min, + args.range_max, + args.range_points) + except ValueError as e: + console.print(f"⛔️[bold red]VALUE ERROR: [bold yellow]{e}") + exit(1) + + try: + potentials = [utils.Potentials[utils.remove_extras_for_class(field)] \ + for field in keyargpairs] + potential_args = [keyargpairs[field] for field in keyargpairs] + except KeyError as e: + console.print(f"⛔️[bold red]KEY ERROR: " + f"[bold yellow][underline]{e}[/underline]" + f" is not a valid potential name") + exit(2) with console.status("[bold green]Plot created"): - # 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]) + if args.together: + # potentials = [utils.Potentials[utils.remove_extras_for_class(field)] \ + # for field in keyargpairs] + # potential_args = [keyargpairs[field] for field in keyargpairs] + try: + force_fields.plot_fields_single(potentials, + potential_args, + points, + args.line_type, + args.y_range) + except ValueError as e: + console.print(f"⛔️[bold red]VALUE ERROR: [bold yellow]{e}") + exit(1) + except TypeError as e: + console.print(f"⛔️[bold red]TYPE ERROR: [bold yellow]{e}") + exit(3) + else: + for pos in range(len(potentials)): + try: + force_fields.plot_field(potentials[pos], + points, + args.line_type, + args.y_range, + **potential_args[pos]) + except ValueError as e: + console.print(f"⛔️[bold red]VALUE ERROR: [bold yellow]{e}") + exit(1) + except TypeError as e: + console.print(f"⛔️[bold redTYPE ERROR: [bold yellow]{e}") + exit(3) + elif args.which == "list": descriptions, max_length = utils.field_descriptions() for field in descriptions: diff --git a/pu_pjr/force_fields_plotting/force_fields.py b/pu_pjr/force_fields_plotting/force_fields.py index 080154c..624404a 100644 --- a/pu_pjr/force_fields_plotting/force_fields.py +++ b/pu_pjr/force_fields_plotting/force_fields.py @@ -38,7 +38,7 @@ def plot_field( # 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() + plt.legend(draggable=True) # Set y range plt.ylim(y_range) @@ -49,4 +49,43 @@ def plot_field( plt.xlabel("r") plt.ylabel("V(r)") + plt.show(block=True) + +def plot_fields_single( + fields: list[utils.Potentials], + args: list[dict[str, float]], + points: list[float] = utils.create_range(), + line_type: str = 'o-', + y_range: list[float] = [-5.0, 10.0], +) -> None: + """Plot a list of potentials on the same plot""" + + # Check that the number of fields and args match + if len(fields) != len(args): + raise ValueError(f"Number of fields ({len(fields)}) " + f"and args ({len(args)}) must match") + + for pos in range(len(fields)): + # plot_field(fields[pos], points, line_type, y_range, **args[pos]) + x = points + y = get_field_data(fields[pos], points, **args[pos]) + + # line at y = 0 + plt.axhline(y=0, color='k') + + line = plt.plot(x, y, line_type, + label=f"{utils.remove_extras_from_field_name(fields[pos].name)}:" + f" {args[pos]}") + + # Assign the same colour to the vertical line as the plot line + if 'cut' in args[pos]: + plt.axvline(x=args[pos]['cut'], color=line[0].get_color(), linestyle='--') + + + plt.legend(draggable=True) + plt.ylim(y_range) + + 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 2a5d491..0a496f5 100644 --- a/pu_pjr/force_fields_plotting/utils.py +++ b/pu_pjr/force_fields_plotting/utils.py @@ -25,7 +25,7 @@ def create_range( ) -> list[float]: """Return a float range between 'start' and 'end'""" # Check end is bigger than start - if start > end: + if start >= end: raise ValueError("start value is bigger than end value") # Check that items is bigger than 1 @@ -75,7 +75,7 @@ def check_inputs_non_neg( ) -> 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}") + raise TypeError(f"Input must be a float: {input}") if input < 0: raise ValueError(f"Input must be non-negative: {input}") diff --git a/pu_pjr/plotting/__main__.py b/pu_pjr/plotting/__main__.py index 656a02f..1562b36 100644 --- a/pu_pjr/plotting/__main__.py +++ b/pu_pjr/plotting/__main__.py @@ -15,7 +15,7 @@ def main(): parser.set_defaults(which="main") parser.add_argument( - "--version", "-v", action="version", version="%(prog)s v0.17.0" + "--version", "-v", action="version", version="%(prog)s v0.18.0" ) # Sub-parser for the "xy" command diff --git a/pyproject.toml b/pyproject.toml index 82aecb3..1ef5897 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "pu-pjr" -version = "0.17.0" +version = "0.18.0" description = "Personal utilities for coding and data analysis" authors = ["Pedro Juan Royo "] readme = "README.md" @@ -20,7 +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" +potentials = "pu_pjr.force_fields_plotting.__main__:main" [build-system] requires = ["poetry-core"] diff --git a/tests/test_files/potentials.pot b/tests/test_files/potentials.pot new file mode 100644 index 0000000..d7c0972 --- /dev/null +++ b/tests/test_files/potentials.pot @@ -0,0 +1,5 @@ +buck_coul A=4545 rho=0.261 C=0 q1=1.939 q2=-0.96 +buck A=4545 rho=0.26 C=0 +lj_cut epsilon=1 sigma=0.6 cut=2.7 +lj_cut epsilon=1 sigma=2 cut=4 +lj_cut epsilon=1 sigma=1.5 cut=4.4 \ No newline at end of file diff --git a/tests/test_force_fields_plotting.py b/tests/test_force_fields_plotting.py index efd3a77..904029f 100644 --- a/tests/test_force_fields_plotting.py +++ b/tests/test_force_fields_plotting.py @@ -2,6 +2,13 @@ import pytest +def test_field_descriptions(): + descriptions, max_length = utils.field_descriptions() + assert max_length == 9 + assert "lj_cut" in descriptions + assert "buck" in descriptions + assert "buck_coul" in descriptions + 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, @@ -10,5 +17,31 @@ def test_create_range(): 2.8, 2.9, 3.0] == \ utils.create_range() assert [0, 1.0] == utils.create_range(0.0, 1.0, 2) + assert [0.0, 0.5, 1.0] == utils.create_range(0.0, 1.0, 3) + with pytest.raises(ValueError): + utils.create_range(1, 0) + with pytest.raises(ValueError): - utils.create_range(1, 0) \ No newline at end of file + utils.create_range(1, 1) + +def test_remove_extras_from_field_name(): + assert "lj/cut" == utils.remove_extras_from_field_name("lj_cut") + assert "buck" == utils.remove_extras_from_field_name("buck") + assert "buck/coul" == utils.remove_extras_from_field_name("buck_coul") + +def test_remove_extras_for_class(): + assert "lj_cut" == utils.remove_extras_for_class("lj_cut") + assert "buck" == utils.remove_extras_for_class("buck") + assert "buck_coul" == utils.remove_extras_for_class("buck_coul-4") + assert "lj_cut" == utils.remove_extras_for_class("lj_cut") + assert "buck" == utils.remove_extras_for_class("buck-6") + assert "buck_coul" == utils.remove_extras_for_class("buck_coul") + +def test_parse_keyargpairs(): + assert {"lj_cut": {"epsilon": 1.0, "sigma": 2.0, "cut": 3.0}} == \ + utils.parse_keyargpairs(["lj_cut", "epsilon=1.0", "sigma=2.0", "cut=3.0"]) + assert {"buck": {"A": 1.0, "rho": 2.0, "C": 3.0}} == \ + utils.parse_keyargpairs(["buck", "A=1.0", "rho=2.0", "C=3.0"]) + assert {"buck_coul": {"A": 1.0, "rho": 2.0, "C": 3.0, "q1": 4.0, "q2": 5.0}} == \ + utils.parse_keyargpairs(["buck_coul", "A=1.0", "rho=2.0", "C=3.0", "q1=4.0", + "q2=5.0"]) \ No newline at end of file