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/__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..554acf9 --- /dev/null +++ b/pu_pjr/force_fields_plotting/__main__.py @@ -0,0 +1,181 @@ +import argparse +from rich.console import Console +from rich.text import Text + +from . import utils +from . import force_fields + +def main(): + parser = argparse.ArgumentParser( + description="Plot potentials.", + 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.18.0" + ) + + subparsers = parser.add_subparsers(title="subcommands", dest="which") + list_parser = subparsers.add_parser( + "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", + 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" + ) + 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( + "--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." + ) + + args = parser.parse_args() + # print(args) + + console = Console() + if args.which == "main": + 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"): + 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: + 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() + +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..624404a --- /dev/null +++ b/pu_pjr/force_fields_plotting/force_fields.py @@ -0,0 +1,91 @@ +from matplotlib import pyplot as plt + +from . import utils + +def get_field_data( + field: utils.Potentials, points: list[float], **kwargs +) -> list[float]: + """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.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 potential""" + + x = points + y = get_field_data(field, points, **kwargs) + # print(x, len(x)) + # print(y) + + # 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(draggable=True) + + # 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) + +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 new file mode 100644 index 0000000..0a496f5 --- /dev/null +++ b/pu_pjr/force_fields_plotting/utils.py @@ -0,0 +1,148 @@ +from enum import Enum +import math + +### Constants ### + +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 +) -> 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 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 TypeError(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: + 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/(r)) + + # print(min(values)) + # print(max(values)) + + return values \ No newline at end of file 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 f635f80..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,6 +20,7 @@ pytest = "^7.4.2" [tool.poetry.scripts] quickplot = "pu_pjr.plotting.__main__:main" dirstats = "pu_pjr.dir_stats.__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 new file mode 100644 index 0000000..904029f --- /dev/null +++ b/tests/test_force_fields_plotting.py @@ -0,0 +1,47 @@ +from pu_pjr.force_fields_plotting import utils + +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, + 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) + 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, 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