diff --git a/.github/workflows/Dependabot.yml b/.github/Dependabot.yml similarity index 100% rename from .github/workflows/Dependabot.yml rename to .github/Dependabot.yml diff --git a/pancax/__init__.py b/pancax/__init__.py index 708065a..d885420 100644 --- a/pancax/__init__.py +++ b/pancax/__init__.py @@ -1,6 +1,7 @@ # from .bcs import * from .bcs import DirichletBC, NeumannBC from .bvps import SimpleShearLinearRamp, UniaxialTensionLinearRamp +from .cli import pancax_main from .constitutive_models import \ ConstitutiveModel, \ BoundedProperty, \ @@ -112,6 +113,8 @@ "PronySeries", "SimpleFeFv", "WLF", + # cli + "pancax_main", # "data", "FullFieldData", "FullFieldDataLoader", diff --git a/pancax/__main__.py b/pancax/__main__.py deleted file mode 100644 index 74ba8b8..0000000 --- a/pancax/__main__.py +++ /dev/null @@ -1,130 +0,0 @@ -from importlib import import_module -from pathlib import Path -import argparse -import jax -import os -import sys - - -code_name = """ -██████╗ █████╗ ███╗ ██╗ ██████╗ █████╗ ██╗ ██╗ -██╔══██╗██╔══██╗████╗ ██║██╔════╝██╔══██╗╚██╗██╔╝ -██████╔╝███████║██╔██╗ ██║██║ ███████║ ╚███╔╝ -██╔═══╝ ██╔══██║██║╚██╗██║██║ ██╔══██║ ██╔██╗ -██║ ██║ ██║██║ ╚████║╚██████╗██║ ██║██╔╝ ██╗ -╚═╝ ╚═╝ ╚═╝╚═╝ ╚═══╝ ╚═════╝╚═╝ ╚═╝╚═╝ ╚═╝ - -Developed by Craig M. Hamel - -MIT License - -Copyright (c) 2024 Sandia National Laboratories - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. - -""" - - -parser = argparse.ArgumentParser( - # prog='pancax', - prog=code_name, - description='Physics Informed Neural Network library written in jax', - epilog='Reach out to chamel@sandia.gov for help' -) -parser.add_argument( - '-d', '--debug-nans', - # default='off', - action='store_true', - default=False, - help='Flag to debug nans. Options are on or off.' -) -parser.add_argument( - '-i', '--input-file', - help='Input file for pancax to run.' -) -parser.add_argument( - '-l', '--log-file', - default=os.path.join(os.getcwd(), 'pancax.log'), - help='Log file for pancax to write stdout and stderr to.' -) -parser.add_argument( - '-p', '--precision', - default='double', - help='Floating point precision to use. Options are single or double.' -) -parser.add_argument( - '-v', '--verbose', - action='store_true', - default=False, - help='Whether or not to print to console.' -) -args = parser.parse_args() - -print(code_name) - -# switch on debug state -# if args.debug_nans == 'on': -if args.debug_nans: - print('Debugging NaNs') - jax.config.update('jax_debug_nans', True) -else: - jax.config.update('jax_debug_nans', False) - -# switch on different precision levels -if args.precision == 'double': - print('Using double precision') - jax.config.update("jax_enable_x64", True) -elif args.precision == 'single': - print('Using single precision') -else: - raise ValueError('Precision needs to be \'single\' or \'double\'.') - -# read in the input file and make sure it exists -input_file = Path(args.input_file) -if not input_file.is_file(): - raise FileNotFoundError(f'Input file {input_file} does not exist.') - -log_file = Path(args.log_file) - -# switch over file types, e.g. python or yaml in the future -print(f'Running {input_file}') -print(f'Writing output to {log_file}') -# NOTE below will only work on linux/mac - -with open(log_file, 'w') as log: - # direct output to log - if not args.verbose: - sys.stdout = log - sys.stderr = log - print(code_name, flush=True) - # read input file and run it - if input_file.suffix == '.py': - try: - # with jax.profiler.trace("/tmp/jax-trace", - # create_perfetto_link=True): - import_module(str(input_file).replace('/', '.')) - except ModuleNotFoundError: - # things bug out currently if we use above in the bootstrapper - print('Finished running') - else: - raise ValueError('Only python files are supported currently!') - - # Reset stdout and stderr to their original values - sys.stdout = sys.__stdout__ - sys.stderr = sys.__stderr__ diff --git a/pancax/cli/__init__.py b/pancax/cli/__init__.py new file mode 100644 index 0000000..296e2fb --- /dev/null +++ b/pancax/cli/__init__.py @@ -0,0 +1,6 @@ +from .cli import pancax_main + + +__all__ = [ + "pancax_main" +] diff --git a/pancax/cli/cli.py b/pancax/cli/cli.py new file mode 100644 index 0000000..371cc83 --- /dev/null +++ b/pancax/cli/cli.py @@ -0,0 +1,134 @@ +from importlib import import_module +from pathlib import Path +import argparse +# import jax +import os +import sys + + +code_name = """ +██████╗ █████╗ ███╗ ██╗ ██████╗ █████╗ ██╗ ██╗ +██╔══██╗██╔══██╗████╗ ██║██╔════╝██╔══██╗╚██╗██╔╝ +██████╔╝███████║██╔██╗ ██║██║ ███████║ ╚███╔╝ +██╔═══╝ ██╔══██║██║╚██╗██║██║ ██╔══██║ ██╔██╗ +██║ ██║ ██║██║ ╚████║╚██████╗██║ ██║██╔╝ ██╗ +╚═╝ ╚═╝ ╚═╝╚═╝ ╚═══╝ ╚═════╝╚═╝ ╚═╝╚═╝ ╚═╝ + +Developed by Craig M. Hamel + +MIT License + +Copyright (c) 2024 Sandia National Laboratories + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + +""" + + +def pancax_main(): + parser = argparse.ArgumentParser( + prog='pancax', + # prog=code_name, + description='Physics Informed Neural Network library written in jax', + epilog='Reach out to chamel@sandia.gov for help' + ) + parser.add_argument( + '-d', '--debug-nans', + # default='off', + action='store_true', + default=False, + help='Flag to debug nans. Options are on or off.' + ) + parser.add_argument( + '-i', '--input-file', + help='Input file for pancax to run.' + ) + parser.add_argument( + '-l', '--log-file', + default=os.path.join(os.getcwd(), 'pancax.log'), + help='Log file for pancax to write stdout and stderr to.' + ) + parser.add_argument( + '-p', '--precision', + default='double', + help='Floating point precision to use. Options are single or double.' + ) + parser.add_argument( + '-v', '--verbose', + action='store_true', + default=False, + help='Whether or not to print to console.' + ) + args = parser.parse_args() + + print(code_name) + + # import jax here to speed up --help + import jax + + # switch on debug state + # if args.debug_nans == 'on': + if args.debug_nans: + print('Debugging NaNs') + jax.config.update('jax_debug_nans', True) + else: + jax.config.update('jax_debug_nans', False) + + # switch on different precision levels + if args.precision == 'double': + print('Using double precision') + jax.config.update("jax_enable_x64", True) + elif args.precision == 'single': + print('Using single precision') + else: + raise ValueError('Precision needs to be \'single\' or \'double\'.') + + # read in the input file and make sure it exists + input_file = Path(args.input_file) + if not input_file.is_file(): + raise FileNotFoundError(f'Input file {input_file} does not exist.') + + log_file = Path(args.log_file) + + # switch over file types, e.g. python or yaml in the future + print(f'Running {input_file}') + print(f'Writing output to {log_file}') + # NOTE below will only work on linux/mac + + with open(log_file, 'w') as log: + # direct output to log + if not args.verbose: + sys.stdout = log + sys.stderr = log + print(code_name, flush=True) + # read input file and run it + if input_file.suffix == '.py': + try: + # with jax.profiler.trace("/tmp/jax-trace", + # create_perfetto_link=True): + import_module(str(input_file).replace('/', '.')) + except ModuleNotFoundError: + # things bug out currently if we use above in the bootstrapper + print('Finished running') + else: + raise ValueError('Only python files are supported currently!') + + # Reset stdout and stderr to their original values + sys.stdout = sys.__stdout__ + sys.stderr = sys.__stderr__ diff --git a/pyproject.toml b/pyproject.toml index 4655f9d..5e52928 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = 'pancax' -version = '0.0.12' +version = '0.0.13' authors = [ {name = 'Craig M. Hamel, email = '} ] @@ -16,15 +16,15 @@ dependencies = [ [project.optional-dependencies] cpu = [ 'chex', - 'equinox==0.12.1', - 'jax==0.5.0', + 'equinox==0.13', + 'jax==0.6.2', 'jaxtyping', 'optax' ] cuda = [ 'chex', - 'equinox==0.12.1', - 'jax[cuda12]==0.5.0', + 'equinox==0.13', + 'jax[cuda12]==0.6.2', 'jaxtyping', 'optax' ] @@ -54,6 +54,9 @@ viz = [ requires = ["hatchling"] build-backend = "hatchling.build" +[project.scripts] +pancax = "pancax.cli:pancax_main" + # [tool.setuptools] [tool.hatch.build.targets.wheel] packages = [