pyproject.toml

[build-system]
requires = ['flit_core >=3.4,<4']
build-backend = 'flit_core.buildapi'

[project]
name = 'tensorial'
dynamic = ["version", "description"]
authors = [
    { name = 'Martin Uhrin', email = 'martin@digitalbaker.net' },
]
readme = 'README.rst'
license = { file = 'LICENSE.txt' }
classifiers = [
    'Development Status :: 4 - Beta',
    'License :: OSI Approved :: GNU Lesser General Public License v3 or later (LGPLv3+)',
    'Programming Language :: Python',
    'Programming Language :: Python :: 3.10',
    'Programming Language :: Python :: 3.11',
    'Programming Language :: Python :: 3.12',
    'Programming Language :: Python :: 3.13',
    'Programming Language :: Python :: 3.14',
]
keywords = ['machine learning', 'e3nn-jax', 'physics']
requires-python = '>=3.10'
dependencies = [
    "ase",
    'beartype',
    "flax>=0.10.7",
    "hydra-core",
    'jax<0.9.0', # Control the JAX version, as previous updates have broken things
    'jaxlib<0.9.0',
    'jaxtyping<0.4.0',
    'jraph',
    'orbax-checkpoint',
    'pytray',
    'e3nn-jax',
    "equinox>=0.12.0",
    "reax>=0.6.2,<0.7",
    "rich",
    "tensorboardX",
    "typing-extensions"
]
[project.scripts]
tensorial = "tensorial.reaxkit.cli:main_cli"

[project.urls]
Home = 'https://github.com/camml-lab/tensorial'
Source = 'https://github.com/camml-lab/tensorial'

[project.optional-dependencies]
docs = [
    'docutils',
    'jupyter', # For running doc examples
    'nbsphinx', # Jupyter notebooks in docs
    'pandoc',
    'pydata_sphinx_theme',
    'sphinx',
    'sphinx-autobuild',
]
dev = [
    'flit',
    'ipython',
    'pandas', # Needed because we optionally support this in some files
    'pytest',
    'pytest-cov',
    'pytest-env',
    "pytest-codspeed",
    'pre-commit',
    'pylint',
    'black',
    'tensorboardX'
]

[tool.flit.module]
name = 'tensorial'

[project.entry-points."reax.plugins.batch_sizers"]
native = "tensorial._provides:get_batch_sizers"

[tool.flit.sdist]
exclude = [
    '.github/',
    'docs/',
    'examples/',
    'tests/',
]

[project.entry-points."reax.plugins.metrics"]
"atomic/num_species" = "tensorial.gcnn.atomic:NumSpecies"
"atomic/all_atomic_numbers" = "tensorial.gcnn.atomic:AllAtomicNumbers"
"atomic/avg_num_neighbours" = "tensorial.gcnn.atomic:AvgNumNeighbours"
"atomic/force_std" = "tensorial.gcnn.atomic:ForceStd"
"atomic/energy_per_atom_lstsq" = "tensorial.gcnn.atomic:EnergyPerAtomLstsq"

[tool.flynt]
line-length = 100
fail-on-change = true

[tool.isort]
profile = "black"
force_sort_within_sections = true
include_trailing_comma = true
line_length = 100
multi_line_output = 3

[tool.pylint.format]
max-line-length = 100

[tool.black]
line-length = 100

[tool.pylint.messages_control]
disable = [
    # Unfortunately jaxtyping decorator creates a function that seems to mistakenly be identified as
    # not returning anything, so we have to disable the error below for now
    'assignment-from-no-return',
    'duplicate-code',
    'import-outside-toplevel',
    'missing-docstring',
    'locally-disabled',
    'too-few-public-methods',
    'too-many-arguments',
    'too-many-instance-attributes',
    'use-dict-literal',
    'fixme',
]

[tool.pylint.design]
max-locals = 20
max-positional-arguments = 11

[pytest]
log_cli = "True"
log_cli_level = "DEBUG"

[tool.pytest_env]
XLA_PYTHON_CLIENT_PREALLOCATE = false # Don't allow JAX to preallocate memory

[tool.pytest.ini_options]
minversion = '6.0'
testpaths = [
    'test',
]
filterwarnings = [
    'ignore::DeprecationWarning:frozendict:',
]
#addopts = "--jaxtyping-packages=tensorial,beartype.beartype"

[tool.yapf]
align_closing_bracket_with_visual_indent = true
based_on_style = 'google'
coalesce_brackets = true
column_limit = 100
dedent_closing_brackets = true
indent_dictionary_value = false
split_arguments_when_comma_terminated = true

[tool.tox]
legacy_tox_ini = """
[tox]
envlist = py313

[testenv]
usedevelop = true

[testenv:py{310,311,312,313,314}]
description = Run the unit tests
extras =
    dev
commands = pytest {posargs}

[testenv:pre-commit]
description = Run the style checks and formatting
extras =
    dev
commands = pre-commit run {posargs}

[pytest]
filterwarnings =
    ignore::DeprecationWarning:distutils:
"""