├── pymbolic ├── py.typed ├── interop │ ├── __init__.py │ ├── matchpy │ │ ├── mapper.py │ │ └── tofrom.py │ ├── sympy.py │ └── symengine.py ├── maxima.py ├── sympy_interface.py ├── version.py ├── imperative │ ├── __init__.py │ ├── instruction.py │ ├── analysis.py │ ├── transform.py │ ├── utils.py │ └── statement.py ├── mapper │ ├── analysis.py │ ├── persistent_hash.py │ ├── constant_converter.py │ ├── cse_tagger.py │ ├── flop_counter.py │ ├── flattener.py │ ├── substitutor.py │ ├── coefficient.py │ ├── constant_folder.py │ ├── collector.py │ ├── distributor.py │ ├── dependency.py │ ├── graphviz.py │ └── c_code.py ├── functions.py ├── __init__.py ├── traits.py ├── geometric_algebra │ └── primitives.py ├── rational.py ├── compiler.py ├── cse.py └── typing.py ├── TODO ├── doc ├── algorithms.rst ├── upload-docs.sh ├── geometric-algebra.rst ├── primitives.rst ├── mappers.rst ├── utilities.rst ├── conf.py ├── index.rst ├── Makefile └── misc.rst ├── requirements.txt ├── LITERATURE ├── .gitignore ├── .test-py3.yml ├── .github ├── dependabot.yml └── workflows │ ├── autopush.yml │ └── ci.yml ├── MANIFEST.in ├── .editorconfig ├── CITATION.cff ├── LICENSE ├── test ├── simple.py ├── test_persistent_hash.py ├── test_pattern_match.py ├── test_matchpy.py ├── test_sympy.py └── test_maxima.py ├── .gitlab-ci.yml ├── README.rst └── pyproject.toml /pymbolic/py.typed: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pymbolic/interop/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /TODO: -------------------------------------------------------------------------------- 1 | Build out traits 2 | 3 | add traits types to expressions 4 | 5 | simplification 6 | -------------------------------------------------------------------------------- /doc/algorithms.rst: -------------------------------------------------------------------------------- 1 | Algorithms 2 | ========== 3 | 4 | .. automodule:: pymbolic.algorithm 5 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # test_matchpy_interop.py 2 | matchpy 3 | astunparse; python_version < '3.9' 4 | -------------------------------------------------------------------------------- /doc/upload-docs.sh: -------------------------------------------------------------------------------- 1 | #! /bin/sh 2 | 3 | rsync --verbose --archive --delete _build/html/ doc-upload:doc/pymbolic 4 | -------------------------------------------------------------------------------- /doc/geometric-algebra.rst: -------------------------------------------------------------------------------- 1 | Geometric Algebra 2 | ================= 3 | 4 | .. automodule:: pymbolic.geometric_algebra 5 | 6 | .. vim: sw=4 7 | -------------------------------------------------------------------------------- /LITERATURE: -------------------------------------------------------------------------------- 1 | [Bosch]: S. Bosch, Algebra, 3rd Edition, Springer, 1999 2 | [Davenport]: J.H. Davenport, Y. Siret, E. Tournier, Computer Algebra, 3 | Academic Press, 1988 4 | -------------------------------------------------------------------------------- /doc/primitives.rst: -------------------------------------------------------------------------------- 1 | Primitives (Basic Objects) 2 | ========================== 3 | 4 | .. automodule:: pymbolic.typing 5 | .. automodule:: pymbolic.primitives 6 | 7 | .. vim: sw=4 8 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | build 2 | .*.sw[po] 3 | .sw[po] 4 | *~ 5 | *.pyc 6 | *.pyo 7 | *.egg-info 8 | MANIFEST 9 | dist 10 | setuptools*egg 11 | setuptools*tar.gz 12 | setuptools.pth 13 | _build 14 | 15 | .cache 16 | -------------------------------------------------------------------------------- /.test-py3.yml: -------------------------------------------------------------------------------- 1 | name: py3 2 | channels: 3 | - conda-forge 4 | - defaults 5 | dependencies: 6 | - conda-forge::numpy 7 | - conda-forge::sympy 8 | - python 9 | - python-symengine 10 | # - pexpect 11 | # - maxima 12 | -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | updates: 3 | # Set update schedule for GitHub Actions 4 | - package-ecosystem: "github-actions" 5 | directory: "/" 6 | schedule: 7 | interval: "weekly" 8 | 9 | # vim: sw=4 10 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include LITERATURE 2 | include LICENSE 3 | 4 | include README.rst 5 | 6 | include doc/*.rst 7 | include doc/Makefile 8 | include doc/*.py 9 | include doc/conf.py 10 | include doc/_static/*.css 11 | include doc/_templates/*.html 12 | 13 | include test/*.py 14 | -------------------------------------------------------------------------------- /pymbolic/maxima.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from warnings import warn 4 | 5 | from pymbolic.interop.maxima import * # noqa: F403 6 | 7 | 8 | warn("pymbolic.maxima is deprecated. Use pymbolic.interop.maxima instead", 9 | DeprecationWarning, stacklevel=1) 10 | -------------------------------------------------------------------------------- /pymbolic/sympy_interface.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from warnings import warn 4 | 5 | from pymbolic.interop.sympy import * # noqa: F403 6 | 7 | 8 | warn("pymbolic.sympy_interface is deprecated. Use pymbolic.interop.sympy instead", 9 | DeprecationWarning, stacklevel=1) 10 | -------------------------------------------------------------------------------- /pymbolic/version.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from importlib import metadata 4 | 5 | 6 | def _parse_version(version: str) -> tuple[tuple[int, ...], str]: 7 | import re 8 | 9 | m = re.match(r"^([0-9.]+)([a-z0-9]*?)$", version) 10 | assert m is not None 11 | 12 | return tuple(int(nr) for nr in m.group(1).split(".")), m.group(2) 13 | 14 | 15 | VERSION_TEXT = metadata.version("pymbolic") 16 | VERSION, VERSION_STATUS = _parse_version(VERSION_TEXT) 17 | -------------------------------------------------------------------------------- /.github/workflows/autopush.yml: -------------------------------------------------------------------------------- 1 | name: Gitlab mirror 2 | on: 3 | push: 4 | branches: 5 | - main 6 | 7 | jobs: 8 | autopush: 9 | name: Automatic push to gitlab.tiker.net 10 | if: startsWith(github.repository, 'inducer/') 11 | runs-on: ubuntu-latest 12 | steps: 13 | - uses: actions/checkout@v6 14 | - run: | 15 | curl -L -O https://tiker.net/ci-support-v0 16 | . ./ci-support-v0 17 | mirror_github_to_gitlab 18 | 19 | env: 20 | GITLAB_AUTOPUSH_KEY: ${{ secrets.GITLAB_AUTOPUSH_KEY }} 21 | 22 | # vim: sw=4 23 | -------------------------------------------------------------------------------- /pymbolic/interop/matchpy/mapper.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from typing import TYPE_CHECKING, Any 4 | 5 | 6 | if TYPE_CHECKING: 7 | from collections.abc import Callable 8 | 9 | from pymbolic.interop.matchpy import PymbolicOp 10 | 11 | 12 | class Mapper: 13 | def __init__(self) -> None: 14 | self.cache: dict[PymbolicOp, Any] = {} 15 | 16 | def rec(self, expr: PymbolicOp) -> Any: 17 | if expr in self.cache: 18 | return self.cache[expr] 19 | 20 | method: Callable[[PymbolicOp], Any] = getattr(self, expr._mapper_method) 21 | 22 | return method(expr) 23 | 24 | __call__ = rec 25 | -------------------------------------------------------------------------------- /.editorconfig: -------------------------------------------------------------------------------- 1 | # https://editorconfig.org/ 2 | # https://github.com/editorconfig/editorconfig-vim 3 | # https://github.com/editorconfig/editorconfig-emacs 4 | 5 | root = true 6 | 7 | [*] 8 | indent_style = space 9 | end_of_line = lf 10 | charset = utf-8 11 | trim_trailing_whitespace = true 12 | insert_final_newline = true 13 | 14 | [*.py] 15 | indent_size = 4 16 | 17 | [*.rst] 18 | indent_size = 4 19 | 20 | [*.cpp] 21 | indent_size = 2 22 | 23 | [*.hpp] 24 | indent_size = 2 25 | 26 | # There may be one in doc/ 27 | [Makefile] 28 | indent_style = tab 29 | 30 | # https://github.com/microsoft/vscode/issues/1679 31 | [*.md] 32 | trim_trailing_whitespace = false 33 | -------------------------------------------------------------------------------- /CITATION.cff: -------------------------------------------------------------------------------- 1 | cff-version: 1.2.0 2 | message: "If you use this software, please cite it as below." 3 | authors: 4 | - family-names: "Kloeckner" 5 | given-names: "Andreas" 6 | orcid: "https://orcid.org/0000-0003-1228-519X" 7 | - family-names: "Wala" 8 | given-names: "Matt" 9 | - family-names: "Fernando" 10 | given-names: "Isuru" 11 | - family-names: "Kulkarni" 12 | given-names: "Kaushik" 13 | - family-names: "Fikl" 14 | given-names: "Alex" 15 | - family-names: "Weiner" 16 | given-names: "Zach" 17 | - family-names: "Kempf" 18 | given-names: "Dominic" 19 | - family-names: "Ham" 20 | given-names: "David A." 21 | - family-names: "Mitchell" 22 | given-names: "Lawrence" 23 | - family-names: "Wilcox" 24 | given-names: "Lucas C" 25 | - family-names: "Diener" 26 | given-names: "Matthias" 27 | - family-names: "Kapyshin" 28 | given-names: "Pavlo" 29 | - family-names: "Raksi" 30 | given-names: "Reno" 31 | - family-names: "Gibson" 32 | given-names: "Thomas H." 33 | title: "pymbolic" 34 | version: 2022.1 35 | doi: 10.5281/zenodo.6533945 36 | date-released: 2022-05-08 37 | url: "https://github.com/inducer/pymbolic" 38 | license: MIT 39 | -------------------------------------------------------------------------------- /doc/mappers.rst: -------------------------------------------------------------------------------- 1 | Mappers 2 | ======= 3 | 4 | .. automodule:: pymbolic.mapper 5 | 6 | More specialized mappers 7 | ------------------------ 8 | 9 | Converting to strings and code 10 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 11 | 12 | .. automodule:: pymbolic.mapper.stringifier 13 | 14 | .. automodule:: pymbolic.mapper.c_code 15 | 16 | .. automodule:: pymbolic.mapper.graphviz 17 | 18 | Some minimal mathematics 19 | ^^^^^^^^^^^^^^^^^^^^^^^^ 20 | 21 | .. automodule:: pymbolic.mapper.evaluator 22 | 23 | .. automodule:: pymbolic.mapper.differentiator 24 | 25 | .. automodule:: pymbolic.mapper.distributor 26 | 27 | .. automodule:: pymbolic.mapper.collector 28 | 29 | .. automodule:: pymbolic.mapper.constant_folder 30 | .. automodule:: pymbolic.mapper.substitutor 31 | 32 | Finding expression properties 33 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 34 | 35 | .. automodule:: pymbolic.mapper.dependency 36 | 37 | .. automodule:: pymbolic.mapper.flop_counter 38 | 39 | .. autoclass:: FlopCounter 40 | .. autoclass:: CSEAwareFlopCounter 41 | 42 | Analysis tools 43 | ^^^^^^^^^^^^^^ 44 | 45 | .. automodule:: pymbolic.mapper.analysis 46 | 47 | Simplification 48 | ^^^^^^^^^^^^^^ 49 | 50 | .. automodule:: pymbolic.mapper.flattener 51 | 52 | 53 | .. vim: sw=4 54 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | pymbolic is licensed to you under the MIT/X Consortium license: 2 | 3 | Copyright (c) 2009-16 Andreas Klöckner and Contributors. 4 | 5 | Permission is hereby granted, free of charge, to any person 6 | obtaining a copy of this software and associated documentation 7 | files (the "Software"), to deal in the Software without 8 | restriction, including without limitation the rights to use, 9 | copy, modify, merge, publish, distribute, sublicense, and/or sell 10 | copies of the Software, and to permit persons to whom the 11 | Software is furnished to do so, subject to the following 12 | conditions: 13 | 14 | The above copyright notice and this permission notice shall be 15 | included in all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 18 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES 19 | OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 20 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT 21 | HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 22 | WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 23 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 24 | OTHER DEALINGS IN THE SOFTWARE. 25 | -------------------------------------------------------------------------------- /pymbolic/imperative/__init__.py: -------------------------------------------------------------------------------- 1 | """Imperative program representation""" 2 | from __future__ import annotations 3 | 4 | 5 | __copyright__ = "Copyright (C) 2015 Matt Wala, Andreas Kloeckner" 6 | 7 | __license__ = """ 8 | Permission is hereby granted, free of charge, to any person obtaining a copy 9 | of this software and associated documentation files (the "Software"), to deal 10 | in the Software without restriction, including without limitation the rights 11 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 12 | copies of the Software, and to permit persons to whom the Software is 13 | furnished to do so, subject to the following conditions: 14 | 15 | The above copyright notice and this permission notice shall be included in 16 | all copies or substantial portions of the Software. 17 | 18 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 19 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 20 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 21 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 22 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 23 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 24 | THE SOFTWARE. 25 | """ 26 | -------------------------------------------------------------------------------- /doc/utilities.rst: -------------------------------------------------------------------------------- 1 | Utilities for dealing with expressions 2 | ====================================== 3 | 4 | Parser 5 | ------ 6 | 7 | .. currentmodule:: pymbolic 8 | 9 | .. autofunction:: parse 10 | 11 | The parser is also relatively easy to extend. See the source code of the following 12 | class. 13 | 14 | .. automodule:: pymbolic.parser 15 | 16 | Compiler 17 | -------- 18 | 19 | .. automodule:: pymbolic.compiler 20 | 21 | Interoperability with other symbolic systems 22 | ============================================ 23 | 24 | Interoperability with :mod:`sympy` 25 | ---------------------------------- 26 | 27 | .. automodule:: pymbolic.interop.sympy 28 | 29 | Interoperability with ``symengine`` 30 | ----------------------------------- 31 | 32 | .. automodule:: pymbolic.interop.symengine 33 | 34 | Interoperability with Maxima 35 | ---------------------------- 36 | 37 | .. automodule:: pymbolic.interop.maxima 38 | 39 | Interoperability with Python's :mod:`ast` module 40 | ------------------------------------------------ 41 | 42 | .. automodule:: pymbolic.interop.ast 43 | 44 | Interoperability with :mod:`matchpy.functions` module 45 | ----------------------------------------------------- 46 | 47 | .. automodule:: pymbolic.interop.matchpy 48 | 49 | Visualizing Expressions 50 | ======================= 51 | 52 | .. autofunction:: pymbolic.imperative.utils.get_dot_dependency_graph 53 | -------------------------------------------------------------------------------- /test/simple.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | from pymbolic import parse, var 27 | from pymbolic.mapper.dependency import DependencyMapper 28 | 29 | 30 | x = var("x") 31 | y = var("y") 32 | 33 | expr2 = 3*x+5-y 34 | expr = parse("3*x+5-y") 35 | 36 | print(expr) 37 | print(expr2) 38 | 39 | dm = DependencyMapper() 40 | print(dm(expr)) 41 | -------------------------------------------------------------------------------- /.gitlab-ci.yml: -------------------------------------------------------------------------------- 1 | Python 3: 2 | script: | 3 | PY_EXE=python3 4 | # pytest tries to import this, but it doesn't find symengine 5 | rm pymbolic/interop/symengine.py 6 | EXTRA_INSTALL="numpy sympy pexpect" 7 | curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project.sh 8 | . ./build-and-test-py-project.sh 9 | tags: 10 | - python3 11 | - maxima 12 | except: 13 | - tags 14 | artifacts: 15 | reports: 16 | junit: test/pytest.xml 17 | 18 | Python 3 Conda: 19 | script: | 20 | CONDA_ENVIRONMENT=.test-py3.yml 21 | curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project-within-miniconda.sh 22 | . ./build-and-test-py-project-within-miniconda.sh 23 | 24 | tags: 25 | - linux 26 | except: 27 | - tags 28 | artifacts: 29 | reports: 30 | junit: test/pytest.xml 31 | 32 | Documentation: 33 | script: 34 | - EXTRA_INSTALL="numpy sympy symengine" 35 | - curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-docs.sh 36 | - ". ./build-docs.sh" 37 | tags: 38 | - linux 39 | 40 | Ruff: 41 | script: | 42 | pipx install ruff 43 | ruff check 44 | tags: 45 | - docker-runner 46 | except: 47 | - tags 48 | 49 | Downstream: 50 | parallel: 51 | matrix: 52 | - DOWNSTREAM_PROJECT: [loopy, pytential, pytato] 53 | tags: 54 | - large-node 55 | - "docker-runner" 56 | script: | 57 | curl -L -O https://tiker.net/ci-support-v0 58 | . ./ci-support-v0 59 | test_downstream "$DOWNSTREAM_PROJECT" 60 | -------------------------------------------------------------------------------- /pymbolic/imperative/instruction.py: -------------------------------------------------------------------------------- 1 | """Instruction types""" 2 | from __future__ import annotations 3 | 4 | 5 | __copyright__ = "Copyright (C) 2015 Matt Wala, Andreas Kloeckner" 6 | 7 | __license__ = """ 8 | Permission is hereby granted, free of charge, to any person obtaining a copy 9 | of this software and associated documentation files (the "Software"), to deal 10 | in the Software without restriction, including without limitation the rights 11 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 12 | copies of the Software, and to permit persons to whom the Software is 13 | furnished to do so, subject to the following conditions: 14 | 15 | The above copyright notice and this permission notice shall be included in 16 | all copies or substantial portions of the Software. 17 | 18 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 19 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 20 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 21 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 22 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 23 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 24 | THE SOFTWARE. 25 | """ 26 | 27 | from warnings import warn 28 | 29 | 30 | warn("pymbolic.imperative.instruction was imported. This has been renamed " 31 | "to pymbolic.imperative.statement", DeprecationWarning, stacklevel=1) 32 | 33 | from pymbolic.imperative.statement import ( # noqa: F401 34 | Assignment, 35 | ConditionalAssignment, 36 | ConditionalInstruction, 37 | Instruction, 38 | Nop, 39 | ) 40 | -------------------------------------------------------------------------------- /test/test_persistent_hash.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2023 University of Illinois Board of Trustees" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | 27 | from pymbolic.mapper.persistent_hash import PersistentHashWalkMapper 28 | 29 | 30 | def test_persistent_hash_simple() -> None: 31 | import hashlib 32 | 33 | from testlib import generate_random_expression 34 | expr = generate_random_expression(seed=(333)) 35 | 36 | key_hash = hashlib.sha256() 37 | 38 | phwm = PersistentHashWalkMapper(key_hash) 39 | phwm(expr) 40 | 41 | assert key_hash.hexdigest() == \ 42 | "1a1cd91483015333f2a9b06ab049a8edabc72aafc1f9b6d7cd831a39068e50da" 43 | -------------------------------------------------------------------------------- /pymbolic/imperative/analysis.py: -------------------------------------------------------------------------------- 1 | """Fusion and other user-facing code transforms""" 2 | from __future__ import annotations 3 | 4 | 5 | __copyright__ = "Copyright (C) 2015 Matt Wala, Andreas Kloeckner" 6 | 7 | __license__ = """ 8 | Permission is hereby granted, free of charge, to any person obtaining a copy 9 | of this software and associated documentation files (the "Software"), to deal 10 | in the Software without restriction, including without limitation the rights 11 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 12 | copies of the Software, and to permit persons to whom the Software is 13 | furnished to do so, subject to the following conditions: 14 | 15 | The above copyright notice and this permission notice shall be included in 16 | all copies or substantial portions of the Software. 17 | 18 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 19 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 20 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 21 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 22 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 23 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 24 | THE SOFTWARE. 25 | """ 26 | 27 | 28 | from typing import TYPE_CHECKING 29 | 30 | 31 | if TYPE_CHECKING: 32 | from collections.abc import Iterable 33 | 34 | from pymbolic.imperative.statement import StatementLike 35 | 36 | 37 | def get_all_used_insn_ids(insn_stream: Iterable[StatementLike]): 38 | return frozenset(insn.id for insn in insn_stream) 39 | 40 | 41 | def get_all_used_identifiers(insn_stream: Iterable[StatementLike]): 42 | result: set[str] = set() 43 | for insn in insn_stream: 44 | result |= insn.get_read_variables() 45 | result |= insn.get_written_variables() 46 | 47 | return result 48 | -------------------------------------------------------------------------------- /test/test_pattern_match.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2022 University of Illinois Board of Trustees" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | import pymbolic.primitives as p 27 | 28 | 29 | def test_pattern_match(): 30 | from pymbolic import var 31 | 32 | x = var("x") 33 | xp1 = (x+1) 34 | u = xp1**5 + x 35 | 36 | match u: 37 | case p.Sum((p.Power(base, exp), other_term)): 38 | assert base is xp1 39 | assert exp == 5 40 | assert other_term is x 41 | 42 | case _: 43 | raise AssertionError() 44 | 45 | 46 | if __name__ == "__main__": 47 | import sys 48 | if len(sys.argv) > 1: 49 | exec(sys.argv[1]) 50 | else: 51 | from pytest import main 52 | main([__file__]) 53 | 54 | # vim: fdm=marker 55 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | Pymbolic: Easy Expression Trees and Term Rewriting 2 | ================================================== 3 | 4 | .. image:: https://gitlab.tiker.net/inducer/pymbolic/badges/main/pipeline.svg 5 | :alt: Gitlab Build Status 6 | :target: https://gitlab.tiker.net/inducer/pymbolic/commits/main 7 | .. image:: https://github.com/inducer/pymbolic/actions/workflows/ci.yml/badge.svg 8 | :alt: Github Build Status 9 | :target: https://github.com/inducer/pymbolic/actions/workflows/ci.yml 10 | .. image:: https://badge.fury.io/py/pymbolic.svg 11 | :alt: Python Package Index Release Page 12 | :target: https://pypi.org/project/pymbolic 13 | .. image:: https://zenodo.org/badge/2016193.svg 14 | :alt: Zenodo DOI for latest release 15 | :target: https://zenodo.org/badge/latestdoi/2016193 16 | 17 | Pymbolic is a small expression tree and symbolic manipulation library. Two 18 | things set it apart from other libraries of its kind: 19 | 20 | * Users can easily write their own symbolic operations, simply by deriving 21 | from the builtin visitor classes. 22 | * Users can easily add their own symbolic entities to do calculations 23 | with. 24 | 25 | Pymbolic currently understands regular arithmetic expressions, derivatives, 26 | sparse polynomials, fractions, term substitution, expansion. It automatically 27 | performs constant folding, and it can compile its expressions into Python 28 | bytecode for fast(er) execution. 29 | 30 | If you are looking for a full-blown Computer Algebra System, look at 31 | `sympy `__ or 32 | `PyGinac `__. If you are looking for a 33 | basic, small and extensible set of symbolic operations, pymbolic may 34 | well be for you. 35 | 36 | Resources: 37 | 38 | * `PyPI package `__ 39 | * `Documentation `__ 40 | * `Source code (GitHub) `__ 41 | -------------------------------------------------------------------------------- /pymbolic/mapper/analysis.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = """Copyright (C) 2022 University of Illinois Board of Trustees""" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | from typing import TYPE_CHECKING 27 | 28 | from typing_extensions import override 29 | 30 | from pymbolic.mapper import CachedWalkMapper 31 | 32 | 33 | if TYPE_CHECKING: 34 | from pymbolic.typing import Expression 35 | 36 | 37 | __doc__ = """ 38 | .. autoclass:: NodeCountMapper 39 | :show-inheritance: 40 | 41 | .. autofunction:: get_num_nodes 42 | """ 43 | 44 | 45 | # {{{ NodeCountMapper 46 | 47 | class NodeCountMapper(CachedWalkMapper[[]]): 48 | """ 49 | Counts the number of nodes in an expression tree. Nodes that occur 50 | repeatedly as well as :class:`~pymbolic.primitives.CommonSubexpression` 51 | nodes are only counted once. 52 | 53 | .. attribute:: count 54 | 55 | The number of nodes. 56 | """ 57 | 58 | count: int 59 | 60 | def __init__(self) -> None: 61 | super().__init__() 62 | self.count = 0 63 | 64 | @override 65 | def post_visit(self, expr: object) -> None: 66 | self.count += 1 67 | 68 | 69 | def get_num_nodes(expr: Expression) -> int: 70 | """ 71 | :returns: the number of nodes in *expr*. Nodes that occur 72 | repeatedly as well as :class:`~pymbolic.primitives.CommonSubexpression` 73 | nodes are only counted once. 74 | """ 75 | 76 | ncm = NodeCountMapper() 77 | ncm(expr) 78 | 79 | return ncm.count 80 | 81 | # }}} 82 | -------------------------------------------------------------------------------- /pymbolic/mapper/persistent_hash.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | 27 | from warnings import warn 28 | 29 | from pymbolic.mapper import WalkMapper 30 | 31 | 32 | class PersistentHashWalkMapper(WalkMapper): 33 | """A subclass of :class:`pymbolic.mapper.WalkMapper` for constructing 34 | persistent hash keys for use with 35 | :class:`pytools.persistent_dict.PersistentDict`. 36 | """ 37 | 38 | def __init__(self, key_hash): 39 | self.key_hash = key_hash 40 | 41 | warn("PersistentHashWalkMapper is deprecated. " 42 | "Since they are dataclasses, expression objects should now " 43 | "support persistent hashing natively without any help. " 44 | "It will be removed in 2026.", 45 | DeprecationWarning, stacklevel=2) 46 | 47 | def visit(self, expr): 48 | self.key_hash.update(type(expr).__name__.encode("utf8")) 49 | return True 50 | 51 | def map_variable(self, expr): 52 | self.key_hash.update(expr.name.encode("utf8")) 53 | 54 | def map_constant(self, expr): 55 | import sys 56 | if "numpy" in sys.modules: 57 | import numpy as np 58 | if isinstance(expr, np.generic): 59 | # Makes a Python scalar from a numpy one. 60 | expr = expr.item() 61 | 62 | self.key_hash.update(repr(expr).encode("utf8")) 63 | 64 | def map_comparison(self, expr): 65 | if self.visit(expr): 66 | self.rec(expr.left) 67 | self.key_hash.update(repr(expr.operator).encode("utf8")) 68 | self.rec(expr.right) 69 | -------------------------------------------------------------------------------- /pymbolic/functions.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | from typing import TYPE_CHECKING 27 | 28 | import pymbolic.primitives as p 29 | 30 | 31 | if TYPE_CHECKING: 32 | from pymbolic.typing import ArithmeticExpression 33 | 34 | 35 | def sin(x: ArithmeticExpression) -> ArithmeticExpression: 36 | return p.Call(p.Lookup(p.Variable("math"), "sin"), (x,)) 37 | 38 | 39 | def cos(x: ArithmeticExpression) -> ArithmeticExpression: 40 | return p.Call(p.Lookup(p.Variable("math"), "cos"), (x,)) 41 | 42 | 43 | def tan(x: ArithmeticExpression) -> ArithmeticExpression: 44 | return p.Call(p.Lookup(p.Variable("math"), "tan"), (x,)) 45 | 46 | 47 | def log(x: ArithmeticExpression) -> ArithmeticExpression: 48 | return p.Call(p.Lookup(p.Variable("math"), "log"), (x,)) 49 | 50 | 51 | def exp(x: ArithmeticExpression) -> ArithmeticExpression: 52 | return p.Call(p.Lookup(p.Variable("math"), "exp"), (x,)) 53 | 54 | 55 | def sinh(x: ArithmeticExpression) -> ArithmeticExpression: 56 | return p.Call(p.Lookup(p.Variable("math"), "sinh"), (x,)) 57 | 58 | 59 | def cosh(x: ArithmeticExpression) -> ArithmeticExpression: 60 | return p.Call(p.Lookup(p.Variable("math"), "cosh"), (x,)) 61 | 62 | 63 | def tanh(x: ArithmeticExpression) -> ArithmeticExpression: 64 | return p.Call(p.Lookup(p.Variable("math"), "tanh"), (x,)) 65 | 66 | 67 | def expm1(x: ArithmeticExpression) -> ArithmeticExpression: 68 | return p.Call(p.Lookup(p.Variable("math"), "expm1"), (x,)) 69 | 70 | 71 | def fabs(x: ArithmeticExpression) -> ArithmeticExpression: 72 | return p.Call(p.Lookup(p.Variable("math"), "fabs"), (x,)) 73 | 74 | 75 | def sign(x: ArithmeticExpression) -> ArithmeticExpression: 76 | return p.Call(p.Lookup(p.Variable("math"), "copysign"), (1, x,)) 77 | -------------------------------------------------------------------------------- /pymbolic/mapper/constant_converter.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2017 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | from pymbolic.mapper import IdentityMapper 27 | 28 | 29 | class ConstantToNumpyConversionMapper(IdentityMapper[[]]): 30 | """Because of `this numpy bug `__, 31 | sized :mod:`numpy` number (i.e. ones with definite bit width, such as 32 | :class:`numpy.complex64`) have a low likelihood of surviving expression 33 | construction. 34 | 35 | This mapper ensures that all occurring numerical constants are of the 36 | expected type. 37 | """ 38 | 39 | def __init__(self, real_type, complex_type=None, integer_type=None): 40 | import numpy as np 41 | self.real_type = real_type 42 | 43 | if complex_type is None: 44 | if real_type is np.float32: 45 | complex_type = np.complex64 46 | elif real_type is np.float64: 47 | complex_type = np.complex128 48 | elif real_type is np.float128: 49 | complex_type = np.complex256 50 | else: 51 | raise TypeError("unable to determine corresponding complex type " 52 | f"for '{real_type.__name__}'") 53 | 54 | self.complex_type = complex_type 55 | 56 | self.integer_type = integer_type 57 | 58 | def map_constant(self, expr): 59 | if expr.imag: 60 | return self.complex_type(expr) 61 | 62 | expr = expr.real 63 | 64 | if int(expr) == expr and not isinstance(expr, float): 65 | if self.integer_type is not None: 66 | return self.integer_type(expr) 67 | else: 68 | return expr 69 | else: 70 | return self.real_type(expr) 71 | -------------------------------------------------------------------------------- /doc/conf.py: -------------------------------------------------------------------------------- 1 | from importlib import metadata 2 | from urllib.request import urlopen 3 | 4 | 5 | _conf_url = "https://raw.githubusercontent.com/inducer/sphinxconfig/main/sphinxconfig.py" 6 | with urlopen(_conf_url) as _inf: 7 | exec(compile(_inf.read(), _conf_url, "exec"), globals()) 8 | 9 | copyright = "2013-24, Andreas Kloeckner" 10 | release = metadata.version("pymbolic") 11 | version = ".".join(release.split(".")[:2]) 12 | 13 | # List of patterns, relative to source directory, that match files and 14 | # directories to ignore when looking for source files. 15 | exclude_patterns = ["_build"] 16 | 17 | intersphinx_mapping = { 18 | "galgebra": ("https://galgebra.readthedocs.io/en/latest/", None), 19 | "mako": ("https://docs.makotemplates.org/en/latest/", None), 20 | "matchpy": ("https://matchpy.readthedocs.io/en/latest/", None), 21 | "numpy": ("https://numpy.org/doc/stable/", None), 22 | "python": ("https://docs.python.org/3", None), 23 | "sympy": ("https://docs.sympy.org/dev/", None), 24 | "pytools": ("https://documen.tician.de/pytools/", None), 25 | "typing_extensions": 26 | ("https://typing-extensions.readthedocs.io/en/latest/", None), 27 | "constantdict": 28 | ("https://matthiasdiener.github.io/constantdict/", None) 29 | } 30 | 31 | autodoc_type_aliases = { 32 | "Expression": "Expression", 33 | "ArithmeticExpression": "ArithmeticExpression", 34 | } 35 | 36 | nitpick_ignore_regex = [ 37 | # Sphinx started complaining about these in 8.2.1(-ish) 38 | # -AK, 2025-02-24 39 | ["py:class", r"TypeAliasForwardRef"], 40 | ["py:class", r"ast.expr"], 41 | ["py:class", r"onp.*"], 42 | ["py:class", r"optype.*"], 43 | ["py:class", r"MultiVector\[ArithmeticExpression\]"], 44 | ] 45 | 46 | sphinxconfig_missing_reference_aliases = { 47 | # numpy 48 | "NDArray": "obj:numpy.typing.NDArray", 49 | "DTypeLike": "obj:numpy.typing.DTypeLike", 50 | "np.inexact": "class:numpy.inexact", 51 | "np.generic": "class:numpy.generic", 52 | "np.dtype": "class:numpy.dtype", 53 | "np.ndarray": "class:numpy.ndarray", 54 | # matchpy typing 55 | "ReplacementRule": "class:matchpy.functions.ReplacementRule", 56 | # pytools typing 57 | "T": "class:pytools.T", 58 | "ShapeT": "class:pytools.obj_array.ShapeT", 59 | "ObjectArray": "class:pytools.obj_array.ObjectArray", 60 | "ObjectArray1D": "class:pytools.obj_array.ObjectArray", 61 | # pymbolic typing 62 | "ArithmeticExpression": "data:pymbolic.typing.ArithmeticExpression", 63 | "Comparison": "class:pymbolic.primitives.Comparison", 64 | "Expression": "data:pymbolic.typing.Expression", 65 | "ExpressionNode": "class:pymbolic.primitives.ExpressionNode", 66 | "FromMatchpyT": "obj:pymbolic.interop.matchpy.FromMatchpyT", 67 | "LogicalAnd": "class:pymbolic.primitives.LogicalAnd", 68 | "LogicalNot": "class:pymbolic.primitives.LogicalNot", 69 | "LogicalOr": "class:pymbolic.primitives.LogicalOr", 70 | "Lookup": "class:pymbolic.primitives.Lookup", 71 | "ToMatchpyT": "obj:pymbolic.interop.matchpy.ToMatchpyT", 72 | "_Expression": "data:pymbolic.typing.Expression", 73 | "p.AlgebraicLeaf": "class:pymbolic.primitives.AlgebraicLeaf", 74 | "prim.Variable": "class:pymbolic.primitives.Variable", 75 | "P.args": "obj:pymbolic.mapper.P", 76 | "P.kwargs": "obj:pymbolic.mapper.P", 77 | } 78 | 79 | 80 | def setup(app): 81 | app.connect("missing-reference", process_autodoc_missing_reference) # noqa: F821 82 | -------------------------------------------------------------------------------- /doc/index.rst: -------------------------------------------------------------------------------- 1 | Welcome to pymbolic! 2 | ==================== 3 | 4 | Pymbolic is a simple and extensible package for precise manipulation of 5 | symbolic expressions in Python. It doesn't try to compete with :mod:`sympy` as 6 | a computer algebra system. Pymbolic emphasizes providing an extensible 7 | expression tree and a flexible, extensible way to manipulate it. 8 | 9 | A taste of :mod:`pymbolic` 10 | -------------------------- 11 | 12 | Follow along on a simple example. Let's import :mod:`pymbolic` and create a 13 | symbol, *x* in this case. 14 | 15 | .. doctest:: 16 | 17 | >>> import pymbolic as pmbl 18 | 19 | >>> x = pmbl.var("x") 20 | >>> x 21 | Variable('x') 22 | 23 | Next, let's create an expression using *x*: 24 | 25 | .. doctest:: 26 | 27 | >>> u = (x+1)**5 28 | >>> u 29 | Power(Sum((Variable('x'), 1)), 5) 30 | >>> print(u) 31 | (x + 1)**5 32 | 33 | Note the two ways an expression can be printed, namely :func:`repr` and 34 | :class:`str`. :mod:`pymbolic` purposefully distinguishes the two. 35 | 36 | :mod:`pymbolic` does not perform any manipulations on expressions 37 | you put in. It has a few of those built in, but that's not really the point: 38 | 39 | .. doctest:: 40 | 41 | >>> print(pmbl.differentiate(u, 'x')) 42 | 5*(x + 1)**4 43 | 44 | .. _custom-manipulation: 45 | 46 | Manipulating expressions 47 | ^^^^^^^^^^^^^^^^^^^^^^^^ 48 | 49 | The point is for you to be able to easily write so-called *mappers* to 50 | manipulate expressions. Suppose we would like all sums replaced by 51 | products: 52 | 53 | .. doctest:: 54 | 55 | >>> from pymbolic.mapper import IdentityMapper 56 | >>> class MyMapper(IdentityMapper): 57 | ... def map_sum(self, expr): 58 | ... return pmbl.primitives.Product(expr.children) 59 | ... 60 | >>> print(u) 61 | (x + 1)**5 62 | >>> print(MyMapper()(u)) 63 | (x*1)**5 64 | 65 | Custom Objects 66 | ^^^^^^^^^^^^^^ 67 | 68 | You can also easily define your own objects to use inside an expression: 69 | 70 | .. doctest:: 71 | 72 | >>> from pymbolic import ExpressionNode, expr_dataclass 73 | >>> from pymbolic.typing import Expression 74 | >>> 75 | >>> @expr_dataclass() 76 | ... class FancyOperator(ExpressionNode): 77 | ... operand: Expression 78 | ... 79 | >>> u 80 | Power(Sum((Variable('x'), 1)), 5) 81 | >>> 17*FancyOperator(u) 82 | Product((17, FancyOperator(Power(Sum((..., 1)), 5)))) 83 | 84 | As a final example, we can now derive from *MyMapper* to multiply all 85 | *FancyOperator* instances by 2. 86 | 87 | .. doctest:: 88 | 89 | >>> FancyOperator.mapper_method 90 | 'map_fancy_operator' 91 | >>> class MyMapper2(MyMapper): 92 | ... def map_fancy_operator(self, expr): 93 | ... return 2*FancyOperator(self.rec(expr.operand)) 94 | ... 95 | >>> MyMapper2()(FancyOperator(u)) 96 | Product((2, FancyOperator(Power(Product((..., 1)), 5)))) 97 | 98 | .. automodule:: pymbolic 99 | 100 | Pymbolic around the web 101 | ----------------------- 102 | 103 | * `PyPI package `__ 104 | * `Documentation `__ 105 | * `Source code (GitHub) `__ 106 | 107 | Contents 108 | -------- 109 | 110 | .. toctree:: 111 | :maxdepth: 2 112 | 113 | primitives 114 | mappers 115 | utilities 116 | algorithms 117 | geometric-algebra 118 | misc 119 | 🚀 Github 120 | 💾 Download Releases 121 | 122 | * :ref:`genindex` 123 | * :ref:`modindex` 124 | 125 | .. vim: sw=4 126 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | on: 3 | push: 4 | branches: 5 | - main 6 | pull_request: 7 | paths-ignore: 8 | - 'doc/*.rst' 9 | schedule: 10 | - cron: '17 3 * * 0' 11 | 12 | concurrency: 13 | group: ${{ github.head_ref || github.ref_name }} 14 | cancel-in-progress: true 15 | 16 | jobs: 17 | typos: 18 | name: Typos 19 | runs-on: ubuntu-latest 20 | steps: 21 | - uses: actions/checkout@v6 22 | - uses: crate-ci/typos@master 23 | 24 | ruff: 25 | name: Ruff 26 | runs-on: ubuntu-latest 27 | steps: 28 | - uses: actions/checkout@v6 29 | with: 30 | submodules: true 31 | - uses: actions/setup-python@v6 32 | - name: "Main Script" 33 | run: | 34 | pip install ruff 35 | ruff check 36 | 37 | basedpyright: 38 | runs-on: ubuntu-latest 39 | 40 | steps: 41 | - uses: actions/checkout@v6 42 | - uses: actions/setup-python@v6 43 | with: 44 | python-version: '3.x' 45 | - name: "Main Script" 46 | run: | 47 | curl -L -O https://tiker.net/ci-support-v0 48 | . ./ci-support-v0 49 | build_py_project_in_venv 50 | pip install -e .[test] 51 | python -m pip install numpy pexpect sympy scipy scipy-stubs optype 52 | python -m pip install basedpyright 53 | basedpyright 54 | 55 | pytest: 56 | name: Pytest on Py${{ matrix.python-version }} 57 | runs-on: ubuntu-latest 58 | strategy: 59 | matrix: 60 | python-version: ["3.10", "3.12", "3.x"] 61 | steps: 62 | - uses: actions/checkout@v6 63 | - 64 | uses: actions/setup-python@v6 65 | with: 66 | python-version: ${{ matrix.python-version }} 67 | - name: "Main Script" 68 | run: | 69 | EXTRA_INSTALL="numpy sympy pexpect" 70 | curl -L -O https://tiker.net/ci-support-v0 71 | . ./ci-support-v0 72 | 73 | build_py_project_in_venv 74 | 75 | # https://github.com/inducer/pymbolic/pull/66#issuecomment-950371315 76 | pip install symengine || true 77 | 78 | test_py_project 79 | 80 | docs: 81 | name: Documentation 82 | runs-on: ubuntu-latest 83 | steps: 84 | - uses: actions/checkout@v6 85 | - 86 | uses: actions/setup-python@v6 87 | with: 88 | python-version: '3.x' 89 | - name: "Main Script" 90 | run: | 91 | EXTRA_INSTALL="numpy sympy symengine" 92 | curl -L -O https://tiker.net/ci-support-v0 93 | . ./ci-support-v0 94 | build_py_project_in_venv 95 | build_docs 96 | 97 | downstream_tests: 98 | strategy: 99 | matrix: 100 | downstream_project: [loopy, pytential, pytato] 101 | fail-fast: false 102 | name: Tests for downstream project ${{ matrix.downstream_project }} 103 | runs-on: ubuntu-latest 104 | steps: 105 | - uses: actions/checkout@v6 106 | - name: "Main Script" 107 | env: 108 | DOWNSTREAM_PROJECT: ${{ matrix.downstream_project }} 109 | run: | 110 | curl -L -O https://tiker.net/ci-support-v0 111 | . ./ci-support-v0 112 | test_downstream "$DOWNSTREAM_PROJECT" 113 | 114 | # vim: sw=4 115 | -------------------------------------------------------------------------------- /pymbolic/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | 27 | from functools import partial 28 | 29 | from pytools import module_getattr_for_deprecations 30 | 31 | from . import compiler, parser, primitives 32 | from .compiler import compile 33 | from .mapper import ( 34 | dependency, 35 | differentiator, 36 | distributor, 37 | evaluator, 38 | flattener, 39 | stringifier, 40 | substitutor, 41 | ) 42 | from .mapper.differentiator import differentiate, differentiate as diff 43 | from .mapper.distributor import distribute, distribute as expand 44 | from .mapper.evaluator import evaluate, evaluate_kw 45 | from .mapper.flattener import flatten 46 | from .mapper.substitutor import substitute 47 | from .parser import parse 48 | from .primitives import ( # noqa: N813 49 | ExpressionNode, 50 | Variable, 51 | Variable as var, 52 | disable_subscript_by_getitem, 53 | expr_dataclass, 54 | flattened_product, 55 | flattened_sum, 56 | linear_combination, 57 | make_common_subexpression as cse, 58 | make_sym_vector, 59 | quotient, 60 | subscript, 61 | variables, 62 | ) 63 | from .typing import ( 64 | ArithmeticExpression, 65 | Bool, 66 | Expression, 67 | Expression as _TypingExpression, 68 | Number, 69 | RealNumber, 70 | Scalar, 71 | ) 72 | from pymbolic.version import VERSION_TEXT as __version__ # noqa: F401,N811 73 | 74 | 75 | __all__ = ( 76 | "ArithmeticExpression", 77 | "Bool", 78 | "Expression", 79 | "ExpressionNode", 80 | "Number", 81 | "RealNumber", 82 | "Scalar", 83 | "Variable", 84 | "compile", 85 | "compiler", 86 | "cse", 87 | "dependency", 88 | "diff", 89 | "differentiate", 90 | "differentiator", 91 | "disable_subscript_by_getitem", 92 | "distribute", 93 | "distributor", 94 | "evaluate", 95 | "evaluate_kw", 96 | "evaluator", 97 | "expand", 98 | "expr_dataclass", 99 | "flatten", 100 | "flattened_product", 101 | "flattened_sum", 102 | "flattener", 103 | "linear_combination", 104 | "make_sym_vector", 105 | "parse", 106 | "parser", 107 | "primitives", 108 | "quotient", 109 | "stringifier", 110 | "subscript", 111 | "substitute", 112 | "substitutor", 113 | "var", 114 | "variables", 115 | ) 116 | 117 | __getattr__ = partial(module_getattr_for_deprecations, __name__, { 118 | "ExpressionT": ("pymbolic.typing.Expression", _TypingExpression, 2026), 119 | "ArithmeticExpressionT": ("ArithmeticExpression", ArithmeticExpression, 2026), 120 | "BoolT": ("Bool", Bool, 2026), 121 | "ScalarT": ("Scalar", Scalar, 2026), 122 | }) 123 | -------------------------------------------------------------------------------- /pymbolic/mapper/cse_tagger.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | from typing import TYPE_CHECKING 27 | 28 | from typing_extensions import Self, override 29 | 30 | import pymbolic.primitives as prim 31 | from pymbolic.mapper import IdentityMapper, WalkMapper 32 | 33 | 34 | if TYPE_CHECKING: 35 | from collections.abc import Callable, Hashable 36 | 37 | from pymbolic.typing import Expression 38 | 39 | 40 | class CSEWalkMapper(WalkMapper[[]]): 41 | subexpr_histogram: dict[Hashable, int] 42 | 43 | def __init__(self) -> None: 44 | self.subexpr_histogram = {} 45 | 46 | @override 47 | def visit(self, expr: object) -> bool: 48 | self.subexpr_histogram[expr] = self.subexpr_histogram.get(expr, 0) + 1 49 | return True 50 | 51 | 52 | class CSETagMapper(IdentityMapper[[]]): 53 | subexpr_histogram: dict[Hashable, int] 54 | 55 | def __init__(self, walk_mapper: CSEWalkMapper) -> None: 56 | self.subexpr_histogram = walk_mapper.subexpr_histogram 57 | 58 | def _map_subexpr(self, expr: prim.ExpressionNode, /) -> Expression: 59 | if self.subexpr_histogram.get(expr, 0) > 1: 60 | return prim.CommonSubexpression(expr, scope=prim.cse_scope.EVALUATION) 61 | else: 62 | return getattr(IdentityMapper, expr.mapper_method)(self, expr) 63 | 64 | map_call: Callable[[Self, prim.Call], Expression] = _map_subexpr 65 | map_sum: Callable[[Self, prim.Sum], Expression] = _map_subexpr 66 | map_product: Callable[[Self, prim.Product], Expression] = _map_subexpr 67 | map_quotient: Callable[[Self, prim.Quotient], Expression] = _map_subexpr 68 | map_floor_div: Callable[[Self, prim.FloorDiv], Expression] = _map_subexpr 69 | map_remainder: Callable[[Self, prim.Remainder], Expression] = _map_subexpr 70 | map_power: Callable[[Self, prim.Power], Expression] = _map_subexpr 71 | 72 | map_left_shift: Callable[[Self, prim.LeftShift], Expression] = _map_subexpr 73 | map_right_shift: Callable[[Self, prim.RightShift], Expression] = _map_subexpr 74 | 75 | map_bitwise_not: Callable[[Self, prim.BitwiseNot], Expression] = _map_subexpr 76 | map_bitwise_or: Callable[[Self, prim.BitwiseOr], Expression] = _map_subexpr 77 | map_bitwise_xor: Callable[[Self, prim.BitwiseXor], Expression] = _map_subexpr 78 | map_bitwise_and: Callable[[Self, prim.BitwiseAnd], Expression] = _map_subexpr 79 | 80 | map_comparison: Callable[[Self, prim.Comparison], Expression] = _map_subexpr 81 | 82 | map_logical_not: Callable[[Self, prim.LogicalNot], Expression] = _map_subexpr 83 | map_logical_and: Callable[[Self, prim.LogicalAnd], Expression] = _map_subexpr 84 | map_logical_or: Callable[[Self, prim.LogicalOr], Expression] = _map_subexpr 85 | 86 | map_if: Callable[[Self, prim.If], Expression] = _map_subexpr 87 | map_if_positive: Callable[[Self, prim.If], Expression] = _map_subexpr 88 | -------------------------------------------------------------------------------- /pymbolic/traits.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | from functools import reduce 27 | 28 | from . import algorithm 29 | 30 | 31 | class NoTraitsError(Exception): 32 | pass 33 | 34 | 35 | class NoCommonTraitsError(Exception): 36 | pass 37 | 38 | 39 | def traits(x): 40 | try: 41 | return x.traits() 42 | except AttributeError: 43 | if isinstance(x, complex | float): 44 | return FieldTraits() 45 | elif isinstance(x, int): 46 | return IntegerTraits() 47 | else: 48 | raise NoTraitsError from None 49 | 50 | 51 | def common_traits(*args): 52 | def common_traits_two(t_x, t_y): 53 | if isinstance(t_y, t_x.__class__): 54 | return t_y 55 | elif isinstance(t_x, t_y.__class__): 56 | return t_x 57 | else: 58 | raise NoCommonTraitsError( 59 | "No common traits type between '{}' and '{}'".format( 60 | t_x.__class__.__name__, 61 | t_y.__class__.__name__)) 62 | 63 | return reduce(common_traits_two, (traits(arg) for arg in args)) 64 | 65 | 66 | class Traits: 67 | pass 68 | 69 | 70 | class IntegralDomainTraits(Traits): 71 | pass 72 | 73 | 74 | class EuclideanRingTraits(IntegralDomainTraits): 75 | @classmethod 76 | def norm(cls, x): 77 | """Returns the algebraic norm of the element x. 78 | 79 | "Norm" is used as in the definition of a Euclidean ring, 80 | see [Bosch], p. 42 81 | """ 82 | raise NotImplementedError 83 | 84 | @staticmethod 85 | def gcd_extended(q, r): 86 | """Return a tuple (p, a, b) such that p = aq + br, 87 | where p is the greatest common divisor. 88 | """ 89 | return algorithm.extended_euclidean(q, r) 90 | 91 | @staticmethod 92 | def gcd(q, r): 93 | """Returns the greatest common divisor of q and r. 94 | """ 95 | return algorithm.extended_euclidean(q, r)[0] 96 | 97 | @classmethod 98 | def lcm(cls, a, b): 99 | """Returns the least common multiple of a and b. 100 | """ 101 | return a * b / cls.gcd(a, b) 102 | 103 | @staticmethod 104 | def get_unit(x): 105 | """Returns the unit in the prime factor decomposition of x. 106 | """ 107 | raise NotImplementedError 108 | 109 | 110 | class FieldTraits(IntegralDomainTraits): 111 | pass 112 | 113 | 114 | class IntegerTraits(EuclideanRingTraits): 115 | @staticmethod 116 | def norm(x): 117 | return abs(x) 118 | 119 | @staticmethod 120 | def get_unit(x): 121 | if x < 0: 122 | return -1 123 | elif x > 0: 124 | return 1 125 | else: 126 | raise RuntimeError("0 does not have a prime factor decomposition") 127 | -------------------------------------------------------------------------------- /test/test_matchpy.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2022 Kaushik Kulkarni" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | from typing import TYPE_CHECKING, cast 26 | 27 | import pymbolic.interop.matchpy as m 28 | import pymbolic.primitives as p 29 | 30 | 31 | if TYPE_CHECKING: 32 | from pymbolic.typing import Expression 33 | 34 | 35 | def test_replace_with_variadic_op(): 36 | # Replace 'a * c' -> '6' 37 | 38 | from pytools import product 39 | 40 | w1 = p.StarWildcard("w1_star") 41 | a, b, c = p.variables("a b c") 42 | 43 | def a_times_c_is_6(w1_star): 44 | result_args = [6] 45 | for k, v in w1_star.items(): 46 | result_args.extend([k]*v) 47 | 48 | return cast("Expression", product(result_args)) 49 | 50 | a_times_c_pattern = a * c * w1 51 | expr = a * b * c 52 | 53 | rule = m.make_replacement_rule(a_times_c_pattern, a_times_c_is_6) 54 | replaced_expr = m.replace_all(expr, [rule]) 55 | assert replaced_expr == (6 * b) 56 | 57 | 58 | def test_replace_with_non_commutative_op(): 59 | # Replace 'f(a, b, c, ...)' -> 'g(d, ...)' 60 | w1 = p.StarWildcard("w1_star") 61 | a, b, c, d, f, g, x, y, z = p.variables("a b c d f g x y z") 62 | 63 | def replacer(w1_star): 64 | return g(d, *w1_star) 65 | 66 | rule = m.make_replacement_rule(f(a, b, c, w1), replacer) 67 | replaced_expr = m.replace_all(f(a, b, c, x, y, z), [rule]) 68 | assert replaced_expr == g(d, x, y, z) 69 | 70 | 71 | def test_replace_with_ternary_ops(): 72 | from pymbolic import parse 73 | 74 | expr = parse("b < f(c, d)") 75 | rule = m.make_replacement_rule( 76 | p.Variable("f")(p.Variable("c"), p.DotWildcard("w1_")), 77 | lambda w1_: (w1_*42)) 78 | assert m.replace_all(expr, [rule]) == parse("b < (42 * d)") 79 | 80 | expr = parse("b if f(c, d) else g(e)") 81 | rule = m.make_replacement_rule( 82 | p.If(parse("f(c, d)"), 83 | p.DotWildcard("w1_"), 84 | p.DotWildcard("w2_")), 85 | lambda w1_, w2_: p.If(parse("f(d, c)"), w2_, w1_)) 86 | assert m.replace_all(expr, [rule]) == parse("g(e) if f(d, c) else b") 87 | 88 | 89 | def test_make_subexpr_subst(): 90 | from functools import reduce 91 | 92 | from pymbolic import parse 93 | from pymbolic.mapper.flattener import flatten 94 | 95 | subject = parse("a[k]*b[i, j]*c[i, j]*d[k]") 96 | pattern = parse("b[i, j]*c[i, j]") * p.StarWildcard("w1_") 97 | 98 | rule = m.make_replacement_rule( 99 | pattern, 100 | lambda w1_: (parse("subst(i, j)") 101 | * (reduce(lambda acc, x: acc * (x[0] ** x[1]), 102 | w1_.items(), 103 | 1))) 104 | ) 105 | 106 | replaced_expr = m.replace_all(subject, [rule]) 107 | 108 | ref_expr = flatten(parse("subst(i, j)*a[(k,)]*d[(k,)]")) 109 | assert flatten(replaced_expr) == ref_expr 110 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["hatchling"] 3 | build-backend = "hatchling.build" 4 | 5 | [project] 6 | name = "pymbolic" 7 | version = "2025.1" 8 | description = "A package for symbolic computation" 9 | readme = "README.rst" 10 | license = "MIT" 11 | authors = [ 12 | { name = "Andreas Kloeckner", email = "inform@tiker.net" }, 13 | ] 14 | requires-python = ">=3.10" 15 | classifiers = [ 16 | "Development Status :: 4 - Beta", 17 | "Intended Audience :: Developers", 18 | "Intended Audience :: Other Audience", 19 | "Intended Audience :: Science/Research", 20 | "Natural Language :: English", 21 | "Programming Language :: Python", 22 | "Programming Language :: Python :: 3 :: Only", 23 | "Topic :: Scientific/Engineering", 24 | "Topic :: Scientific/Engineering :: Mathematics", 25 | "Topic :: Software Development :: Libraries", 26 | "Topic :: Utilities", 27 | ] 28 | dependencies = [ 29 | "constantdict", 30 | "pytools>=2025.2", 31 | # for TypeIs 32 | "typing-extensions>=4.10", 33 | ] 34 | 35 | [project.optional-dependencies] 36 | matchpy = [ 37 | "matchpy", 38 | ] 39 | numpy = [ 40 | "numpy>=1.6", 41 | ] 42 | test = [ 43 | "basedpyright", 44 | "optype", 45 | "pytest", 46 | "ruff", 47 | ] 48 | 49 | [project.urls] 50 | Documentation = "https://documen.tician.de/pymbolic" 51 | Homepage = "https://github.com/inducer/pymbolic" 52 | 53 | [tool.hatch.build.targets.sdist] 54 | exclude = [ 55 | "/.git*", 56 | "/doc/_build", 57 | "/.editorconfig", 58 | "/run-*.sh", 59 | "/.basedpyright", 60 | ] 61 | 62 | 63 | [tool.ruff] 64 | preview = true 65 | 66 | [tool.ruff.lint] 67 | extend-select = [ 68 | "B", # flake8-bugbear 69 | "C", # flake8-comprehensions 70 | "E", # pycodestyle 71 | "F", # pyflakes 72 | "G", # flake8-logging-format 73 | "I", # flake8-isort 74 | "N", # pep8-naming 75 | "NPY", # numpy 76 | "PGH", # pygrep-hooks 77 | "Q", # flake8-quotes 78 | "RUF", # ruff 79 | "SIM", # flake8-simplify 80 | "TC", # flake8-type-checking 81 | "UP", # pyupgrade 82 | "W", # pycodestyle 83 | ] 84 | extend-ignore = [ 85 | "C409", # remove comprehension within tuple call 86 | "C90", # McCabe complexity 87 | "E226", # missing whitespace around arithmetic operator 88 | "E241", # multiple spaces after comma 89 | "E242", # tab after comma 90 | "E402", # module level import not at the top of file 91 | "UP031", # use f-strings instead of % 92 | "UP032", # use f-strings instead of .format 93 | ] 94 | 95 | [tool.ruff.lint.per-file-ignores] 96 | "experiments/traversal-benchmark.py" = ["E501"] 97 | "doc/conf.py" = ["I002"] 98 | "experiments/*.py" = ["I002"] 99 | 100 | [tool.ruff.lint.pep8-naming] 101 | extend-ignore-names = ["map_*"] 102 | 103 | [tool.ruff.lint.flake8-quotes] 104 | inline-quotes = "double" 105 | docstring-quotes = "double" 106 | multiline-quotes = "double" 107 | 108 | [tool.typos.default] 109 | extend-ignore-re = [ 110 | "(?Rm)^.*(#|//)\\s*spellchecker:\\s*disable-line" 111 | ] 112 | 113 | [tool.typos.default.extend-words] 114 | "nd" = "nd" 115 | 116 | [tool.ruff.lint.isort] 117 | known-first-party = ["pytools"] 118 | known-local-folder = ["pymbolic"] 119 | lines-after-imports = 2 120 | combine-as-imports = true 121 | required-imports = ["from __future__ import annotations"] 122 | 123 | [tool.basedpyright] 124 | reportImplicitStringConcatenation = "none" 125 | reportUnnecessaryIsInstance = "none" 126 | reportUnusedCallResult = "none" 127 | reportExplicitAny = "none" 128 | reportUnusedParameter = "hint" 129 | 130 | # This reports even cycles that are qualified by 'if TYPE_CHECKING'. Not what 131 | # we care about at this moment. 132 | # https://github.com/microsoft/pyright/issues/746 133 | reportImportCycles = "none" 134 | 135 | pythonVersion = "3.10" 136 | pythonPlatform = "All" 137 | 138 | [[tool.basedpyright.executionEnvironments]] 139 | root = "test" 140 | reportUnknownArgumentType = "hint" 141 | 142 | [[tool.basedpyright.executionEnvironments]] 143 | root = "pymbolic/interop/symengine.py" 144 | reportMissingTypeStubs = "none" 145 | reportMissingImports = "none" 146 | -------------------------------------------------------------------------------- /pymbolic/mapper/flop_counter.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | from typing import TYPE_CHECKING 27 | 28 | from typing_extensions import Self, override 29 | 30 | import pymbolic.primitives as p 31 | from pymbolic.mapper import CachedMapper, CombineMapper 32 | from pymbolic.typing import ArithmeticExpression 33 | 34 | 35 | if TYPE_CHECKING: 36 | from collections.abc import Callable, Iterable 37 | 38 | 39 | class FlopCounterBase(CombineMapper[ArithmeticExpression, []]): 40 | @override 41 | def combine(self, values: Iterable[ArithmeticExpression]) -> ArithmeticExpression: 42 | return sum(values) 43 | 44 | @override 45 | def map_constant(self, expr: object) -> ArithmeticExpression: 46 | return 0 47 | 48 | @override 49 | def map_variable(self, expr: p.Variable) -> ArithmeticExpression: 50 | return 0 51 | 52 | @override 53 | def map_sum(self, expr: p.Sum | p.Product) -> ArithmeticExpression: 54 | if expr.children: 55 | return len(expr.children) - 1 + sum(self.rec(ch) for ch in expr.children) 56 | else: 57 | return 0 58 | 59 | map_product: Callable[[Self, p.Product], ArithmeticExpression] = map_sum 60 | 61 | @override 62 | def map_quotient(self, expr: p.Quotient | p.FloorDiv) -> ArithmeticExpression: 63 | return 1 + self.rec(expr.numerator) + self.rec(expr.denominator) 64 | 65 | map_floor_div: Callable[[Self, p.FloorDiv], ArithmeticExpression] = map_quotient 66 | 67 | @override 68 | def map_power(self, expr: p.Power) -> ArithmeticExpression: 69 | return 1 + self.rec(expr.base) + self.rec(expr.exponent) 70 | 71 | @override 72 | def map_if(self, expr: p.If) -> ArithmeticExpression: 73 | rec_then = self.rec(expr.then) 74 | rec_else = self.rec(expr.else_) 75 | if isinstance(rec_then, int) and isinstance(rec_else, int): 76 | eval_flops = max(rec_then, rec_else) 77 | else: 78 | eval_flops = p.Max((rec_then, rec_else)) 79 | return self.rec(expr.condition) + eval_flops 80 | 81 | 82 | class FlopCounter(CachedMapper[int, []], FlopCounterBase): # pyright: ignore[reportGeneralTypeIssues] 83 | pass 84 | 85 | 86 | class CSEAwareFlopCounter(FlopCounterBase): 87 | """A flop counter that only counts the contribution from common 88 | subexpressions once. 89 | 90 | .. warning:: 91 | 92 | You must use a fresh mapper for each new evaluation operation for which 93 | reuse may take place. 94 | """ 95 | def __init__(self): 96 | super().__init__() 97 | self.cse_seen_set: set[p.CommonSubexpression] = set() 98 | 99 | @override 100 | def map_common_subexpression(self, 101 | expr: p.CommonSubexpression 102 | ) -> ArithmeticExpression: 103 | if expr in self.cse_seen_set: 104 | return 0 105 | else: 106 | self.cse_seen_set.add(expr) 107 | return self.rec(expr.child) 108 | -------------------------------------------------------------------------------- /pymbolic/geometric_algebra/primitives.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2014 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | # This is experimental, undocumented, and could go away any second. 27 | # Consider yourself warned. 28 | 29 | from abc import ABC, abstractmethod 30 | from typing import TYPE_CHECKING, ClassVar, TypeAlias 31 | 32 | import pytools.obj_array as obj_array 33 | 34 | from pymbolic.geometric_algebra import MultiVector 35 | from pymbolic.primitives import ExpressionNode, Variable, expr_dataclass 36 | 37 | 38 | if TYPE_CHECKING: 39 | from collections.abc import Hashable 40 | 41 | from pymbolic.typing import ( 42 | ArithmeticExpression, 43 | ArithmeticExpressionContainerTc, 44 | Expression, 45 | ) 46 | 47 | 48 | NablaId: TypeAlias = "Hashable" 49 | 50 | 51 | class MultiVectorVariable(Variable): 52 | mapper_method: ClassVar[str] = "map_multivector_variable" 53 | 54 | 55 | # {{{ geometric calculus 56 | 57 | class _GeometricCalculusExpression(ExpressionNode): 58 | def stringifier(self, originating_stringifier=None): 59 | from pymbolic.geometric_algebra.mapper import StringifyMapper 60 | return StringifyMapper() 61 | 62 | 63 | @expr_dataclass() 64 | class NablaComponent(_GeometricCalculusExpression): 65 | """ 66 | .. autoattribute:: ambient_axis 67 | .. autoattribute:: NablaId 68 | """ 69 | ambient_axis: int 70 | nabla_id: NablaId 71 | 72 | 73 | @expr_dataclass() 74 | class Nabla(_GeometricCalculusExpression): 75 | """ 76 | .. autoattribute:: nabla_id 77 | """ 78 | nabla_id: NablaId 79 | 80 | 81 | @expr_dataclass() 82 | class DerivativeSource(_GeometricCalculusExpression): 83 | """ 84 | .. autoattribute:: operand 85 | .. autoattribute:: nabla_id 86 | """ 87 | operand: Expression 88 | nabla_id: Hashable 89 | 90 | 91 | class Derivative(ABC): 92 | """This mechanism cannot be used to take more than one derivative at a time. 93 | 94 | .. autoproperty:: nabla 95 | .. automethod:: dnabla 96 | .. automethod:: resolve 97 | .. automethod:: __call__ 98 | """ 99 | 100 | my_id: str 101 | _next_id: ClassVar[list[int]] = [0] 102 | 103 | def __init__(self) -> None: 104 | self.my_id = f"id{self._next_id[0]}" 105 | self._next_id[0] += 1 106 | 107 | @property 108 | def nabla(self) -> Nabla: 109 | return Nabla(self.my_id) 110 | 111 | def dnabla(self, ambient_dim: int) -> MultiVector[ArithmeticExpression]: 112 | nablas: list[ArithmeticExpression] = [ 113 | NablaComponent(axis, self.my_id) 114 | for axis in range(ambient_dim)] 115 | return MultiVector(obj_array.new_1d(nablas)) 116 | 117 | def __call__( 118 | self, operand: ArithmeticExpressionContainerTc, 119 | ) -> ArithmeticExpressionContainerTc: 120 | from pymbolic.geometric_algebra import componentwise 121 | 122 | def func(coeff: ArithmeticExpression) -> ArithmeticExpression: 123 | return DerivativeSource(coeff, self.my_id) 124 | 125 | return componentwise(func, operand) # pyright: ignore[reportReturnType] 126 | 127 | @staticmethod 128 | @abstractmethod 129 | def resolve( 130 | expr: ArithmeticExpressionContainerTc 131 | ) -> ArithmeticExpressionContainerTc: 132 | # This method will need to be overridden by codes using this 133 | # infrastructure to use the appropriate subclass of DerivativeBinder. 134 | pass 135 | 136 | # }}} 137 | 138 | # vim: foldmethod=marker 139 | -------------------------------------------------------------------------------- /pymbolic/mapper/flattener.py: -------------------------------------------------------------------------------- 1 | """ 2 | .. autoclass:: FlattenMapper 3 | :show-inheritance: 4 | 5 | .. currentmodule:: pymbolic 6 | 7 | .. autofunction:: flatten 8 | """ 9 | 10 | from __future__ import annotations 11 | 12 | 13 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 14 | 15 | __license__ = """ 16 | Permission is hereby granted, free of charge, to any person obtaining a copy 17 | of this software and associated documentation files (the "Software"), to deal 18 | in the Software without restriction, including without limitation the rights 19 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 20 | copies of the Software, and to permit persons to whom the Software is 21 | furnished to do so, subject to the following conditions: 22 | 23 | The above copyright notice and this permission notice shall be included in 24 | all copies or substantial portions of the Software. 25 | 26 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 27 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 28 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 29 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 30 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 31 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 32 | THE SOFTWARE. 33 | """ 34 | from typing import TYPE_CHECKING, cast 35 | 36 | from typing_extensions import override 37 | 38 | import pymbolic.primitives as p 39 | from pymbolic.mapper import IdentityMapper 40 | 41 | 42 | if TYPE_CHECKING: 43 | from pymbolic.typing import ArithmeticOrExpressionT, Expression 44 | 45 | 46 | class FlattenMapper(IdentityMapper[[]]): 47 | """ 48 | Applies :func:`pymbolic.primitives.flattened_sum` 49 | to :class:`~pymbolic.primitives.Sum`" 50 | and :func:`pymbolic.primitives.flattened_product` 51 | to :class:`~pymbolic.primitives.Product`." 52 | Also applies light-duty simplification to other operators. 53 | 54 | This parallels what was done implicitly in the expression node 55 | constructors. 56 | 57 | .. automethod:: is_expr_integer_valued 58 | """ 59 | 60 | def is_expr_integer_valued(self, expr: Expression, /) -> bool: 61 | """A user-supplied method to indicate whether a given *expr* is integer- 62 | valued. This enables additional simplifications that are not valid in 63 | general. The default implementation simply returns *False*. 64 | 65 | .. versionadded :: 2024.1 66 | """ 67 | return False 68 | 69 | @override 70 | def map_sum(self, expr: p.Sum, /) -> Expression: 71 | from pymbolic.primitives import flattened_sum 72 | return flattened_sum([ 73 | self.rec_arith(ch) 74 | for ch in expr.children]) 75 | 76 | @override 77 | def map_product(self, expr: p.Product, /) -> Expression: 78 | from pymbolic.primitives import flattened_product 79 | return flattened_product([ 80 | self.rec_arith(ch) 81 | for ch in expr.children]) 82 | 83 | @override 84 | def map_quotient(self, expr: p.Quotient, /) -> Expression: 85 | r_num = self.rec_arith(expr.numerator) 86 | r_den = self.rec_arith(expr.denominator) 87 | if p.is_zero(r_num): 88 | return 0 89 | if p.is_zero(r_den - 1): 90 | return r_num 91 | 92 | return expr.__class__(r_num, r_den) 93 | 94 | @override 95 | def map_floor_div(self, expr: p.FloorDiv, /) -> Expression: 96 | r_num = self.rec_arith(expr.numerator) 97 | r_den = self.rec_arith(expr.denominator) 98 | if p.is_zero(r_num): 99 | return 0 100 | if p.is_zero(r_den - 1) and self.is_expr_integer_valued(r_num): 101 | # With a denominator of 1, it's the floor function in this case. 102 | return r_num 103 | 104 | return expr.__class__(r_num, r_den) 105 | 106 | @override 107 | def map_remainder(self, expr: p.Remainder, /) -> Expression: 108 | r_num = self.rec_arith(expr.numerator) 109 | r_den = self.rec_arith(expr.denominator) 110 | assert p.is_arithmetic_expression(r_den) 111 | if p.is_zero(r_num): 112 | return 0 113 | if p.is_zero(r_den - 1) and self.is_expr_integer_valued(r_num): 114 | # mod 1 is zero for integers, however 3.1 % 1 == .1 115 | return 0 116 | 117 | return expr.__class__(r_num, r_den) 118 | 119 | @override 120 | def map_power(self, expr: p.Power, /) -> Expression: 121 | r_base = self.rec_arith(expr.base) 122 | r_exp = self.rec_arith(expr.exponent) 123 | 124 | if p.is_zero(r_exp - 1): 125 | return r_base 126 | 127 | return expr.__class__(r_base, r_exp) 128 | 129 | 130 | def flatten(expr: ArithmeticOrExpressionT) -> ArithmeticOrExpressionT: 131 | return cast("ArithmeticOrExpressionT", FlattenMapper()(expr)) 132 | -------------------------------------------------------------------------------- /pymbolic/mapper/substitutor.py: -------------------------------------------------------------------------------- 1 | """ 2 | .. autoclass:: SubstitutionMapper 3 | :show-inheritance: 4 | .. autoclass:: CachedSubstitutionMapper 5 | :show-inheritance: 6 | 7 | .. autofunction:: make_subst_func 8 | .. autofunction:: substitute 9 | """ 10 | 11 | from __future__ import annotations 12 | 13 | 14 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 15 | 16 | __license__ = """ 17 | Permission is hereby granted, free of charge, to any person obtaining a copy 18 | of this software and associated documentation files (the "Software"), to deal 19 | in the Software without restriction, including without limitation the rights 20 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 21 | copies of the Software, and to permit persons to whom the Software is 22 | furnished to do so, subject to the following conditions: 23 | 24 | The above copyright notice and this permission notice shall be included in 25 | all copies or substantial portions of the Software. 26 | 27 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 28 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 29 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 30 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 31 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 32 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 33 | THE SOFTWARE. 34 | """ 35 | 36 | from typing import TYPE_CHECKING, Any, Protocol, TypeVar 37 | 38 | from typing_extensions import override 39 | 40 | from pymbolic.mapper import CachedIdentityMapper, IdentityMapper 41 | 42 | 43 | if TYPE_CHECKING: 44 | from collections.abc import Callable, Set 45 | 46 | import optype 47 | 48 | from pymbolic.primitives import AlgebraicLeaf, Lookup, Subscript, Variable 49 | from pymbolic.typing import Expression 50 | 51 | _KT_co = TypeVar("_KT_co", covariant=True) 52 | _VT_co = TypeVar("_VT_co", covariant=True) 53 | 54 | class CanItems(Protocol[_KT_co, _VT_co]): 55 | def items(self) -> Set[tuple[_KT_co, _VT_co]]: ... 56 | 57 | 58 | class SubstitutionMapper(IdentityMapper[[]]): 59 | subst_func: Callable[[AlgebraicLeaf], Expression | None] 60 | 61 | def __init__( 62 | self, subst_func: Callable[[AlgebraicLeaf], Expression | None] 63 | ) -> None: 64 | super().__init__() 65 | self.subst_func = subst_func 66 | 67 | @override 68 | def map_variable(self, expr: Variable) -> Expression: 69 | result = self.subst_func(expr) 70 | if result is not None: 71 | return result 72 | else: 73 | return expr 74 | 75 | @override 76 | def map_subscript(self, expr: Subscript) -> Expression: 77 | result = self.subst_func(expr) 78 | if result is not None: 79 | return result 80 | else: 81 | return IdentityMapper.map_subscript(self, expr) 82 | 83 | @override 84 | def map_lookup(self, expr: Lookup) -> Expression: 85 | result = self.subst_func(expr) 86 | if result is not None: 87 | return result 88 | else: 89 | return IdentityMapper.map_lookup(self, expr) 90 | 91 | 92 | class CachedSubstitutionMapper(CachedIdentityMapper[[]], SubstitutionMapper): 93 | def __init__( 94 | self, subst_func: Callable[[AlgebraicLeaf], Expression | None] 95 | ) -> None: 96 | SubstitutionMapper.__init__(self, subst_func) 97 | super().__init__(subst_func) 98 | 99 | 100 | def make_subst_func( 101 | # "Any" here avoids the whole Mapping variance disaster 102 | # e.g. https://github.com/python/typing/issues/445 103 | variable_assignments: optype.CanGetitem[Any, Expression], 104 | ) -> Callable[[AlgebraicLeaf], Expression | None]: 105 | import pymbolic.primitives as primitives 106 | 107 | def subst_func(var: AlgebraicLeaf) -> Expression | None: 108 | try: 109 | return variable_assignments[var] 110 | except KeyError: 111 | if isinstance(var, primitives.Variable): 112 | try: 113 | return variable_assignments[var.name] 114 | except KeyError: 115 | return None 116 | else: 117 | return None 118 | 119 | return subst_func 120 | 121 | 122 | def substitute( 123 | expression: Expression, 124 | variable_assignments: CanItems[AlgebraicLeaf | str, Expression] | None 125 | = None, 126 | mapper_cls=CachedSubstitutionMapper, 127 | **kwargs: Expression, 128 | ): 129 | """ 130 | :arg mapper_cls: A :class:`type` of the substitution mapper 131 | whose instance applies the substitution. 132 | """ 133 | if variable_assignments is None: 134 | # "Any" here avoids pointless grief about variance 135 | # e.g. https://github.com/python/typing/issues/445 136 | v_ass_copied: dict[Any, Expression] = {} 137 | else: 138 | v_ass_copied = dict(variable_assignments.items()) 139 | 140 | v_ass_copied.update(kwargs) 141 | 142 | return mapper_cls(make_subst_func(v_ass_copied))(expression) 143 | -------------------------------------------------------------------------------- /pymbolic/rational.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | from sys import intern 27 | 28 | import pymbolic.primitives as primitives 29 | import pymbolic.traits as traits 30 | 31 | 32 | class Rational(primitives.ExpressionNode): 33 | def __init__(self, numerator, denominator=1): 34 | d_unit = traits.traits(denominator).get_unit(denominator) 35 | numerator /= d_unit 36 | denominator /= d_unit 37 | self.Numerator = numerator 38 | self.Denominator = denominator 39 | 40 | def _num(self): 41 | return self.Numerator 42 | numerator = property(_num) 43 | 44 | def _den(self): 45 | return self.Denominator 46 | denominator = property(_den) 47 | 48 | def __bool__(self): 49 | return bool(self.Numerator) 50 | 51 | def __neg__(self): 52 | return Rational(-self.Numerator, self.Denominator) 53 | 54 | def __eq__(self, other): 55 | if not isinstance(other, Rational): 56 | other = Rational(other) 57 | 58 | return self.Numerator == other.Numerator and \ 59 | self.Denominator == other.Denominator 60 | 61 | def __add__(self, other): 62 | newother = Rational(other) if not isinstance(other, Rational) else other 63 | 64 | try: 65 | t = traits.common_traits(self.Denominator, newother.Denominator) 66 | newden = t.lcm(self.Denominator, newother.Denominator) 67 | newnum = self.Numerator * newden/self.Denominator + \ 68 | newother.Numerator * newden/newother.Denominator 69 | gcd = t.gcd(newden, newnum) 70 | return primitives.quotient(newnum/gcd, newden/gcd) 71 | except traits.NoTraitsError: 72 | return primitives.ExpressionNode.__add__(self, other) 73 | except traits.NoCommonTraitsError: 74 | return primitives.ExpressionNode.__add__(self, other) 75 | 76 | __radd__ = __add__ 77 | 78 | def __sub__(self, other): 79 | return self.__add__(-other) 80 | 81 | def __rsub__(self, other): 82 | return (-self).__radd__(other) 83 | 84 | def __mul__(self, other): 85 | newother = Rational(other) if not isinstance(other, Rational) else other 86 | 87 | try: 88 | t = traits.common_traits(self.Numerator, newother.Numerator, 89 | self.Denominator, newother. Denominator) 90 | gcd_1 = t.gcd(self.Numerator, newother.Denominator) 91 | gcd_2 = t.gcd(newother.Numerator, self.Denominator) 92 | 93 | new_num = self.Numerator/gcd_1 * newother.Numerator/gcd_2 94 | new_denom = self.Denominator/gcd_2 * newother.Denominator/gcd_1 95 | 96 | if not (new_denom-1): 97 | return new_num 98 | 99 | return Rational(new_num, new_denom) 100 | except traits.NoTraitsError: 101 | return primitives.ExpressionNode.__mul__(self, other) 102 | except traits.NoCommonTraitsError: 103 | return primitives.ExpressionNode.__mul__(self, other) 104 | 105 | __rmul__ = __mul__ 106 | 107 | def __div__(self, other): 108 | if not isinstance(other, Rational): 109 | other = Rational(other) 110 | 111 | return self.__mul__(Rational(other.Denominator, other.Numerator)) 112 | 113 | def __rdiv__(self, other): 114 | if not isinstance(other, Rational): 115 | other = Rational(other) 116 | 117 | return Rational(self.Denominator, self.Numerator).__rmul__(other) 118 | 119 | def __pow__(self, other): 120 | return Rational(self.Denominator**other, self.Numerator**other) 121 | 122 | def __getinitargs__(self): 123 | return (self.Numerator, self.Denominator) 124 | 125 | def reciprocal(self) -> Rational: 126 | return Rational(self.Denominator, self.Numerator) 127 | 128 | mapper_method = intern("map_rational") 129 | 130 | 131 | if __name__ == "__main__": 132 | one = Rational(1) 133 | print(3 + 1/(1 - 3/(one + 17))) 134 | print(one/3 + 2*one/3) 135 | print(one/3 + 2*one/3 + 0*one/1771) 136 | -------------------------------------------------------------------------------- /pymbolic/imperative/transform.py: -------------------------------------------------------------------------------- 1 | """Imperative program representation: transformations""" 2 | from __future__ import annotations 3 | 4 | 5 | __copyright__ = "Copyright (C) 2015 Matt Wala, Andreas Kloeckner" 6 | 7 | __license__ = """ 8 | Permission is hereby granted, free of charge, to any person obtaining a copy 9 | of this software and associated documentation files (the "Software"), to deal 10 | in the Software without restriction, including without limitation the rights 11 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 12 | copies of the Software, and to permit persons to whom the Software is 13 | furnished to do so, subject to the following conditions: 14 | 15 | The above copyright notice and this permission notice shall be included in 16 | all copies or substantial portions of the Software. 17 | 18 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 19 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 20 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 21 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 22 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 23 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 24 | THE SOFTWARE. 25 | """ 26 | 27 | 28 | from typing import TYPE_CHECKING 29 | 30 | 31 | if TYPE_CHECKING: 32 | from collections.abc import Callable, Sequence 33 | 34 | from pymbolic.imperative.statement import ( 35 | BasicStatementLikeT, 36 | StatementLike, 37 | StatementLikeT, 38 | ) 39 | from pymbolic.primitives import Variable 40 | 41 | 42 | # {{{ fuse statement streams 43 | 44 | def fuse_statement_streams_with_unique_ids( 45 | statements_a: Sequence[BasicStatementLikeT], 46 | statements_b: Sequence[BasicStatementLikeT] 47 | ) -> tuple[list[BasicStatementLikeT], dict[str, str]]: 48 | new_statements = list(statements_a) 49 | from pytools import UniqueNameGenerator 50 | stmt_id_gen = UniqueNameGenerator( 51 | {stmta.id for stmta in new_statements}) 52 | 53 | b_unique_statements: list[BasicStatementLikeT] = [] 54 | old_b_id_to_new_b_id: dict[str, str] = {} 55 | for stmtb in statements_b: 56 | old_id = stmtb.id 57 | new_id = stmt_id_gen(old_id) 58 | old_b_id_to_new_b_id[old_id] = new_id 59 | 60 | b_unique_statements.append( 61 | stmtb.copy(id=new_id)) 62 | 63 | for stmtb in b_unique_statements: 64 | new_statements.append( 65 | stmtb.copy( 66 | depends_on=frozenset( 67 | old_b_id_to_new_b_id[dep_id] 68 | for dep_id in stmtb.depends_on))) 69 | 70 | return new_statements, old_b_id_to_new_b_id 71 | 72 | 73 | def fuse_instruction_streams_with_unique_ids( 74 | insns_a: Sequence[StatementLikeT], 75 | insns_b: Sequence[StatementLikeT] 76 | ): 77 | from warnings import warn 78 | warn("fuse_instruction_streams_with_unique_ids has been renamed to " 79 | "fuse_statement_streams_with_unique_ids", DeprecationWarning, 80 | stacklevel=2) 81 | 82 | return fuse_statement_streams_with_unique_ids(insns_a, insns_b) 83 | 84 | # }}} 85 | 86 | 87 | # {{{ disambiguate_identifiers 88 | 89 | def disambiguate_identifiers( 90 | statements_a: Sequence[StatementLike], 91 | statements_b: Sequence[StatementLike], 92 | should_disambiguate_name: Callable[[str], bool] | None = None, 93 | ): 94 | if should_disambiguate_name is None: 95 | should_disambiguate_name = lambda name: True # noqa: E731 96 | 97 | from pymbolic.imperative.analysis import get_all_used_identifiers 98 | 99 | id_a = get_all_used_identifiers(statements_a) 100 | id_b = get_all_used_identifiers(statements_b) 101 | 102 | from pytools import UniqueNameGenerator 103 | vng = UniqueNameGenerator(id_a | id_b) 104 | 105 | from pymbolic import var 106 | subst_b: dict[str, Variable] = {} 107 | for clash in id_a & id_b: 108 | if should_disambiguate_name(clash): 109 | unclash = vng(clash) 110 | subst_b[clash] = var(unclash) 111 | 112 | from pymbolic.mapper.substitutor import SubstitutionMapper, make_subst_func 113 | subst_map = SubstitutionMapper(make_subst_func(subst_b)) 114 | 115 | statements_b = [ 116 | stmt.map_expressions(subst_map) for stmt in statements_b] 117 | 118 | return statements_b, subst_b 119 | 120 | # }}} 121 | 122 | 123 | # {{{ disambiguate_and_fuse 124 | 125 | def disambiguate_and_fuse( 126 | statements_a: Sequence[StatementLike], 127 | statements_b: Sequence[StatementLike], 128 | should_disambiguate_name: Callable[[str], bool] | None = None, 129 | ): 130 | statements_b, subst_b = disambiguate_identifiers( 131 | statements_a, statements_b, 132 | should_disambiguate_name) 133 | 134 | fused, old_b_id_to_new_b_id = \ 135 | fuse_statement_streams_with_unique_ids( 136 | statements_a, statements_b) 137 | 138 | return fused, subst_b, old_b_id_to_new_b_id 139 | 140 | # }}} 141 | 142 | 143 | # vim: foldmethod=marker 144 | -------------------------------------------------------------------------------- /pymbolic/mapper/coefficient.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2013 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | from collections.abc import Collection, Mapping 27 | from typing import Literal, TypeAlias 28 | 29 | from typing_extensions import override 30 | 31 | import pymbolic.primitives as p 32 | from pymbolic.mapper import Mapper 33 | from pymbolic.typing import ArithmeticExpression 34 | 35 | 36 | CoeffsT: TypeAlias = Mapping[p.AlgebraicLeaf | Literal[1], ArithmeticExpression] 37 | 38 | 39 | class CoefficientCollector(Mapper[CoeffsT, []]): 40 | target_names: Collection[str] | None 41 | 42 | def __init__(self, target_names: Collection[str] | None = None) -> None: 43 | self.target_names = target_names 44 | 45 | @override 46 | def map_sum(self, expr: p.Sum, /) -> CoeffsT: 47 | stride_dicts = [self.rec(ch) for ch in expr.children] 48 | 49 | result: dict[p.AlgebraicLeaf | Literal[1], ArithmeticExpression] = {} 50 | for stride_dict in stride_dicts: 51 | for var, stride in stride_dict.items(): 52 | if var in result: 53 | result[var] += stride 54 | else: 55 | result[var] = stride 56 | 57 | return result 58 | 59 | @override 60 | def map_product(self, expr: p.Product, /) -> CoeffsT: 61 | children_coeffs = [self.rec(child) for child in expr.children] 62 | 63 | idx_of_child_with_vars = None 64 | for i, child_coeffs in enumerate(children_coeffs): 65 | for k in child_coeffs: 66 | if k != 1: 67 | if (idx_of_child_with_vars is not None 68 | and idx_of_child_with_vars != i): 69 | raise RuntimeError("nonlinear expression") 70 | idx_of_child_with_vars = i 71 | 72 | other_coeffs: ArithmeticExpression = 1 73 | for i, child_coeffs in enumerate(children_coeffs): 74 | if i != idx_of_child_with_vars: 75 | assert len(child_coeffs) == 1 76 | other_coeffs *= child_coeffs[1] 77 | 78 | if idx_of_child_with_vars is None: 79 | return {1: other_coeffs} 80 | else: 81 | return { 82 | var: p.flattened_product((other_coeffs, coeff)) 83 | for var, coeff in 84 | children_coeffs[idx_of_child_with_vars].items()} 85 | 86 | @override 87 | def map_quotient(self, expr: p.Quotient, /) -> CoeffsT: 88 | from pymbolic.primitives import Quotient 89 | d_num = dict(self.rec(expr.numerator)) 90 | d_den = self.rec(expr.denominator) 91 | # d_den should look like {1: k} 92 | if len(d_den) > 1 or 1 not in d_den: 93 | raise RuntimeError("nonlinear expression") 94 | val = d_den[1] 95 | for k in d_num: 96 | d_num[k] = p.flattened_product((d_num[k], Quotient(1, val))) 97 | return d_num 98 | 99 | @override 100 | def map_power(self, expr: p.Power, /) -> CoeffsT: 101 | d_base = self.rec(expr.base) 102 | d_exponent = self.rec(expr.exponent) 103 | # d_exponent should look like {1: k} 104 | if len(d_exponent) > 1 or 1 not in d_exponent: 105 | raise RuntimeError("nonlinear expression") 106 | # d_base should look like {1: k} 107 | if len(d_base) > 1 or 1 not in d_base: 108 | raise RuntimeError("nonlinear expression") 109 | return {1: expr} 110 | 111 | @override 112 | def map_constant(self, expr: object, /) -> CoeffsT: 113 | assert p.is_arithmetic_expression(expr) 114 | from pymbolic.primitives import is_zero 115 | return {} if is_zero(expr) else {1: expr} 116 | 117 | @override 118 | def map_variable(self, expr: p.Variable, /) -> CoeffsT: 119 | if self.target_names is None or expr.name in self.target_names: 120 | return {expr: 1} 121 | else: 122 | return {1: expr} 123 | 124 | @override 125 | def map_algebraic_leaf(self, expr: p.AlgebraicLeaf, /) -> CoeffsT: 126 | if self.target_names is None: 127 | return {expr: 1} 128 | else: 129 | return {1: expr} 130 | -------------------------------------------------------------------------------- /pymbolic/mapper/constant_folder.py: -------------------------------------------------------------------------------- 1 | """ 2 | .. autoclass:: ConstantFoldingMapperBase 3 | :show-inheritance: 4 | .. autoclass:: ConstantFoldingMapper 5 | :show-inheritance: 6 | .. autoclass:: CommutativeConstantFoldingMapperBase 7 | :show-inheritance: 8 | .. autoclass:: CommutativeConstantFoldingMapper 9 | :show-inheritance: 10 | """ 11 | from __future__ import annotations 12 | 13 | 14 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 15 | 16 | __license__ = """ 17 | Permission is hereby granted, free of charge, to any person obtaining a copy 18 | of this software and associated documentation files (the "Software"), to deal 19 | in the Software without restriction, including without limitation the rights 20 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 21 | copies of the Software, and to permit persons to whom the Software is 22 | furnished to do so, subject to the following conditions: 23 | 24 | The above copyright notice and this permission notice shall be included in 25 | all copies or substantial portions of the Software. 26 | 27 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 28 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 29 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 30 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 31 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 32 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 33 | THE SOFTWARE. 34 | """ 35 | 36 | 37 | from typing import TYPE_CHECKING 38 | 39 | from typing_extensions import Self, override 40 | 41 | import pymbolic.primitives as prim 42 | from pymbolic.mapper import ( 43 | CSECachingMapperMixin, 44 | IdentityMapper, 45 | Mapper, 46 | ) 47 | from pymbolic.typing import ArithmeticExpression, Expression 48 | 49 | 50 | if TYPE_CHECKING: 51 | from collections.abc import Callable 52 | 53 | 54 | class ConstantFoldingMapperBase(Mapper[Expression, []]): 55 | def is_constant(self, expr: Expression, /) -> bool: 56 | from pymbolic.mapper.dependency import DependencyMapper 57 | return not bool(DependencyMapper()(expr)) 58 | 59 | def evaluate(self, expr: Expression, /) -> Expression | None: 60 | from pymbolic import evaluate 61 | 62 | try: 63 | return evaluate(expr) 64 | except ValueError: 65 | return None 66 | 67 | def fold(self, 68 | expr: prim.Sum | prim.Product, 69 | op: Callable[ 70 | [ArithmeticExpression, ArithmeticExpression], 71 | ArithmeticExpression], 72 | constructor: Callable[ 73 | [tuple[ArithmeticExpression, ...]], 74 | ArithmeticExpression], 75 | ) -> Expression: 76 | klass = type(expr) 77 | 78 | constants: list[ArithmeticExpression] = [] 79 | nonconstants: list[ArithmeticExpression] = [] 80 | 81 | queue = list(expr.children) 82 | while queue: 83 | child = self.rec(queue.pop(0)) 84 | assert prim.is_arithmetic_expression(child) 85 | 86 | if isinstance(child, klass): 87 | assert isinstance(child, (prim.Sum, prim.Product)) 88 | queue = list(child.children) + queue 89 | else: 90 | if self.is_constant(child): 91 | value = self.evaluate(child) 92 | if value is None: 93 | # couldn't evaluate 94 | nonconstants.append(child) 95 | else: 96 | constants.append(value) 97 | else: 98 | nonconstants.append(child) 99 | 100 | if constants: 101 | from functools import reduce 102 | constant = reduce(op, constants) 103 | return constructor((constant, *nonconstants)) 104 | else: 105 | return constructor(tuple(nonconstants)) 106 | 107 | @override 108 | def map_sum(self, expr: prim.Sum, /) -> Expression: 109 | import operator 110 | 111 | from pymbolic.primitives import flattened_sum 112 | 113 | return self.fold(expr, operator.add, flattened_sum) 114 | 115 | 116 | class CommutativeConstantFoldingMapperBase(ConstantFoldingMapperBase): 117 | @override 118 | def map_product(self, expr: prim.Product, /) -> Expression: 119 | import operator 120 | 121 | from pymbolic.primitives import flattened_product 122 | 123 | return self.fold(expr, operator.mul, flattened_product) 124 | 125 | 126 | class ConstantFoldingMapper( 127 | CSECachingMapperMixin[Expression, []], 128 | ConstantFoldingMapperBase, 129 | IdentityMapper[[]]): 130 | 131 | map_common_subexpression_uncached: ( 132 | Callable[[Self, prim.CommonSubexpression], Expression]) = ( 133 | IdentityMapper.map_common_subexpression) 134 | 135 | 136 | class CommutativeConstantFoldingMapper( 137 | CSECachingMapperMixin[Expression, []], 138 | CommutativeConstantFoldingMapperBase, 139 | IdentityMapper[[]]): 140 | 141 | map_common_subexpression_uncached: ( 142 | Callable[[Self, prim.CommonSubexpression], Expression]) = ( 143 | IdentityMapper.map_common_subexpression) 144 | -------------------------------------------------------------------------------- /pymbolic/interop/sympy.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = """ 5 | Copyright (C) 2017 Matt Wala 6 | Copyright (C) 2009-2013 Andreas Kloeckner 7 | """ 8 | 9 | __license__ = """ 10 | Permission is hereby granted, free of charge, to any person obtaining a copy 11 | of this software and associated documentation files (the "Software"), to deal 12 | in the Software without restriction, including without limitation the rights 13 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 14 | copies of the Software, and to permit persons to whom the Software is 15 | furnished to do so, subject to the following conditions: 16 | 17 | The above copyright notice and this permission notice shall be included in 18 | all copies or substantial portions of the Software. 19 | 20 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 25 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 26 | THE SOFTWARE. 27 | """ 28 | 29 | from typing import TYPE_CHECKING, Any 30 | 31 | import sympy as sp 32 | from typing_extensions import override 33 | 34 | import pymbolic.primitives as prim 35 | from pymbolic.interop.common import ( 36 | PymbolicToSympyLikeMapper, 37 | SympyLikeExpression, 38 | SympyLikeToPymbolicMapper, 39 | ) 40 | 41 | 42 | if TYPE_CHECKING: 43 | from sympy.core.numbers import ImaginaryUnit 44 | 45 | from pymbolic.typing import Expression 46 | 47 | __doc__ = """ 48 | .. autofunction:: make_cse 49 | 50 | .. autoclass:: SympyToPymbolicMapper 51 | .. autoclass:: PymbolicToSympyMapper 52 | """ 53 | 54 | 55 | # {{{ sympy -> pymbolic 56 | 57 | class SympyToPymbolicMapper(SympyLikeToPymbolicMapper): 58 | """ 59 | .. automethod:: __call__ 60 | """ 61 | 62 | @override 63 | def function_name(self, expr: object) -> str: 64 | return type(expr).__name__ 65 | 66 | def map_ImaginaryUnit(self, expr: ImaginaryUnit) -> Expression: 67 | return 1j 68 | 69 | def map_Float(self, expr: sp.Float) -> Expression: 70 | return self.to_float(expr) 71 | 72 | def map_NumberSymbol(self, expr: sp.NumberSymbol) -> Expression: 73 | return self.to_float(expr) 74 | 75 | def map_Indexed(self, expr: sp.Indexed) -> Expression: 76 | if len(expr.args) == 2: 77 | return prim.Subscript( 78 | self.rec(expr.args[0].args[0]), 79 | self.rec(expr.args[1]), 80 | ) 81 | 82 | return prim.Subscript( 83 | self.rec(expr.args[0].args[0]), 84 | tuple([self.rec(i) for i in expr.args[1:]]) 85 | ) 86 | 87 | def map_CSE(self, expr: CSE) -> Expression: 88 | return prim.CommonSubexpression( 89 | self.rec(expr.args[0]), expr.prefix, expr.scope) 90 | 91 | def map_Piecewise(self, expr: sp.Piecewise) -> Expression: 92 | # We only handle piecewises with 2 arguments! 93 | if not len(expr.args) == 2: 94 | raise NotImplementedError 95 | 96 | # We only handle if/else cases 97 | if not (expr.args[1][1].is_Boolean and bool(expr.args[1][1]) is True): 98 | raise NotImplementedError 99 | 100 | then = self.rec(expr.args[0][0]) 101 | else_ = self.rec(expr.args[1][0]) 102 | cond = self.rec(expr.args[0][1]) 103 | return prim.If(cond, then, else_) 104 | 105 | # }}} 106 | 107 | 108 | # {{{ pymbolic -> sympy 109 | 110 | class PymbolicToSympyMapper(PymbolicToSympyLikeMapper): 111 | """ 112 | .. automethod:: __call__ 113 | """ 114 | 115 | @property 116 | @override 117 | def sym(self) -> Any: 118 | return sp 119 | 120 | @override 121 | def raise_conversion_error(self, expr: object) -> None: 122 | raise RuntimeError(f"do not know how to translate '{expr}' to sympy") 123 | 124 | @override 125 | def map_subscript(self, expr: prim.Subscript) -> SympyLikeExpression: 126 | return self.sym.Indexed( 127 | self.sym.IndexedBase(self.rec(expr.aggregate)), 128 | *[self.rec(i) for i in expr.index_tuple] 129 | ) 130 | # }}} 131 | 132 | 133 | class CSE(sp.Function): 134 | """ 135 | A function to translate to a :class:`~pymbolic.primitives.CommonSubexpression`. 136 | """ 137 | 138 | nargs: int = 1 139 | 140 | 141 | def make_cse(arg: SympyLikeExpression, 142 | prefix: str | None = None, 143 | scope: str | None = None) -> sp.Function: 144 | """Create an expression compatible with 145 | :class:`~pymbolic.primitives.CommonSubexpression`. 146 | 147 | This is used by the mappers (e.g. :class:`SympyToPymbolicMapper`) to 148 | correctly convert a subexpression to :mod:`pymbolic`'s CSE. 149 | """ 150 | result = CSE(arg) 151 | 152 | # FIXME: make these parts of the CSE class properly? 153 | result.prefix = prefix # pyright: ignore[reportAttributeAccessIssue] 154 | result.scope = scope # pyright: ignore[reportAttributeAccessIssue] 155 | 156 | return result # pyright: ignore[reportReturnType] 157 | 158 | 159 | # vim: fdm=marker 160 | -------------------------------------------------------------------------------- /pymbolic/interop/symengine.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = """ 5 | Copyright (C) 2017 Matt Wala 6 | """ 7 | 8 | __license__ = """ 9 | Permission is hereby granted, free of charge, to any person obtaining a copy 10 | of this software and associated documentation files (the "Software"), to deal 11 | in the Software without restriction, including without limitation the rights 12 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 13 | copies of the Software, and to permit persons to whom the Software is 14 | furnished to do so, subject to the following conditions: 15 | 16 | The above copyright notice and this permission notice shall be included in 17 | all copies or substantial portions of the Software. 18 | 19 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 20 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 21 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 22 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 23 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 24 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 25 | THE SOFTWARE. 26 | """ 27 | 28 | from typing import TYPE_CHECKING, Any 29 | 30 | import symengine as sp 31 | from typing_extensions import override 32 | 33 | import pymbolic.primitives as prim 34 | from pymbolic.interop.common import ( 35 | PymbolicToSympyLikeMapper, 36 | SympyLikeExpression, 37 | SympyLikeToPymbolicMapper, 38 | ) 39 | 40 | 41 | if TYPE_CHECKING: 42 | import optype 43 | 44 | from pymbolic.typing import Expression 45 | 46 | __doc__ = """ 47 | .. autoclass:: SymEngineToPymbolicMapper 48 | .. autoclass:: PymbolicToSymEngineMapper 49 | """ 50 | 51 | 52 | # {{{ symengine -> pymbolic 53 | 54 | class SymEngineToPymbolicMapper(SympyLikeToPymbolicMapper): 55 | """ 56 | .. automethod:: __call__ 57 | """ 58 | 59 | def map_Pow(self, expr: sp.Pow) -> Expression: 60 | # SymEngine likes to use as e**a to express exp(a); we undo that here. 61 | base, exp = expr.args 62 | if base == sp.E: 63 | return prim.Variable("exp")(self.rec(exp)) 64 | else: 65 | return prim.Power(self.rec(base), self.rec(exp)) 66 | 67 | def map_Constant(self, expr: object) -> Expression: 68 | return self.rec(expr.n()) 69 | 70 | def map_Complex(self, expr: object) -> Expression: 71 | return self.rec(expr.n()) 72 | 73 | def map_ComplexDouble(self, expr: optype.CanComplex) -> Expression: 74 | return complex(expr) 75 | 76 | def map_RealDouble(self, expr: optype.CanFloat) -> Expression: 77 | return self.to_float(expr) 78 | 79 | def map_Piecewise(self, expr: sp.Piecewise) -> Expression: 80 | # We only handle piecewises with 2 statements! 81 | if not len(expr.args) == 4: 82 | raise NotImplementedError 83 | 84 | # We only handle if/else cases 85 | if not (expr.args[3].is_Boolean and bool(expr.args[3]) is True): 86 | raise NotImplementedError 87 | 88 | rec_args = [self.rec(arg) for arg in expr.args[:3]] 89 | then, cond, else_ = rec_args 90 | return prim.If(cond, then, else_) 91 | 92 | @override 93 | def function_name(self, expr: object) -> str: 94 | try: 95 | # For FunctionSymbol instances 96 | return expr.get_name() 97 | except AttributeError: 98 | # For builtin functions 99 | return type(expr).__name__ 100 | 101 | @override 102 | def not_supported(self, expr: object) -> Expression: 103 | from symengine.lib.symengine_wrapper import PyFunction 104 | 105 | if isinstance(expr, PyFunction) and self.function_name(expr) == "CSE": 106 | sympy_expr = expr._sympy_() 107 | return prim.CommonSubexpression( 108 | self.rec(expr.args[0]), sympy_expr.prefix, sympy_expr.scope) 109 | elif isinstance(expr, sp.Function) and self.function_name(expr) == "CSE": 110 | return prim.CommonSubexpression( 111 | self.rec(expr.args[0]), scope=prim.cse_scope.EVALUATION) 112 | return SympyLikeToPymbolicMapper.not_supported(self, expr) 113 | 114 | # }}} 115 | 116 | 117 | # {{{ pymbolic -> symengine 118 | 119 | class PymbolicToSymEngineMapper(PymbolicToSympyLikeMapper): 120 | """ 121 | .. automethod:: __call__ 122 | """ 123 | 124 | @property 125 | @override 126 | def sym(self) -> Any: 127 | return sp 128 | 129 | @override 130 | def raise_conversion_error(self, expr: object) -> None: 131 | raise RuntimeError(f"do not know how to translate '{expr}' to symengine") 132 | 133 | 134 | # }}} 135 | 136 | 137 | CSE = sp.Function("CSE") 138 | 139 | 140 | def make_cse(arg: SympyLikeExpression, 141 | prefix: str | None = None, 142 | scope: str | None = None) -> sp.Function: 143 | # SymEngine's classes can't be inherited, but there's a 144 | # mechanism to create one based on SymPy's ones which stores 145 | # the SymPy object inside the C++ object. 146 | # This SymPy object is later retrieved to get the prefix 147 | # These conversions between SymPy and SymEngine are expensive, 148 | # so use it only if necessary. 149 | if prefix is None and scope is None: 150 | return CSE(arg) 151 | 152 | from pymbolic.interop.sympy import make_cse as make_cse_sympy 153 | sympy_result = make_cse_sympy(arg, prefix=prefix, scope=scope) 154 | 155 | return sp.sympify(sympy_result) 156 | 157 | # vim: fdm=marker 158 | -------------------------------------------------------------------------------- /pymbolic/compiler.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | __doc__ = """ 27 | .. autoclass:: CompileMapper 28 | :show-inheritance: 29 | .. autoclass:: CompiledExpression 30 | :members: 31 | """ 32 | 33 | import math 34 | from typing import TYPE_CHECKING, Any 35 | 36 | from typing_extensions import override 37 | 38 | import pymbolic.primitives as prim 39 | from pymbolic.mapper.stringifier import PREC_NONE, StringifyMapper 40 | 41 | 42 | if TYPE_CHECKING: 43 | from collections.abc import Callable, Sequence 44 | 45 | from numpy.typing import NDArray 46 | 47 | from pymbolic.typing import Expression 48 | 49 | 50 | class CompileMapper(StringifyMapper[[]]): 51 | @override 52 | def map_constant(self, expr: object, enclosing_prec: int) -> str: 53 | # work around numpy bug #1137 (locale-sensitive repr) 54 | # https://github.com/numpy/numpy/issues/1735 55 | try: 56 | import numpy 57 | except ImportError: 58 | pass 59 | else: 60 | if isinstance(expr, numpy.floating): 61 | expr = float(expr) 62 | elif isinstance(expr, numpy.complexfloating): 63 | expr = complex(expr) 64 | 65 | return repr(expr) 66 | 67 | @override 68 | def map_numpy_array(self, expr: NDArray[Any], enclosing_prec: int) -> str: 69 | def stringify_leading_dimension(ary: NDArray[Any], /) -> str: 70 | if ary.ndim == 1: 71 | def rec(expr: NDArray[Any], /) -> str: 72 | return self.rec(expr, PREC_NONE) 73 | else: 74 | rec = stringify_leading_dimension 75 | 76 | return "[{}]".format(", ".join(rec(x) for x in ary)) 77 | 78 | return "numpy.array({})".format(stringify_leading_dimension(expr)) 79 | 80 | @override 81 | def map_foreign(self, expr: object, enclosing_prec: int) -> str: 82 | return StringifyMapper.map_foreign(self, expr, enclosing_prec) 83 | 84 | 85 | class CompiledExpression: 86 | """This class encapsulates an expression compiled into Python bytecode 87 | for faster evaluation. 88 | 89 | Its instances (unlike plain lambdas) are pickleable. 90 | 91 | .. automethod:: __call__ 92 | """ 93 | 94 | _Expression: Expression 95 | _Variables: list[prim.Variable] 96 | _code: Callable[..., Any] 97 | 98 | def __init__(self, 99 | expression: Expression, 100 | variables: Sequence[str | prim.Variable] | None = None) -> None: 101 | """ 102 | :arg variables: The first arguments (as strings or 103 | :class:`pymbolic.primitives.Variable` instances) to be used for the 104 | compiled function. All variables used by the expression and not 105 | present here are added in lexicographic order. 106 | """ 107 | if variables is None: 108 | variables = [] 109 | 110 | self._compile(expression, variables) 111 | 112 | def _compile(self, 113 | expression: Expression, 114 | variables: Sequence[str | prim.Variable]) -> None: 115 | self._Expression = expression 116 | self._Variables = [prim.make_variable(v) for v in variables] 117 | ctx = self.context().copy() 118 | 119 | try: 120 | import numpy 121 | except ImportError: 122 | pass 123 | else: 124 | ctx["numpy"] = numpy 125 | 126 | from pymbolic.mapper.dependency import DependencyMapper 127 | used_variables = DependencyMapper( 128 | composite_leaves=False)(self._Expression) 129 | used_variables -= set(self._Variables) 130 | used_variables -= {prim.Variable(key) for key in list(ctx.keys())} 131 | used_variables = list(used_variables) 132 | used_variables.sort() 133 | all_variables = self._Variables + used_variables 134 | 135 | expr_s = CompileMapper()(self._Expression, PREC_NONE) 136 | func_s = "lambda {}: {}".format(",".join(str(v) for v in all_variables), 137 | expr_s) 138 | self._code = eval(func_s, ctx) 139 | 140 | def __getstate__(self) -> tuple[Any, ...]: 141 | return self._Expression, self._Variables 142 | 143 | def __setstate__(self, state: tuple[Any, ...]) -> None: 144 | self._compile(*state) 145 | 146 | def __call__(self, *args: Any) -> Any: 147 | return self._code(*args) 148 | 149 | def context(self) -> dict[str, Any]: 150 | return {"math": math} 151 | 152 | 153 | compile = CompiledExpression 154 | -------------------------------------------------------------------------------- /pymbolic/mapper/collector.py: -------------------------------------------------------------------------------- 1 | """ 2 | .. autoclass:: TermCollector 3 | :show-inheritance: 4 | """ 5 | from __future__ import annotations 6 | 7 | 8 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 9 | 10 | __license__ = """ 11 | Permission is hereby granted, free of charge, to any person obtaining a copy 12 | of this software and associated documentation files (the "Software"), to deal 13 | in the Software without restriction, including without limitation the rights 14 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 15 | copies of the Software, and to permit persons to whom the Software is 16 | furnished to do so, subject to the following conditions: 17 | 18 | The above copyright notice and this permission notice shall be included in 19 | all copies or substantial portions of the Software. 20 | 21 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 22 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 23 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 24 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 25 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 26 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 27 | THE SOFTWARE. 28 | """ 29 | 30 | from typing import TYPE_CHECKING 31 | 32 | from typing_extensions import override 33 | 34 | import pymbolic.primitives as p 35 | from pymbolic.mapper import IdentityMapper 36 | 37 | 38 | if TYPE_CHECKING: 39 | from collections.abc import Sequence, Set 40 | 41 | from pymbolic.mapper.dependency import Dependencies 42 | from pymbolic.typing import ArithmeticExpression, Expression 43 | 44 | 45 | class TermCollector(IdentityMapper[[]]): 46 | """A term collector that assumes that multiplication is commutative. 47 | 48 | Allows specifying *parameters* (a set of 49 | :class:`pymbolic.primitives.Variable` instances) that are viewed as being 50 | coefficients and are not used for term collection. 51 | """ 52 | 53 | parameters: Set[p.AlgebraicLeaf] 54 | 55 | def __init__(self, parameters: Set[p.AlgebraicLeaf] | None = None) -> None: 56 | if parameters is None: 57 | parameters = set() 58 | 59 | self.parameters = parameters 60 | 61 | def get_dependencies(self, expr: Expression) -> Dependencies: 62 | from pymbolic.mapper.dependency import DependencyMapper 63 | return DependencyMapper()(expr) 64 | 65 | def split_term(self, mul_term: Expression) -> tuple[ 66 | Set[tuple[ArithmeticExpression, ArithmeticExpression]], 67 | ArithmeticExpression 68 | ]: 69 | """Returns a pair consisting of: 70 | - a frozenset of (base, exponent) pairs 71 | - a product of coefficients (i.e. constants and parameters) 72 | 73 | The set takes care of order-invariant comparison for us and is hashable. 74 | 75 | The argument `product' has to be fully expanded already. 76 | """ 77 | from pymbolic.primitives import AlgebraicLeaf, Power, Product 78 | 79 | def base(term: Expression) -> ArithmeticExpression: 80 | if isinstance(term, Power): 81 | return term.base 82 | else: 83 | assert p.is_arithmetic_expression(term) 84 | return term 85 | 86 | def exponent(term: Expression) -> ArithmeticExpression: 87 | if isinstance(term, Power): 88 | return term.exponent 89 | else: 90 | return 1 91 | 92 | if isinstance(mul_term, Product): 93 | terms: Sequence[Expression] = mul_term.children 94 | elif (isinstance(mul_term, Power | AlgebraicLeaf) 95 | or not bool(self.get_dependencies(mul_term))): 96 | terms = [mul_term] 97 | else: 98 | raise RuntimeError(f"split_term expects a multiplicative term: {mul_term}") 99 | 100 | base2exp: dict[ArithmeticExpression, ArithmeticExpression] = {} 101 | for term in terms: 102 | mybase = base(term) 103 | myexp = exponent(term) 104 | 105 | if mybase in base2exp: 106 | base2exp[mybase] += myexp 107 | else: 108 | base2exp[mybase] = myexp 109 | 110 | coefficients: list[ArithmeticExpression] = [] 111 | cleaned_base2exp: dict[ArithmeticExpression, ArithmeticExpression] = {} 112 | for item_base, item_exp in base2exp.items(): 113 | term = item_base**item_exp 114 | if self.get_dependencies(term) <= self.parameters: 115 | coefficients.append(term) 116 | else: 117 | cleaned_base2exp[item_base] = item_exp 118 | 119 | base_exp_set = frozenset( 120 | (base, exp) for base, exp in cleaned_base2exp.items()) 121 | coeffs = self.rec(p.flattened_product(coefficients)) 122 | assert p.is_arithmetic_expression(coeffs) 123 | 124 | return base_exp_set, coeffs 125 | 126 | @override 127 | def map_sum(self, expr: p.Sum, /) -> Expression: 128 | term2coeff: dict[ 129 | Set[tuple[ArithmeticExpression, ArithmeticExpression]], 130 | ArithmeticExpression] = {} 131 | for child in expr.children: 132 | term, coeff = self.split_term(child) 133 | term2coeff[term] = term2coeff.get(term, 0) + coeff 134 | 135 | def rep2term( 136 | rep: Set[tuple[ArithmeticExpression, ArithmeticExpression]] 137 | ) -> ArithmeticExpression: 138 | return p.flattened_product([base**exp for base, exp in rep]) 139 | 140 | result = p.flattened_sum([ 141 | coeff*rep2term(termrep) for termrep, coeff in term2coeff.items() 142 | ]) 143 | return result 144 | -------------------------------------------------------------------------------- /test/test_sympy.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2017 Matt Wala" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | import pytest 27 | 28 | import pymbolic.primitives as prim 29 | 30 | 31 | x_, y_, i_, j_ = (prim.Variable(s) for s in ["x", "y", "i", "j"]) 32 | 33 | 34 | # {{{ to pymbolic test 35 | 36 | def _test_to_pymbolic(mapper, sym, use_symengine): 37 | x, y = sym.symbols("x,y") 38 | 39 | assert mapper(sym.Rational(3, 4)) == prim.Quotient(3, 4) 40 | assert mapper(sym.Integer(6)) == 6 41 | 42 | if not use_symengine: 43 | assert mapper(sym.Subs(x**2, (x,), (y,))) == \ 44 | prim.Substitution(x_**2, ("x",), (y_,)) 45 | deriv = sym.Derivative(x**2, x) 46 | assert mapper(deriv) == prim.Derivative(x_**2, ("x",)) 47 | else: 48 | assert mapper(sym.Subs(x**2, (x,), (y,))) == \ 49 | y_**2 50 | deriv = sym.Derivative(x**2, x) 51 | assert mapper(deriv) == 2*x_ 52 | 53 | # functions 54 | assert mapper(sym.Function("f")(x)) == prim.Variable("f")(x_) 55 | assert mapper(sym.exp(x)) == prim.Variable("exp")(x_) 56 | 57 | # indexed accesses 58 | i, j = sym.symbols("i,j") 59 | if not use_symengine: 60 | idx = sym.Indexed(x, i, j) 61 | else: 62 | idx = sym.Function("Indexed")(x, i, j) 63 | assert mapper(idx) == x_[i_, j_] 64 | 65 | # constants 66 | import math 67 | # FIXME: Why isn't this exact? 68 | assert abs(mapper(sym.pi) - math.pi) < 1e-14 69 | assert abs(mapper(sym.E) - math.e) < 1e-14 70 | assert mapper(sym.I) == 1j 71 | 72 | # }}} 73 | 74 | 75 | def test_symengine_to_pymbolic(): 76 | sym = pytest.importorskip("symengine") 77 | from pymbolic.interop.symengine import SymEngineToPymbolicMapper 78 | mapper = SymEngineToPymbolicMapper() 79 | 80 | _test_to_pymbolic(mapper, sym, True) 81 | 82 | 83 | def test_sympy_to_pymbolic(): 84 | sym = pytest.importorskip("sympy") 85 | from pymbolic.interop.sympy import SympyToPymbolicMapper 86 | mapper = SympyToPymbolicMapper() 87 | 88 | _test_to_pymbolic(mapper, sym, False) 89 | 90 | 91 | # {{{ from pymbolic test 92 | 93 | def _test_from_pymbolic(mapper, sym, use_symengine): 94 | x, y = sym.symbols("x,y") 95 | 96 | assert mapper(x_ + y_) == x + y 97 | assert mapper(x_ * y_) == x * y 98 | assert mapper(x_ ** 2) == x ** 2 99 | 100 | assert mapper(prim.Substitution(x_**2, ("x",), (y_,))) == \ 101 | sym.Subs(x**2, (x,), (y,)) 102 | deriv = sym.Derivative(x**2, x) 103 | assert mapper(prim.Derivative(x_**2, ("x",))) == deriv 104 | floordiv = sym.floor(x / y) 105 | assert mapper(prim.FloorDiv(x_, y_)) == floordiv 106 | 107 | if use_symengine: 108 | assert mapper(x_[0]) == sym.Function("Indexed")("x", 0) 109 | else: 110 | i, j = sym.symbols("i,j") 111 | assert mapper(x_[i_, j_]) == sym.Indexed(x, i, j) 112 | 113 | assert mapper(prim.Variable("f")(x_)) == sym.Function("f")(x) 114 | 115 | # }}} 116 | 117 | 118 | def test_pymbolic_to_symengine(): 119 | sym = pytest.importorskip("symengine") 120 | from pymbolic.interop.symengine import PymbolicToSymEngineMapper 121 | mapper = PymbolicToSymEngineMapper() 122 | 123 | _test_from_pymbolic(mapper, sym, True) 124 | 125 | 126 | def test_pymbolic_to_sympy(): 127 | sym = pytest.importorskip("sympy") 128 | from pymbolic.interop.sympy import PymbolicToSympyMapper 129 | mapper = PymbolicToSympyMapper() 130 | 131 | _test_from_pymbolic(mapper, sym, False) 132 | 133 | 134 | # {{{ roundtrip tests 135 | 136 | def _test_roundtrip(forward, backward, sym, use_symengine): 137 | exprs = [ 138 | 2 + x_, 139 | 2 * x_, 140 | x_ ** 2, 141 | x_[0], 142 | x_[i_, j_], 143 | prim.Variable("f")(x_), 144 | prim.If(prim.Comparison(x_, "<=", y_), 1, 0), 145 | ] 146 | 147 | for expr in exprs: 148 | assert expr == backward(forward(expr)) 149 | 150 | # }}} 151 | 152 | 153 | def test_pymbolic_to_sympy_roundtrip(): 154 | sym = pytest.importorskip("sympy") 155 | from pymbolic.interop.sympy import PymbolicToSympyMapper, SympyToPymbolicMapper 156 | forward = PymbolicToSympyMapper() 157 | backward = SympyToPymbolicMapper() 158 | 159 | _test_roundtrip(forward, backward, sym, False) 160 | 161 | 162 | def test_pymbolic_to_symengine_roundtrip(): 163 | sym = pytest.importorskip("symengine") 164 | from pymbolic.interop.symengine import ( 165 | PymbolicToSymEngineMapper, 166 | SymEngineToPymbolicMapper, 167 | ) 168 | forward = PymbolicToSymEngineMapper() 169 | backward = SymEngineToPymbolicMapper() 170 | 171 | _test_roundtrip(forward, backward, sym, True) 172 | 173 | 174 | if __name__ == "__main__": 175 | import sys 176 | if len(sys.argv) > 1: 177 | exec(sys.argv[1]) 178 | else: 179 | from pytest import main 180 | main([__file__]) 181 | -------------------------------------------------------------------------------- /pymbolic/imperative/utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = """ 5 | Copyright (C) 2013 Andreas Kloeckner 6 | Copyright (C) 2014 Matt Wala 7 | """ 8 | 9 | __license__ = """ 10 | Permission is hereby granted, free of charge, to any person obtaining a copy 11 | of this software and associated documentation files (the "Software"), to deal 12 | in the Software without restriction, including without limitation the rights 13 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 14 | copies of the Software, and to permit persons to whom the Software is 15 | furnished to do so, subject to the following conditions: 16 | 17 | The above copyright notice and this permission notice shall be included in 18 | all copies or substantial portions of the Software. 19 | 20 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 25 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 26 | THE SOFTWARE. 27 | """ 28 | 29 | import logging 30 | 31 | 32 | logger = logging.getLogger(__name__) 33 | 34 | 35 | # {{{ graphviz / dot export 36 | 37 | def _default_preamble_hook(): 38 | # Sets default attributes for nodes and edges. 39 | yield 'node [shape="box"];' 40 | yield 'edge [dir="back"];' 41 | 42 | 43 | def get_dot_dependency_graph( 44 | statements, use_stmt_ids=None, 45 | preamble_hook=_default_preamble_hook, 46 | additional_lines_hook=list, 47 | statement_stringifier=None, 48 | 49 | # deprecated 50 | use_insn_ids=None,): 51 | """Return a string in the `dot `__ language depicting 52 | dependencies among kernel statements. 53 | 54 | :arg statements: A sequence of statements, each of which is stringified by 55 | calling *statement_stringifier*. 56 | :arg statement_stringifier: The function to use for stringifying the 57 | statements. The default stringifier uses :class:`str` and escapes all 58 | double quotes (``"``) in the string representation. 59 | :arg preamble_hook: A function that returns an iterable of lines 60 | to add at the beginning of the graph 61 | :arg additional_lines_hook: A function that returns an iterable 62 | of lines to add at the end of the graph 63 | 64 | """ 65 | if statement_stringifier is None: 66 | def statement_stringifier(s): 67 | return str(s).replace('"', r'\"') 68 | 69 | if use_stmt_ids is not None and use_insn_ids is not None: 70 | raise TypeError("may not specify both use_stmt_ids and use_insn_ids") 71 | 72 | if use_insn_ids is not None: 73 | use_stmt_ids = use_insn_ids 74 | from warnings import warn 75 | warn("'use_insn_ids' is deprecated. Use 'use_stmt_ids' instead.", 76 | DeprecationWarning, stacklevel=2) 77 | 78 | def get_node_attrs(stmt): 79 | if use_stmt_ids: 80 | stmt_label = stmt.id 81 | tooltip = statement_stringifier(stmt) 82 | else: 83 | stmt_label = statement_stringifier(stmt) 84 | tooltip = stmt.id 85 | 86 | return f'label="{stmt_label}",shape="box",tooltip="{tooltip}"' 87 | 88 | lines = list(preamble_hook()) 89 | lines.append("rankdir=BT;") 90 | dep_graph = {} 91 | 92 | # maps (oriented) edge onto annotation string 93 | annotation_dep_graph = {} 94 | 95 | for stmt in statements: 96 | lines.append('"{}" [{}];'.format(stmt.id, get_node_attrs(stmt))) 97 | for dep in stmt.depends_on: 98 | dep_graph.setdefault(stmt.id, set()).add(dep) 99 | 100 | if 0: 101 | for dep in stmt.then_depends_on: 102 | annotation_dep_graph[stmt.id, dep] = "then" 103 | for dep in stmt.else_depends_on: 104 | annotation_dep_graph[stmt.id, dep] = "else" 105 | 106 | # {{{ O(n^3) (i.e. slow) transitive reduction 107 | 108 | # first, compute transitive closure by fixed point iteration 109 | while True: 110 | changed_something = False 111 | 112 | for stmt_1 in dep_graph: 113 | for stmt_2 in dep_graph.get(stmt_1, set()).copy(): 114 | for stmt_3 in dep_graph.get(stmt_2, set()).copy(): 115 | if stmt_3 not in dep_graph.get(stmt_1, set()): 116 | changed_something = True 117 | dep_graph[stmt_1].add(stmt_3) 118 | 119 | if not changed_something: 120 | break 121 | 122 | for stmt_1 in dep_graph: 123 | for stmt_2 in dep_graph.get(stmt_1, set()).copy(): 124 | for stmt_3 in dep_graph.get(stmt_2, set()).copy(): 125 | if stmt_3 in dep_graph.get(stmt_1, set()): 126 | dep_graph[stmt_1].remove(stmt_3) 127 | 128 | # }}} 129 | 130 | for stmt_1 in dep_graph: 131 | for stmt_2 in dep_graph.get(stmt_1, set()): 132 | lines.append(f"{stmt_1} -> {stmt_2}") 133 | 134 | for (stmt_1, stmt_2), annot in annotation_dep_graph.items(): 135 | lines.append(f'{stmt_2} -> {stmt_1} [label="{annot}", style="dashed"]') 136 | 137 | lines.extend(additional_lines_hook()) 138 | 139 | return "digraph code {\n%s\n}" % ("\n".join(lines)) 140 | 141 | # }}} 142 | 143 | 144 | # {{{ graphviz / dot interactive show 145 | 146 | def show_dot(dot_code, output_to=None): 147 | from warnings import warn 148 | warn("pymbolic.imperative.utils.show_dot is deprecated. " 149 | "It will stop working in July 2023. " 150 | "Please use pytools.graphviz.show_dot instead.", 151 | DeprecationWarning, stacklevel=2) 152 | 153 | from pytools.graphviz import show_dot 154 | return show_dot(dot_code, output_to) 155 | # }}} 156 | 157 | # vim: fdm=marker 158 | -------------------------------------------------------------------------------- /test/test_maxima.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | import contextlib 27 | 28 | import pytest 29 | 30 | from pymbolic.interop.maxima import MaximaKernel 31 | 32 | 33 | # {{{ check for maxima 34 | 35 | def _check_maxima() -> None: 36 | global MAXIMA_UNAVAILABLE 37 | 38 | import os 39 | executable = os.environ.get("PYMBOLIC_MAXIMA_EXECUTABLE", "maxima") 40 | 41 | try: 42 | knl = MaximaKernel(executable=executable) 43 | MAXIMA_UNAVAILABLE = False 44 | knl.shutdown() 45 | except (ImportError, RuntimeError) as exc: 46 | MAXIMA_UNAVAILABLE = True # pyright: ignore[reportConstantRedefinition] 47 | print(f"MAXIMA_UNAVAILABLE: {exc}") 48 | 49 | 50 | _check_maxima() 51 | 52 | # }}} 53 | 54 | 55 | @pytest.mark.skipif(MAXIMA_UNAVAILABLE, reason="maxima cannot be launched") 56 | def test_kernel() -> None: 57 | knl = MaximaKernel() 58 | 59 | knl.exec_str("k:1/(sqrt((x0-(a+t*b))^2+(y0-(c+t*d))^2+(z0-(e+t*f))^2))") 60 | knl.eval_str("sum(diff(k, t,deg)*t^deg,deg,0,6)") 61 | assert knl.eval_str("2+2").strip() == "4" 62 | 63 | knl.shutdown() 64 | 65 | 66 | @pytest.fixture 67 | def knl(request: pytest.FixtureRequest) -> MaximaKernel: 68 | knl = MaximaKernel() 69 | request.addfinalizer(knl.shutdown) 70 | return knl 71 | 72 | 73 | @pytest.mark.skipif(MAXIMA_UNAVAILABLE, reason="maxima cannot be launched") 74 | def test_setup(knl: MaximaKernel) -> None: 75 | knl.clean_eval_str_with_setup( 76 | ["k:1/(sqrt((x0-(a+t*b))^2+(y0-(c+t*d))^2+(z0-(e+t*f))^2))"], 77 | "sum(diff(k, t,deg)*t^deg,deg,0,6)") 78 | 79 | 80 | @pytest.mark.skipif(MAXIMA_UNAVAILABLE, reason="maxima cannot be launched") 81 | def test_error(knl: MaximaKernel) -> None: 82 | from pymbolic.interop.maxima import MaximaError 83 | with contextlib.suppress(MaximaError): 84 | knl.eval_str("))(!") 85 | 86 | with contextlib.suppress(MaximaError): 87 | knl.eval_str("integrate(1/(x^3*(a+b*x)^(1/3)),x)") 88 | 89 | 90 | @pytest.mark.skipif(MAXIMA_UNAVAILABLE, reason="maxima cannot be launched") 91 | def test_strict_round_trip(knl: MaximaKernel) -> None: 92 | from pymbolic import parse 93 | from pymbolic.primitives import Quotient 94 | 95 | exprs = [ 96 | 2j, 97 | parse("x**y"), 98 | Quotient(1, 2), 99 | parse("exp(x)") 100 | ] 101 | for expr in exprs: 102 | result = knl.eval_expr(expr) 103 | round_trips_correctly = result == expr 104 | if not round_trips_correctly: 105 | print("ORIGINAL:") 106 | print("") 107 | print(expr) 108 | print("") 109 | print("POST-MAXIMA:") 110 | print("") 111 | print(result) 112 | assert round_trips_correctly 113 | 114 | 115 | @pytest.mark.skipif(MAXIMA_UNAVAILABLE, reason="maxima cannot be launched") 116 | def test_lax_round_trip(knl: MaximaKernel) -> None: 117 | from pymbolic.interop.maxima import MaximaParser 118 | k_setup = [ 119 | "k:1/(sqrt((x0-(a+t*b))^2+(y0-(c+t*d))^2))", 120 | "result:sum(diff(k, t,deg)*t^deg,deg,0,4)", 121 | ] 122 | parsed = MaximaParser()( 123 | knl.clean_eval_str_with_setup(k_setup, "result")) 124 | 125 | assert knl.clean_eval_expr_with_setup( 126 | [*k_setup, ("result2", parsed)], 127 | "ratsimp(result-result2)") == 0 128 | 129 | 130 | @pytest.mark.skipif(MAXIMA_UNAVAILABLE, reason="maxima cannot be launched") 131 | def test_parse_matrix(knl: MaximaKernel) -> None: 132 | z = knl.clean_eval_str_with_setup([ 133 | "A:matrix([1,2+0.3*dt], [3,4])", 134 | "B:matrix([1,1], [0,1])", 135 | ], 136 | "A.B") 137 | 138 | from pymbolic.interop.maxima import MaximaParser 139 | print(MaximaParser()(z)) 140 | 141 | 142 | @pytest.mark.skipif(MAXIMA_UNAVAILABLE, reason="maxima cannot be launched") 143 | def test_diff() -> None: 144 | from pymbolic import parse 145 | from pymbolic.interop.maxima import diff 146 | diff(parse("sqrt(x**2+y**2)"), parse("x")) 147 | 148 | 149 | @pytest.mark.skipif(MAXIMA_UNAVAILABLE, reason="maxima cannot be launched") 150 | def test_long_command(knl: MaximaKernel) -> None: 151 | # Seems to fail with an encoding error on pexpect 4.8 and Py3.8. 152 | # -AK, 2020-07-13 153 | # from pymbolic.interop.maxima import set_debug 154 | # set_debug(4) 155 | knl.eval_str("+".join(["1"]*16384)) 156 | 157 | 158 | @pytest.mark.skipif(MAXIMA_UNAVAILABLE, reason="maxima cannot be launched") 159 | def test_restart(knl: MaximaKernel) -> None: 160 | knl = MaximaKernel() 161 | knl.restart() 162 | knl.eval_str("1") 163 | knl.shutdown() 164 | 165 | 166 | if __name__ == "__main__": 167 | import sys 168 | if len(sys.argv) > 1: 169 | exec(sys.argv[1]) 170 | else: 171 | from pytest import main 172 | main([__file__]) 173 | -------------------------------------------------------------------------------- /doc/Makefile: -------------------------------------------------------------------------------- 1 | # Makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line. 5 | SPHINXOPTS = 6 | SPHINXBUILD = python `which sphinx-build` 7 | PAPER = 8 | BUILDDIR = _build 9 | 10 | # Internal variables. 11 | PAPEROPT_a4 = -D latex_paper_size=a4 12 | PAPEROPT_letter = -D latex_paper_size=letter 13 | ALLSPHINXOPTS = -d $(BUILDDIR)/doctrees $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) . 14 | # the i18n builder cannot share the environment and doctrees with the others 15 | I18NSPHINXOPTS = $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) . 16 | 17 | .PHONY: help clean html dirhtml singlehtml pickle json htmlhelp qthelp devhelp epub latex latexpdf text man changes linkcheck doctest gettext 18 | 19 | help: 20 | @echo "Please use \`make ' where is one of" 21 | @echo " html to make standalone HTML files" 22 | @echo " dirhtml to make HTML files named index.html in directories" 23 | @echo " singlehtml to make a single large HTML file" 24 | @echo " pickle to make pickle files" 25 | @echo " json to make JSON files" 26 | @echo " htmlhelp to make HTML files and a HTML help project" 27 | @echo " qthelp to make HTML files and a qthelp project" 28 | @echo " devhelp to make HTML files and a Devhelp project" 29 | @echo " epub to make an epub" 30 | @echo " latex to make LaTeX files, you can set PAPER=a4 or PAPER=letter" 31 | @echo " latexpdf to make LaTeX files and run them through pdflatex" 32 | @echo " text to make text files" 33 | @echo " man to make manual pages" 34 | @echo " texinfo to make Texinfo files" 35 | @echo " info to make Texinfo files and run them through makeinfo" 36 | @echo " gettext to make PO message catalogs" 37 | @echo " changes to make an overview of all changed/added/deprecated items" 38 | @echo " linkcheck to check all external links for integrity" 39 | @echo " doctest to run all doctests embedded in the documentation (if enabled)" 40 | 41 | clean: 42 | -rm -rf $(BUILDDIR)/* 43 | 44 | html: 45 | $(SPHINXBUILD) -b html $(ALLSPHINXOPTS) $(BUILDDIR)/html 46 | @echo 47 | @echo "Build finished. The HTML pages are in $(BUILDDIR)/html." 48 | 49 | dirhtml: 50 | $(SPHINXBUILD) -b dirhtml $(ALLSPHINXOPTS) $(BUILDDIR)/dirhtml 51 | @echo 52 | @echo "Build finished. The HTML pages are in $(BUILDDIR)/dirhtml." 53 | 54 | singlehtml: 55 | $(SPHINXBUILD) -b singlehtml $(ALLSPHINXOPTS) $(BUILDDIR)/singlehtml 56 | @echo 57 | @echo "Build finished. The HTML page is in $(BUILDDIR)/singlehtml." 58 | 59 | pickle: 60 | $(SPHINXBUILD) -b pickle $(ALLSPHINXOPTS) $(BUILDDIR)/pickle 61 | @echo 62 | @echo "Build finished; now you can process the pickle files." 63 | 64 | json: 65 | $(SPHINXBUILD) -b json $(ALLSPHINXOPTS) $(BUILDDIR)/json 66 | @echo 67 | @echo "Build finished; now you can process the JSON files." 68 | 69 | htmlhelp: 70 | $(SPHINXBUILD) -b htmlhelp $(ALLSPHINXOPTS) $(BUILDDIR)/htmlhelp 71 | @echo 72 | @echo "Build finished; now you can run HTML Help Workshop with the" \ 73 | ".hhp project file in $(BUILDDIR)/htmlhelp." 74 | 75 | qthelp: 76 | $(SPHINXBUILD) -b qthelp $(ALLSPHINXOPTS) $(BUILDDIR)/qthelp 77 | @echo 78 | @echo "Build finished; now you can run "qcollectiongenerator" with the" \ 79 | ".qhcp project file in $(BUILDDIR)/qthelp, like this:" 80 | @echo "# qcollectiongenerator $(BUILDDIR)/qthelp/pymbolic.qhcp" 81 | @echo "To view the help file:" 82 | @echo "# assistant -collectionFile $(BUILDDIR)/qthelp/pymbolic.qhc" 83 | 84 | devhelp: 85 | $(SPHINXBUILD) -b devhelp $(ALLSPHINXOPTS) $(BUILDDIR)/devhelp 86 | @echo 87 | @echo "Build finished." 88 | @echo "To view the help file:" 89 | @echo "# mkdir -p $$HOME/.local/share/devhelp/pymbolic" 90 | @echo "# ln -s $(BUILDDIR)/devhelp $$HOME/.local/share/devhelp/pymbolic" 91 | @echo "# devhelp" 92 | 93 | epub: 94 | $(SPHINXBUILD) -b epub $(ALLSPHINXOPTS) $(BUILDDIR)/epub 95 | @echo 96 | @echo "Build finished. The epub file is in $(BUILDDIR)/epub." 97 | 98 | latex: 99 | $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex 100 | @echo 101 | @echo "Build finished; the LaTeX files are in $(BUILDDIR)/latex." 102 | @echo "Run \`make' in that directory to run these through (pdf)latex" \ 103 | "(use \`make latexpdf' here to do that automatically)." 104 | 105 | latexpdf: 106 | $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex 107 | @echo "Running LaTeX files through pdflatex..." 108 | $(MAKE) -C $(BUILDDIR)/latex all-pdf 109 | @echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex." 110 | 111 | text: 112 | $(SPHINXBUILD) -b text $(ALLSPHINXOPTS) $(BUILDDIR)/text 113 | @echo 114 | @echo "Build finished. The text files are in $(BUILDDIR)/text." 115 | 116 | man: 117 | $(SPHINXBUILD) -b man $(ALLSPHINXOPTS) $(BUILDDIR)/man 118 | @echo 119 | @echo "Build finished. The manual pages are in $(BUILDDIR)/man." 120 | 121 | texinfo: 122 | $(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo 123 | @echo 124 | @echo "Build finished. The Texinfo files are in $(BUILDDIR)/texinfo." 125 | @echo "Run \`make' in that directory to run these through makeinfo" \ 126 | "(use \`make info' here to do that automatically)." 127 | 128 | info: 129 | $(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo 130 | @echo "Running Texinfo files through makeinfo..." 131 | make -C $(BUILDDIR)/texinfo info 132 | @echo "makeinfo finished; the Info files are in $(BUILDDIR)/texinfo." 133 | 134 | gettext: 135 | $(SPHINXBUILD) -b gettext $(I18NSPHINXOPTS) $(BUILDDIR)/locale 136 | @echo 137 | @echo "Build finished. The message catalogs are in $(BUILDDIR)/locale." 138 | 139 | changes: 140 | $(SPHINXBUILD) -b changes $(ALLSPHINXOPTS) $(BUILDDIR)/changes 141 | @echo 142 | @echo "The overview file is in $(BUILDDIR)/changes." 143 | 144 | linkcheck: 145 | $(SPHINXBUILD) -b linkcheck $(ALLSPHINXOPTS) $(BUILDDIR)/linkcheck 146 | @echo 147 | @echo "Link check complete; look for any errors in the above output " \ 148 | "or in $(BUILDDIR)/linkcheck/output.txt." 149 | 150 | doctest: 151 | $(SPHINXBUILD) -b doctest $(ALLSPHINXOPTS) $(BUILDDIR)/doctest 152 | @echo "Testing of doctests in the sources finished, look at the " \ 153 | "results in $(BUILDDIR)/doctest/output.txt." 154 | -------------------------------------------------------------------------------- /pymbolic/mapper/distributor.py: -------------------------------------------------------------------------------- 1 | """ 2 | .. autoclass:: ExpressionCallable 3 | .. autoclass:: DistributeMapper 4 | :show-inheritance: 5 | 6 | .. autofunction:: distribute 7 | """ 8 | from __future__ import annotations 9 | 10 | 11 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 12 | 13 | __license__ = """ 14 | Permission is hereby granted, free of charge, to any person obtaining a copy 15 | of this software and associated documentation files (the "Software"), to deal 16 | in the Software without restriction, including without limitation the rights 17 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 18 | copies of the Software, and to permit persons to whom the Software is 19 | furnished to do so, subject to the following conditions: 20 | 21 | The above copyright notice and this permission notice shall be included in 22 | all copies or substantial portions of the Software. 23 | 24 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 25 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 26 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 27 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 28 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 29 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 30 | THE SOFTWARE. 31 | """ 32 | 33 | from typing import TYPE_CHECKING, Protocol 34 | 35 | from typing_extensions import override 36 | 37 | import pymbolic.primitives as p 38 | from pymbolic.mapper import IdentityMapper 39 | from pymbolic.mapper.collector import TermCollector 40 | from pymbolic.mapper.constant_folder import CommutativeConstantFoldingMapper 41 | 42 | 43 | if TYPE_CHECKING: 44 | from pymbolic.typing import ArithmeticExpression, Expression 45 | 46 | 47 | class ExpressionCallable(Protocol): 48 | """Inherits: :class:`typing.Protocol`. 49 | 50 | .. automethod:: __call__ 51 | """ 52 | 53 | def __call__(self, expr: Expression, /) -> Expression: ... 54 | 55 | 56 | class DistributeMapper(IdentityMapper[[]]): 57 | """Example usage: 58 | 59 | .. doctest:: 60 | 61 | >>> import pymbolic.primitives as p 62 | >>> x = p.Variable("x") 63 | >>> expr = (x+1)**7 64 | >>> from pymbolic.mapper.distributor import DistributeMapper as DM 65 | >>> print DM()(expr) # doctest: +SKIP 66 | 7*x**6 + 21*x**5 + 21*x**2 + 35*x**3 + 1 + 35*x**4 + 7*x + x**7 67 | """ 68 | 69 | collector: ExpressionCallable 70 | const_folder: ExpressionCallable 71 | 72 | def __init__(self, 73 | collector: ExpressionCallable | None = None, 74 | const_folder: ExpressionCallable | None = None) -> None: 75 | if collector is None: 76 | collector = TermCollector() 77 | if const_folder is None: 78 | const_folder = CommutativeConstantFoldingMapper() 79 | 80 | self.collector = collector 81 | self.const_folder = const_folder 82 | 83 | def collect(self, expr: Expression) -> Expression: 84 | return self.collector(self.const_folder(expr)) 85 | 86 | @override 87 | def map_sum(self, expr: p.Sum, /) -> Expression: 88 | res = IdentityMapper.map_sum(self, expr) 89 | if isinstance(res, p.Sum): 90 | return self.collect(res) 91 | else: 92 | return res 93 | 94 | @override 95 | def map_product(self, expr: p.Product, /) -> Expression: 96 | def dist(prod: ArithmeticExpression) -> ArithmeticExpression: 97 | if not isinstance(prod, p.Product): 98 | return prod 99 | 100 | leading: list[ArithmeticExpression] = [] 101 | for i in prod.children: 102 | if isinstance(i, p.Sum): 103 | break 104 | else: 105 | leading.append(i) 106 | 107 | if len(leading) == len(prod.children): 108 | # no more sums found 109 | result = p.flattened_product(prod.children) 110 | return result 111 | else: 112 | sum = prod.children[len(leading)] 113 | assert isinstance(sum, p.Sum) 114 | 115 | rest = prod.children[len(leading)+1:] 116 | rest = dist(p.Product(rest)) if rest else 1 117 | 118 | result = self.collect(p.flattened_sum([ 119 | p.flattened_product(leading) * dist(sumchild*rest) 120 | for sumchild in sum.children 121 | ])) 122 | assert p.is_arithmetic_expression(result) 123 | 124 | return result 125 | 126 | return dist(IdentityMapper.map_product(self, expr)) 127 | 128 | @override 129 | def map_quotient(self, expr: p.Quotient, /) -> Expression: 130 | if p.is_zero(expr.numerator - 1): 131 | return expr 132 | else: 133 | # not the smartest thing we can do, but at least *something* 134 | return p.flattened_product([ 135 | type(expr)(1, self.rec_arith(expr.denominator)), 136 | self.rec_arith(expr.numerator) 137 | ]) 138 | 139 | @override 140 | def map_power(self, expr: p.Power, /) -> Expression: 141 | from pymbolic.primitives import Sum 142 | 143 | newbase = self.rec(expr.base) 144 | if isinstance(newbase, p.Product): 145 | return self.rec(p.flattened_product([ 146 | child**expr.exponent for child in newbase.children 147 | ])) 148 | 149 | if isinstance(expr.exponent, int): 150 | if isinstance(newbase, Sum): 151 | return self.rec(p.flattened_product(expr.exponent*(newbase,))) 152 | else: 153 | return IdentityMapper.map_power(self, expr) 154 | else: 155 | return IdentityMapper.map_power(self, expr) 156 | 157 | 158 | def distribute( 159 | expr: Expression, 160 | parameters: frozenset[p.AlgebraicLeaf] | None = None, 161 | commutative: bool = True 162 | ) -> Expression: 163 | if parameters is None: 164 | parameters = frozenset() 165 | 166 | if commutative: 167 | return DistributeMapper(TermCollector(parameters))(expr) 168 | else: 169 | return DistributeMapper(lambda x: x)(expr) 170 | -------------------------------------------------------------------------------- /doc/misc.rst: -------------------------------------------------------------------------------- 1 | Installation 2 | ============ 3 | 4 | This command should install :mod:`pymbolic`:: 5 | 6 | pip install pymbolic 7 | 8 | You may need to run this with :command:`sudo` if you are not in a virtual environment 9 | (not recommended). If you don't already have `pip `__, 10 | run this beforehand:: 11 | 12 | python -m ensurepip 13 | 14 | For a more manual installation, download the source, unpack it, and run:: 15 | 16 | pip install . 17 | 18 | This should also install all the required dependencies (see ``pyproject.toml`` 19 | for a complete list). 20 | 21 | For development, you may want to install in `editable mode 22 | `__:: 23 | 24 | pip install --no-build-isolation --editable .[test] 25 | 26 | Why pymbolic when there's already sympy? 27 | ======================================== 28 | 29 | (This is extracted from an email I (Andreas) sent to Aaron Meurer and Anthony 30 | Scopatz.) 31 | 32 | So why not use :mod:`sympy` as an AST for DSLs and code generation? It's a good 33 | question. As you read the points I make below, please bear in mind that I'm not 34 | saying this to 'attack' sympy or to diminish the achievement that it is. Very 35 | much on the contrary--as I said above, sympy does a fantastic job being a 36 | computer algebra. I just don't think it's as much in its element as an IR for 37 | code generation. Personally, I think that's perfectly fine--IMO, the tradeoffs 38 | are different for IRs and efficient computer algebra. In a sense, pymbolic 39 | competes much harder with Python's ast module for being a usable program 40 | representation than with Sympy for being a CAS. 41 | 42 | At any rate, to answer your question, here goes: 43 | 44 | * First, specifically *because* sympy is smart about its input, and will 45 | rewrite it behind your back. pymbolic is *intended* to be a dumb and 46 | static expression tree, and it will leave its input alone unless you 47 | explicitly tell it not to. In terms of floating point math or around 48 | custom node types that may or may not obey the same rules as scalars, 49 | I feel like 'leave it alone' is a safer default. 50 | 51 | * Pickling: https://github.com/sympy/sympy/issues/4297 52 | 53 | The very moment code generation starts taking more than a second or 54 | so, you'll want to implement a caching mechanism, likely using Pickle. 55 | 56 | * Extensibility of transformation constructs: sympy's built-in traversal 57 | behaviors (e.g. taking derivatives, conversion to string, code 58 | generation) aren't particularly easy to extend. It's important to 59 | understand what I'm talking about here: I would like to be able to 60 | make something that, say, is *like* taking a derivative (or 61 | evaluating, or...), but behaves just a bit differently for a few node 62 | types. This is a need that I've found to be very common in code 63 | generation. In (my understanding of) sympy, these behaviors are 64 | attached to method names, so the only way I could conceivably obtain a 65 | tweaked "diff" would be to temporarily monkeypatch "diff" for my node 66 | type, which is kind of a nonstarter. (unless I'm missing something) 67 | 68 | Pymbolic's "mapper" mechanism does somewhat better here--you 69 | simply inherit from the base behavior, implement/override a few 70 | methods, and you're done. 71 | 72 | This part is a bit of a red herring though, since this can be 73 | implemented for sympy (and, in fact, `I have 74 | `__). 75 | Also, I noticed that sympy's codegen module implements something similar (e.g. 76 | `here 77 | `__). 78 | The remaining issue is that most of sympy's behaviors aren't available to 79 | extend in this style. 80 | 81 | * Representation of code-like constructs, such as: 82 | 83 | * Indexing 84 | 85 | * Bit shifts and other bitwise ops: 86 | 87 | * Distinguishing floor-div and true-div 88 | 89 | * Attribute Access 90 | 91 | * I should also mention that pymbolic, aside from maintenance and bug 92 | fixes, is effectively 'finished'. It's pretty tiny, it's not 93 | ambitious, and it's not going to change much going forward. And that 94 | is precisely what I want from a package that provides the core data 95 | structure for something complicated and compiler-ish that I'm building 96 | on top. 97 | 98 | User-visible changes 99 | ==================== 100 | 101 | Version 2015.3 102 | -------------- 103 | 104 | .. note:: 105 | 106 | This version is currently under development. You can get snapshots from 107 | Pymbolic's `git repository `__ 108 | 109 | * Add :mod:`pymbolic.geometric_algebra`. 110 | * First documented version. 111 | 112 | .. _license: 113 | 114 | License 115 | ======= 116 | 117 | :mod:`pymbolic` is licensed to you under the MIT/X Consortium license: 118 | 119 | Copyright (c) 2008-13 Andreas Klöckner 120 | 121 | Permission is hereby granted, free of charge, to any person 122 | obtaining a copy of this software and associated documentation 123 | files (the "Software"), to deal in the Software without 124 | restriction, including without limitation the rights to use, 125 | copy, modify, merge, publish, distribute, sublicense, and/or sell 126 | copies of the Software, and to permit persons to whom the 127 | Software is furnished to do so, subject to the following 128 | conditions: 129 | 130 | The above copyright notice and this permission notice shall be 131 | included in all copies or substantial portions of the Software. 132 | 133 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 134 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES 135 | OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 136 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT 137 | HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 138 | WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 139 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 140 | OTHER DEALINGS IN THE SOFTWARE. 141 | 142 | Glossary 143 | ======== 144 | 145 | .. glossary:: 146 | 147 | mix-in 148 | See `Wikipedia article `__. 149 | 150 | Be sure to mention the mix-in before the base class being mixed in the 151 | list of base classes. This way, the mix-in can override base class 152 | behavior. 153 | -------------------------------------------------------------------------------- /pymbolic/cse.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | 4 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 5 | 6 | __license__ = """ 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in 15 | all copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 | THE SOFTWARE. 24 | """ 25 | 26 | from typing import TYPE_CHECKING, cast 27 | 28 | from typing_extensions import Self, override 29 | 30 | import pymbolic.primitives as prim 31 | from pymbolic.mapper import IdentityMapper, P, WalkMapper 32 | 33 | 34 | if TYPE_CHECKING: 35 | from collections.abc import Callable, Hashable, Iterable, Sequence, Set 36 | 37 | from pymbolic.typing import Expression 38 | 39 | 40 | COMMUTATIVE_CLASSES = (prim.Sum, prim.Product) 41 | 42 | 43 | class NormalizedKeyGetter: 44 | def __call__(self, expr: object, /) -> Hashable: 45 | if isinstance(expr, COMMUTATIVE_CLASSES): 46 | kid_count: dict[Expression, int] = {} 47 | for child in expr.children: 48 | kid_count[child] = kid_count.get(child, 0) + 1 49 | 50 | return type(expr), frozenset(kid_count.items()) 51 | else: 52 | return expr 53 | 54 | 55 | class UseCountMapper(WalkMapper[P]): 56 | subexpr_counts: dict[Hashable, int] 57 | get_key: Callable[[object], Hashable] 58 | 59 | def __init__(self, get_key: Callable[[object], Hashable]) -> None: 60 | self.subexpr_counts = {} 61 | self.get_key = get_key 62 | 63 | @override 64 | def visit(self, expr: object, /, *args: P.args, **kwargs: P.kwargs) -> bool: 65 | key = self.get_key(expr) 66 | 67 | if key in self.subexpr_counts: 68 | self.subexpr_counts[key] += 1 69 | 70 | # do not re-traverse (and thus re-count subexpressions) 71 | return False 72 | else: 73 | self.subexpr_counts[key] = 1 74 | 75 | # continue traversing 76 | return True 77 | 78 | @override 79 | def map_common_subexpression(self, expr: prim.CommonSubexpression, /, 80 | *args: P.args, **kwargs: P.kwargs) -> None: 81 | # For existing CSEs, reuse has already been decided. 82 | # Add to 83 | 84 | key = self.get_key(expr) 85 | if key in self.subexpr_counts: 86 | self.subexpr_counts[key] += 1 87 | else: 88 | # This order reversal matters: Since get_key removes the outer 89 | # CSE, need to traverse first, then add to counter. 90 | 91 | self.rec(expr.child, *args, **kwargs) 92 | self.subexpr_counts[key] = 1 93 | 94 | 95 | class CSEMapper(IdentityMapper[[]]): 96 | to_eliminate: Set[Hashable] 97 | get_key: Callable[[object], Hashable] 98 | canonical_subexprs: dict[Hashable, Expression] 99 | 100 | def __init__(self, 101 | to_eliminate: Set[Hashable], 102 | get_key: Callable[[object], Hashable]) -> None: 103 | self.to_eliminate = to_eliminate 104 | self.get_key = get_key 105 | self.canonical_subexprs = {} 106 | 107 | def get_cse(self, expr: prim.ExpressionNode, /, key: Hashable = None) -> Expression: 108 | if key is None: 109 | key = self.get_key(expr) 110 | 111 | try: 112 | return self.canonical_subexprs[key] 113 | except KeyError: 114 | new_expr = cast("Expression", prim.make_common_subexpression( 115 | getattr(IdentityMapper, expr.mapper_method)(self, expr) 116 | )) 117 | self.canonical_subexprs[key] = new_expr 118 | return new_expr 119 | 120 | @override 121 | def map_sum(self, 122 | expr: (prim.Sum | prim.Product | prim.Power 123 | | prim.Quotient | prim.Remainder | prim.FloorDiv 124 | | prim.Call 125 | ), / 126 | ) -> Expression: 127 | key = self.get_key(expr) 128 | if key in self.to_eliminate: 129 | result = self.get_cse(expr, key) 130 | return result 131 | else: 132 | return getattr(IdentityMapper, expr.mapper_method)(self, expr) 133 | 134 | map_product: Callable[[Self, prim.Product], Expression] = map_sum 135 | map_power: Callable[[Self, prim.Power], Expression] = map_sum 136 | map_quotient: Callable[[Self, prim.Quotient], Expression] = map_sum 137 | map_remainder: Callable[[Self, prim.Remainder], Expression] = map_sum 138 | map_floor_div: Callable[[Self, prim.FloorDiv], Expression] = map_sum 139 | map_call: Callable[[Self, prim.Call], Expression] = map_sum 140 | 141 | @override 142 | def map_common_subexpression(self, expr: prim.CommonSubexpression, /) -> Expression: 143 | # Avoid creating CSE(CSE(...)) 144 | if type(expr) is prim.CommonSubexpression: 145 | return prim.make_common_subexpression( 146 | self.rec(expr.child), expr.prefix, expr.scope 147 | ) 148 | else: 149 | # expr is of a derived CSE type 150 | result = self.rec(expr.child) 151 | if type(result) is prim.CommonSubexpression: 152 | result = result.child 153 | 154 | return type(expr)(result, expr.prefix, expr.scope, 155 | **expr.get_extra_properties()) 156 | 157 | @override 158 | def map_substitution(self, expr: prim.Substitution, /) -> Expression: 159 | return type(expr)( 160 | expr.child, 161 | expr.variables, 162 | tuple([self.rec(v) for v in expr.values])) 163 | 164 | 165 | def tag_common_subexpressions(exprs: Iterable[Expression]) -> Sequence[Expression]: 166 | get_key = NormalizedKeyGetter() 167 | ucm = UseCountMapper(get_key) 168 | 169 | if isinstance(exprs, prim.ExpressionNode): 170 | raise TypeError("exprs should be an iterable of expressions") 171 | 172 | for expr in exprs: 173 | ucm(expr) 174 | 175 | to_eliminate = {subexpr_key 176 | for subexpr_key, count in ucm.subexpr_counts.items() 177 | if count > 1} 178 | 179 | cse_mapper = CSEMapper(to_eliminate, get_key) 180 | result = [cse_mapper(expr) for expr in exprs] 181 | return result 182 | -------------------------------------------------------------------------------- /pymbolic/mapper/dependency.py: -------------------------------------------------------------------------------- 1 | """ 2 | .. autoclass:: DependencyMapper 3 | :show-inheritance: 4 | .. autoclass:: CachedDependencyMapper 5 | :show-inheritance: 6 | """ 7 | 8 | from __future__ import annotations 9 | 10 | 11 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 12 | 13 | __license__ = """ 14 | Permission is hereby granted, free of charge, to any person obtaining a copy 15 | of this software and associated documentation files (the "Software"), to deal 16 | in the Software without restriction, including without limitation the rights 17 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 18 | copies of the Software, and to permit persons to whom the Software is 19 | furnished to do so, subject to the following conditions: 20 | 21 | The above copyright notice and this permission notice shall be included in 22 | all copies or substantial portions of the Software. 23 | 24 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 25 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 26 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 27 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 28 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 29 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 30 | THE SOFTWARE. 31 | """ 32 | 33 | from collections.abc import Set 34 | from typing import TYPE_CHECKING, Literal, TypeAlias 35 | 36 | from typing_extensions import override 37 | 38 | import pymbolic.primitives as p 39 | from pymbolic.mapper import CachedMapper, Collector, CSECachingMapperMixin, P 40 | 41 | 42 | Dependency: TypeAlias = p.AlgebraicLeaf | p.CommonSubexpression 43 | Dependencies: TypeAlias = Set[Dependency] 44 | 45 | if not TYPE_CHECKING: 46 | DependenciesT: TypeAlias = Dependencies 47 | 48 | 49 | class DependencyMapper( 50 | CSECachingMapperMixin[Dependencies, P], 51 | Collector[p.AlgebraicLeaf | p.CommonSubexpression, P], 52 | ): 53 | """Maps an expression to the :class:`set` of expressions it 54 | is based on. The ``include_*`` arguments to the constructor 55 | determine which types of objects occur in this output set. 56 | If all are *False*, only :class:`pymbolic.primitives.Variable` 57 | instances are included. 58 | """ 59 | 60 | include_subscripts: bool 61 | include_lookups: bool 62 | include_calls: bool | Literal["descend_args"] 63 | include_cses: bool 64 | 65 | def __init__( 66 | self, 67 | include_subscripts: bool = True, 68 | include_lookups: bool = True, 69 | include_calls: bool | Literal["descend_args"] = True, 70 | include_cses: bool = False, 71 | composite_leaves: bool | None = None, 72 | ) -> None: 73 | """ 74 | :arg composite_leaves: Setting this is equivalent to setting 75 | all preceding ``include_*`` flags. 76 | """ 77 | 78 | if composite_leaves is False: 79 | include_subscripts = False 80 | include_lookups = False 81 | include_calls = False 82 | 83 | if composite_leaves is True: 84 | include_subscripts = True 85 | include_lookups = True 86 | include_calls = True 87 | 88 | assert include_calls in [True, False, "descend_args"] 89 | 90 | self.include_subscripts = include_subscripts 91 | self.include_lookups = include_lookups 92 | self.include_calls = include_calls 93 | self.include_cses = include_cses 94 | 95 | @override 96 | def map_variable( 97 | self, expr: p.Variable, /, *args: P.args, **kwargs: P.kwargs 98 | ) -> Dependencies: 99 | return {expr} 100 | 101 | @override 102 | def map_call( 103 | self, expr: p.Call, /, *args: P.args, **kwargs: P.kwargs 104 | ) -> Dependencies: 105 | if self.include_calls == "descend_args": 106 | return self.combine([ 107 | self.rec(child, *args, **kwargs) for child in expr.parameters 108 | ]) 109 | elif self.include_calls: 110 | return {expr} 111 | else: 112 | return super().map_call(expr, *args, **kwargs) 113 | 114 | @override 115 | def map_call_with_kwargs( 116 | self, expr: p.CallWithKwargs, /, *args: P.args, **kwargs: P.kwargs 117 | ) -> Dependencies: 118 | if self.include_calls == "descend_args": 119 | return self.combine( 120 | [self.rec(child, *args, **kwargs) for child in expr.parameters] 121 | + [ 122 | self.rec(val, *args, **kwargs) 123 | for _name, val in expr.kw_parameters.items() 124 | ] 125 | ) 126 | elif self.include_calls: 127 | return {expr} 128 | else: 129 | return super().map_call_with_kwargs(expr, *args, **kwargs) 130 | 131 | @override 132 | def map_lookup( 133 | self, expr: p.Lookup, /, *args: P.args, **kwargs: P.kwargs 134 | ) -> Dependencies: 135 | if self.include_lookups: 136 | return {expr} 137 | else: 138 | return super().map_lookup(expr, *args, **kwargs) 139 | 140 | @override 141 | def map_subscript( 142 | self, expr: p.Subscript, /, *args: P.args, **kwargs: P.kwargs 143 | ) -> Dependencies: 144 | if self.include_subscripts: 145 | return {expr} 146 | else: 147 | return super().map_subscript(expr, *args, **kwargs) 148 | 149 | @override 150 | def map_common_subexpression_uncached( 151 | self, expr: p.CommonSubexpression, /, *args: P.args, **kwargs: P.kwargs 152 | ) -> Dependencies: 153 | if self.include_cses: 154 | return {expr} 155 | else: 156 | return Collector.map_common_subexpression(self, expr, *args, **kwargs) 157 | 158 | @override 159 | def map_slice( 160 | self, expr: p.Slice, /, *args: P.args, **kwargs: P.kwargs 161 | ) -> Dependencies: 162 | return self.combine([ 163 | self.rec(child, *args, **kwargs) 164 | for child in expr.children 165 | if child is not None 166 | ]) 167 | 168 | @override 169 | def map_nan(self, expr: p.NaN, /, 170 | *args: P.args, **kwargs: P.kwargs) -> Dependencies: 171 | return set() 172 | 173 | 174 | class CachedDependencyMapper(CachedMapper[Dependencies, P], 175 | DependencyMapper[P]): 176 | def __init__( 177 | self, 178 | include_subscripts: bool = True, 179 | include_lookups: bool = True, 180 | include_calls: bool | Literal["descend_args"] = True, 181 | include_cses: bool = False, 182 | composite_leaves: bool | None = None, 183 | ) -> None: 184 | CachedMapper.__init__(self) 185 | DependencyMapper.__init__( 186 | self, 187 | include_subscripts=include_subscripts, 188 | include_lookups=include_lookups, 189 | include_calls=include_calls, 190 | include_cses=include_cses, 191 | composite_leaves=composite_leaves, 192 | ) 193 | -------------------------------------------------------------------------------- /pymbolic/mapper/graphviz.py: -------------------------------------------------------------------------------- 1 | """ 2 | .. autoclass:: GraphvizMapper 3 | :show-inheritance: 4 | """ 5 | from __future__ import annotations 6 | 7 | 8 | __copyright__ = "Copyright (C) 2015 Andreas Kloeckner" 9 | 10 | __license__ = """ 11 | Permission is hereby granted, free of charge, to any person obtaining a copy 12 | of this software and associated documentation files (the "Software"), to deal 13 | in the Software without restriction, including without limitation the rights 14 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 15 | copies of the Software, and to permit persons to whom the Software is 16 | furnished to do so, subject to the following conditions: 17 | 18 | The above copyright notice and this permission notice shall be included in 19 | all copies or substantial portions of the Software. 20 | 21 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 22 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 23 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 24 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 25 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 26 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 27 | THE SOFTWARE. 28 | """ 29 | 30 | from typing import TYPE_CHECKING 31 | 32 | from typing_extensions import Self, override 33 | 34 | from pymbolic.mapper import WalkMapper 35 | 36 | 37 | if TYPE_CHECKING: 38 | from collections.abc import Callable, Hashable 39 | 40 | import pymbolic.primitives as prim 41 | from pymbolic.geometric_algebra.primitives import Nabla, NablaComponent 42 | 43 | 44 | class GraphvizMapper(WalkMapper[[]]): 45 | """Produces code for `dot `__ that yields 46 | an expression tree of the traversed expression(s). 47 | 48 | .. automethod:: get_dot_code 49 | 50 | .. versionadded:: 2015.1 51 | """ 52 | 53 | lines: list[str] 54 | parent_stack: list[Hashable] 55 | 56 | next_unique_id: int 57 | nodes_visited: set[int] 58 | common_subexpressions: dict[prim.CommonSubexpression, prim.CommonSubexpression] 59 | 60 | def __init__(self) -> None: 61 | self.lines = [] 62 | self.parent_stack = [] 63 | 64 | self.next_unique_id = -1 65 | self.nodes_visited = set() 66 | self.common_subexpressions = {} 67 | 68 | def get_dot_code(self) -> str: 69 | """Return the dot source code for a previously traversed expression.""" 70 | 71 | lines = "\n".join(f" {line}" for line in self.lines) 72 | return f"digraph expression {{\n{lines}\n}}" 73 | 74 | def get_id(self, expr: object) -> str: 75 | """Generate a unique node ID for dot for *expr*""" 76 | 77 | return f"id{id(expr)}" 78 | 79 | def map_leaf(self, expr: prim.ExpressionNode): 80 | sid = self.get_id(expr) 81 | sexpr = str(expr).replace("\\", "\\\\") 82 | self.lines.append(f'{sid} [label="{sexpr}", shape=box];') 83 | 84 | if self.visit(expr, node_printed=True): 85 | self.post_visit(expr) 86 | 87 | def generate_unique_id(self): 88 | self.next_unique_id += 1 89 | return f"uid{self.next_unique_id}" 90 | 91 | @override 92 | def visit(self, 93 | expr: object, /, 94 | node_printed: bool = False, 95 | node_id: str | None = None) -> bool: 96 | # {{{ print connectivity 97 | 98 | if node_id is None: 99 | node_id = self.get_id(expr) 100 | 101 | if self.parent_stack: 102 | sid = self.get_id(self.parent_stack[-1]) 103 | self.lines.append(f"{sid} -> {node_id};") 104 | 105 | # }}} 106 | 107 | if id(expr) in self.nodes_visited: 108 | return False 109 | self.nodes_visited.add(id(expr)) 110 | 111 | if not node_printed: 112 | sid = self.get_id(expr) 113 | self.lines.append(f'{sid} [label="{type(expr).__name__}"];') 114 | 115 | self.parent_stack.append(expr) 116 | return True 117 | 118 | @override 119 | def post_visit(self, expr: object, /) -> None: 120 | self.parent_stack.pop(-1) 121 | 122 | @override 123 | def map_sum(self, expr: prim.Sum, /) -> None: 124 | sid = self.get_id(expr) 125 | self.lines.append(f'{sid} [label="+",shape=circle];') 126 | if not self.visit(expr, node_printed=True): 127 | return 128 | 129 | for child in expr.children: 130 | self.rec(child) 131 | 132 | self.post_visit(expr) 133 | 134 | @override 135 | def map_product(self, expr: prim.Product, /) -> None: 136 | sid = self.get_id(expr) 137 | self.lines.append(f'{sid} [label="*",shape=circle];') 138 | if not self.visit(expr, node_printed=True): 139 | return 140 | 141 | for child in expr.children: 142 | self.rec(child) 143 | 144 | self.post_visit(expr) 145 | 146 | @override 147 | def map_variable(self, expr: prim.Variable, /) -> None: 148 | # Shared nodes for variables do not make for pretty graphs. 149 | # So we generate our own unique IDs for them. 150 | 151 | node_id = self.generate_unique_id() 152 | 153 | self.lines.append(f'{node_id} [label="{expr.name}",shape=box];') 154 | if not self.visit(expr, node_printed=True, node_id=node_id): 155 | return 156 | 157 | self.post_visit(expr) 158 | 159 | @override 160 | def map_lookup(self, expr: prim.Lookup, /) -> None: 161 | sid = self.get_id(expr) 162 | self.lines.append(f'{sid} [label="Lookup[{expr.name}]",shape=box];') 163 | if not self.visit(expr, node_printed=True): 164 | return 165 | 166 | self.rec(expr.aggregate) 167 | self.post_visit(expr) 168 | 169 | @override 170 | def map_constant(self, expr: object) -> None: 171 | # Some constants (Python ints notably) are shared among small (and 172 | # common) values. While accurate, this doesn't make for pretty 173 | # trees. So we generate our own unique IDs for them. 174 | 175 | node_id = self.generate_unique_id() 176 | 177 | self.lines.append(f'{node_id} [label="{expr}",shape=ellipse];') 178 | if not self.visit(expr, node_printed=True, node_id=node_id): 179 | return 180 | 181 | self.post_visit(expr) 182 | 183 | @override 184 | def map_call(self, expr: prim.Call) -> None: 185 | from pymbolic.primitives import Variable 186 | 187 | if not isinstance(expr.function, Variable): 188 | return super().map_call(expr) 189 | 190 | sid = self.get_id(expr) 191 | self.lines.append(f'{sid} [label="Call[{expr.function}]",shape=box];') 192 | if not self.visit(expr, node_printed=True): 193 | return 194 | 195 | for child in expr.parameters: 196 | self.rec(child) 197 | 198 | self.post_visit(expr) 199 | 200 | @override 201 | def map_common_subexpression(self, expr: prim.CommonSubexpression) -> None: 202 | try: 203 | expr = self.common_subexpressions[expr] 204 | except KeyError: 205 | self.common_subexpressions[expr] = expr 206 | 207 | if not self.visit(expr): 208 | return 209 | 210 | self.rec(expr.child) 211 | 212 | self.post_visit(expr) 213 | 214 | # {{{ geometric algebra 215 | 216 | map_nabla_component: Callable[[Self, NablaComponent], None] = map_leaf 217 | map_nabla: Callable[[Self, Nabla], None] = map_leaf 218 | 219 | # }}} 220 | -------------------------------------------------------------------------------- /pymbolic/typing.py: -------------------------------------------------------------------------------- 1 | """ 2 | Typing helpers 3 | -------------- 4 | 5 | .. |br| raw:: html 6 | 7 |
8 | 9 | .. currentmodule:: pymbolic 10 | 11 | .. autodata:: Bool 12 | 13 | Supported boolean types. |br| 14 | 15 | .. autodata:: RealNumber 16 | 17 | Supported real number types (integer and floating point). 18 | 19 | Mainly distinguished from :data:`Number` by having a total ordering, i.e. 20 | not including the complex numbers. |br| 21 | 22 | .. autodata:: Number 23 | 24 | Supported number types. |br| 25 | 26 | .. autodata:: Scalar 27 | 28 | Supported scalar types. |br| 29 | 30 | .. autodata:: ArithmeticExpression 31 | 32 | A narrower type alias than :class:`~pymbolic.typing.Expression` that is returned 33 | by arithmetic operators, to allow continue doing arithmetic with the result. |br| 34 | 35 | .. currentmodule:: pymbolic.typing 36 | 37 | .. autodata:: Integer 38 | 39 | Supported integer types. |br| 40 | 41 | .. autodata:: Expression 42 | 43 | A union of types that are considered as part of an expression tree. |br| 44 | 45 | .. note:: 46 | 47 | For backward compatibility, ``pymbolic.Expression`` will alias 48 | :class:`pymbolic.primitives.ExpressionNode` for now. Once its 49 | deprecation period is up, it will be removed, and then, in the further 50 | future, ``pymbolic.Expression`` may become this type alias. 51 | 52 | .. autoclass:: ArithmeticOrExpressionT 53 | 54 | .. autodata:: ArithmeticExpressionContainer 55 | 56 | A type alias for allowable expressions that can act as an operand in 57 | arithmetic. This includes numbers and expressions, but also 58 | :class:`~pymbolic.geometric_algebra.MultiVector` and (object) :mod:`numpy` 59 | arrays. |br| 60 | 61 | .. autoclass:: ArithmeticExpressionContainerTc 62 | 63 | .. autofunction:: not_none 64 | """ 65 | 66 | from __future__ import annotations 67 | 68 | 69 | __copyright__ = "Copyright (C) 2024 University of Illinois Board of Trustees" 70 | 71 | __license__ = """ 72 | Permission is hereby granted, free of charge, to any person obtaining a copy 73 | of this software and associated documentation files (the "Software"), to deal 74 | in the Software without restriction, including without limitation the rights 75 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 76 | copies of the Software, and to permit persons to whom the Software is 77 | furnished to do so, subject to the following conditions: 78 | 79 | The above copyright notice and this permission notice shall be included in 80 | all copies or substantial portions of the Software. 81 | 82 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 83 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 84 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 85 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 86 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 87 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 88 | THE SOFTWARE. 89 | """ 90 | 91 | from functools import partial 92 | from typing import TYPE_CHECKING, TypeAlias, TypeVar, Union 93 | 94 | from pytools import T, module_getattr_for_deprecations, obj_array 95 | 96 | 97 | # FIXME: This is a lie. Many more constant types (e.g. numpy and such) 98 | # are in practical use and completely fine. We cannot really add in numpy 99 | # as a special case (because pymbolic doesn't have a hard numpy dependency), 100 | # and there isn't a usable numerical tower that we could rely on. As such, 101 | # code abusing what constants are allowable will have to type-ignore those 102 | # statements. Better ideas would be most welcome. 103 | # 104 | # References: 105 | # https://github.com/python/mypy/issues/3186 106 | # https://discuss.python.org/t/numeric-generics-where-do-we-go-from-pep-3141-and-present-day-mypy/17155/14 107 | 108 | # FIXME: Maybe we should define scalar types via a protocol? 109 | # Pymbolic doesn't particularly restrict scalars to these, this 110 | # just reflects common usage. See typeshed for possible inspiration: 111 | # https://github.com/python/typeshed/blob/119cd09655dcb4ed7fb2021654ba809b8d88846f/stdlib/numbers.pyi 112 | 113 | if TYPE_CHECKING: 114 | from pymbolic import ExpressionNode 115 | from pymbolic.geometric_algebra import MultiVector 116 | 117 | # Experience with depending packages showed that including Decimal and Fraction 118 | # from the stdlib was more trouble than it's worth because those types don't cleanly 119 | # interoperate with each other. 120 | # (e.g. 'Unsupported operand types for * ("Decimal" and "Fraction")') 121 | # And leaving them out doesn't really make any of this more precise. 122 | 123 | _StdlibInexactNumberT = float | complex 124 | 125 | 126 | if TYPE_CHECKING: 127 | # Yes, type-checking pymbolic will require numpy. That's OK. 128 | import numpy as np 129 | 130 | Bool: TypeAlias = bool | np.bool_ 131 | Integer: TypeAlias = int | np.integer 132 | RealNumber: TypeAlias = Integer | float | np.floating 133 | InexactNumber: TypeAlias = _StdlibInexactNumberT | np.inexact 134 | else: 135 | try: 136 | import numpy as np 137 | except ImportError: 138 | Bool = bool 139 | Integer: TypeAlias = int 140 | RealNumber: TypeAlias = Integer | float 141 | InexactNumber: TypeAlias = _StdlibInexactNumberT 142 | else: 143 | Bool = Union[bool, np.bool_] # noqa: UP007 144 | Integer: TypeAlias = Union[int, np.integer] # noqa: UP007 145 | RealNumber: TypeAlias = Union[Integer, float, np.floating] # noqa: UP007 146 | InexactNumber: TypeAlias = _StdlibInexactNumberT | np.inexact 147 | 148 | # NOTE: these are Union because Sphinx seems to understand that better and 149 | # prints a nice "alias of ..." blurb 150 | 151 | Number: TypeAlias = Union[Integer, InexactNumber] # noqa: UP007 152 | Scalar: TypeAlias = Union[Number, Bool] # noqa: UP007 153 | 154 | # NOTE: These need to be Union because they will get used like 155 | # `ArithmeticExpression | None`, which does not work if it's a string. 156 | 157 | ArithmeticExpression: TypeAlias = Union[Number, "ExpressionNode"] 158 | Expression: TypeAlias = Union[Scalar, "ExpressionNode", tuple["Expression", ...]] 159 | 160 | ArithmeticOrExpressionT = TypeVar( 161 | "ArithmeticOrExpressionT", 162 | ArithmeticExpression, 163 | Expression) 164 | """A :class:`~typing.TypeVar` that can be either an 165 | :class:`~pymbolic.ArithmeticExpression` or an :class:`~pymbolic.typing.Expression`. 166 | """ 167 | 168 | ArithmeticExpressionContainer: TypeAlias = Union[ 169 | ArithmeticExpression, 170 | "MultiVector[ArithmeticExpression]", 171 | obj_array.ObjectArray1D[ArithmeticExpression], 172 | obj_array.ObjectArray2D[ArithmeticExpression], 173 | obj_array.ObjectArrayND[ArithmeticExpression]] 174 | 175 | ArithmeticExpressionContainerTc = TypeVar( 176 | "ArithmeticExpressionContainerTc", 177 | ArithmeticExpression, 178 | "MultiVector[ArithmeticExpression]", 179 | obj_array.ObjectArray1D[ArithmeticExpression], 180 | obj_array.ObjectArray2D[ArithmeticExpression], 181 | obj_array.ObjectArrayND[ArithmeticExpression], 182 | ) 183 | """A :class:`~typing.TypeVar` constrained to the types in 184 | :class:`ArithmeticExpressionContainer`. Note that this does not use 185 | :class:`ArithmeticExpressionContainer` as a bound. 186 | """ 187 | 188 | __getattr__ = partial(module_getattr_for_deprecations, __name__, { 189 | "ArithmeticExpressionT": ("ArithmeticExpression", ArithmeticExpression, 2026), 190 | "ExpressionT": ("Expression", Expression, 2026), 191 | "IntegerT": ("Integer", Integer, 2026), 192 | "ScalarT": ("Scalar", Scalar, 2026), 193 | "BoolT": ("Bool", Bool, 2026), 194 | }) 195 | 196 | 197 | def not_none(x: T | None) -> T: 198 | """Backward compatible :func:`operator.is_not_none`.""" 199 | assert x is not None 200 | return x 201 | -------------------------------------------------------------------------------- /pymbolic/imperative/statement.py: -------------------------------------------------------------------------------- 1 | """Instruction types""" 2 | from __future__ import annotations 3 | 4 | 5 | __copyright__ = "Copyright (C) 2015 Matt Wala, Andreas Kloeckner" 6 | 7 | __license__ = """ 8 | Permission is hereby granted, free of charge, to any person obtaining a copy 9 | of this software and associated documentation files (the "Software"), to deal 10 | in the Software without restriction, including without limitation the rights 11 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 12 | copies of the Software, and to permit persons to whom the Software is 13 | furnished to do so, subject to the following conditions: 14 | 15 | The above copyright notice and this permission notice shall be included in 16 | all copies or substantial portions of the Software. 17 | 18 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 19 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 20 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 21 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 22 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 23 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 24 | THE SOFTWARE. 25 | """ 26 | 27 | from dataclasses import dataclass, replace 28 | from sys import intern 29 | from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar, cast 30 | 31 | from typing_extensions import Self, override 32 | 33 | from pymbolic.typing import Expression, not_none 34 | 35 | 36 | if TYPE_CHECKING: 37 | from collections.abc import Callable, Set 38 | 39 | from pymbolic.primitives import Variable 40 | 41 | 42 | # {{{ statement classes 43 | 44 | class BasicStatementLike(Protocol): 45 | @property 46 | def id(self) -> str: ... 47 | @property 48 | def depends_on(self) -> Set[str]: ... 49 | 50 | def copy(self, **kwargs: object) -> Self: ... 51 | 52 | 53 | BasicStatementLikeT = TypeVar("BasicStatementLikeT", bound=BasicStatementLike) 54 | 55 | 56 | class StatementLike(BasicStatementLike, Protocol): 57 | def get_written_variables(self) -> Set[str]: ... 58 | 59 | def get_read_variables(self) -> Set[str]: ... 60 | 61 | def map_expressions(self, 62 | mapper: Callable[[Expression], Expression], 63 | include_lhs: bool = True 64 | ) -> Self: ... 65 | 66 | 67 | StatementLikeT = TypeVar("StatementLikeT", bound=StatementLike) 68 | 69 | 70 | @dataclass(frozen=True) 71 | class Statement: 72 | """ 73 | .. autoattribute:: depends_on 74 | .. autoattribute:: id 75 | 76 | .. automethod:: get_written_variables 77 | .. automethod:: get_read_variables 78 | """ 79 | id: str 80 | """ 81 | A string, a unique identifier for this instruction. 82 | """ 83 | 84 | depends_on: Set[str] 85 | """A :class:`frozenset` of instruction ids that are required to be 86 | executed within this execution context before this instruction can be 87 | executed.""" 88 | 89 | def __post_init__(self): 90 | if isinstance(self.id, str): 91 | object.__setattr__(self, "id", intern(self.id)) 92 | 93 | def get_written_variables(self) -> Set[str]: 94 | """Returns a :class:`frozenset` of variables being written by this 95 | instruction. 96 | """ 97 | return frozenset() 98 | 99 | def get_read_variables(self) -> Set[str]: 100 | """Returns a :class:`frozenset` of variables being read by this 101 | instruction. 102 | """ 103 | return frozenset() 104 | 105 | def map_expressions(self, 106 | mapper: Callable[[Expression], Expression], 107 | include_lhs: bool = True 108 | ) -> Self: 109 | """Returns a new copy of *self* with all expressions 110 | replaced by ``mapepr(expr)`` for every 111 | :class:`pymbolic.primitives.Expression` 112 | contained in *self*. 113 | """ 114 | return self 115 | 116 | def get_dependency_mapper(self, 117 | include_calls: bool | Literal["descend_args"] = True, 118 | ): 119 | from pymbolic.mapper.dependency import DependencyMapper 120 | return DependencyMapper[[]]( 121 | include_subscripts=False, 122 | include_lookups=False, 123 | include_calls=include_calls) 124 | 125 | def copy(self, **kwargs: Any) -> Self: # pyright: ignore[reportAny] 126 | return replace(self, **kwargs) 127 | 128 | # }}} 129 | 130 | 131 | # {{{ statement with condition 132 | 133 | @dataclass(frozen=True) 134 | class ConditionalStatement(Statement): 135 | __doc__ = not_none(Statement.__doc__) + """ 136 | .. autoattribute:: condition 137 | """ 138 | 139 | condition: Expression 140 | """The instruction condition as a :mod:`pymbolic` expression (`True` if the 141 | instruction is unconditionally executed)""" 142 | 143 | def _condition_printing_suffix(self): 144 | if self.condition is True: 145 | return "" 146 | return " if " + str(self.condition) 147 | 148 | @override 149 | def __str__(self): 150 | return (super().__str__() 151 | + self._condition_printing_suffix()) 152 | 153 | @override 154 | def get_read_variables(self) -> Set[str]: 155 | dep_mapper = self.get_dependency_mapper() 156 | return ( 157 | super().get_read_variables() 158 | | frozenset( 159 | cast("Variable", dep).name for dep in dep_mapper(self.condition))) 160 | 161 | # }}} 162 | 163 | 164 | # {{{ assignment 165 | 166 | @dataclass(frozen=True) 167 | class Assignment(Statement): 168 | """ 169 | .. attribute:: lhs 170 | .. attribute:: rhs 171 | """ 172 | 173 | lhs: Expression 174 | rhs: Expression 175 | 176 | @override 177 | def get_written_variables(self): 178 | from pymbolic.primitives import Subscript, Variable 179 | if isinstance(self.lhs, Variable): 180 | return frozenset([self.lhs.name]) 181 | elif isinstance(self.lhs, Subscript): 182 | assert isinstance(self.lhs.aggregate, Variable) 183 | return frozenset([self.lhs.aggregate.name]) 184 | else: 185 | raise TypeError("unexpected type of LHS") 186 | 187 | @override 188 | def get_read_variables(self) -> Set[str]: 189 | result = super().get_read_variables() 190 | get_deps = self.get_dependency_mapper() 191 | 192 | def get_vars(expr: Expression): 193 | return frozenset(cast("Variable", dep).name for dep in get_deps(expr)) 194 | 195 | result = get_vars(self.rhs) | get_vars(self.lhs) 196 | 197 | return result 198 | 199 | @override 200 | def map_expressions(self, 201 | mapper: Callable[[Expression], Expression], 202 | include_lhs: bool = True 203 | ) -> Self: 204 | return (super() 205 | .map_expressions(mapper, include_lhs=include_lhs) 206 | .copy( 207 | lhs=mapper(self.lhs) if include_lhs else self.lhs, 208 | rhs=mapper(self.rhs))) 209 | 210 | @override 211 | def __str__(self): 212 | result = "{assignee} <- {expr}".format( 213 | assignee=str(self.lhs), 214 | expr=str(self.rhs),) 215 | 216 | return result 217 | 218 | # }}} 219 | 220 | 221 | # {{{ conditional assignment 222 | 223 | @dataclass(frozen=True) 224 | class ConditionalAssignment(ConditionalStatement, Assignment): 225 | @override 226 | def map_expressions(self, 227 | mapper: Callable[[Expression], Expression], 228 | include_lhs: bool = True 229 | ) -> Self: 230 | return (super() 231 | .map_expressions(mapper, include_lhs=include_lhs) 232 | .copy(condition=mapper(self.condition))) 233 | 234 | # }}} 235 | 236 | 237 | # {{{ nop 238 | 239 | @dataclass(frozen=True) 240 | class Nop(Statement): 241 | def __str__(self): 242 | return "nop" 243 | 244 | # }}} 245 | 246 | 247 | Instruction = Statement 248 | ConditionalInstruction = ConditionalStatement 249 | 250 | 251 | # vim: foldmethod=marker 252 | -------------------------------------------------------------------------------- /pymbolic/interop/matchpy/tofrom.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from dataclasses import dataclass 4 | from typing import TYPE_CHECKING 5 | 6 | import multiset 7 | import numpy as np 8 | from matchpy import Expression as MatchpyExpression 9 | from typing_extensions import override 10 | 11 | import pymbolic.interop.matchpy as m 12 | import pymbolic.primitives as p 13 | from pymbolic.interop.matchpy.mapper import Mapper as BaseMatchPyMapper 14 | from pymbolic.mapper import Mapper as BasePymMapper 15 | 16 | 17 | if TYPE_CHECKING: 18 | from collections.abc import Callable 19 | 20 | from pymbolic.typing import Expression as PbExpression, Scalar as PbScalar 21 | 22 | 23 | # {{{ to matchpy 24 | 25 | class ToMatchpyExpressionMapper(BasePymMapper[MatchpyExpression, []]): 26 | """ 27 | Mapper to convert instances of :class:`pymbolic.primitives.Expression` to 28 | :class:`pymbolic.interop.matchpy.PymbolicOperation`. 29 | """ 30 | @override 31 | def map_constant(self, expr: object) -> m.Scalar: 32 | if np.isscalar(expr): 33 | assert p.is_constant(expr) 34 | return m.Scalar(expr) # pyright: ignore[reportUnknownArgumentType] 35 | 36 | raise NotImplementedError(expr) 37 | 38 | @override 39 | def map_variable(self, expr: p.Variable) -> m.Variable: 40 | return m.Variable(m.Id(expr.name)) 41 | 42 | @override 43 | def map_call(self, expr: p.Call) -> m.Call: 44 | return m.Call(self.rec(expr.function), 45 | m.TupleOp(tuple(self.rec(p) 46 | for p in expr.parameters))) 47 | 48 | @override 49 | def map_subscript(self, expr: p.Subscript) -> m.Subscript: 50 | return m.Subscript(self.rec(expr.aggregate), 51 | m.TupleOp(tuple(self.rec(idx) 52 | for idx in expr.index_tuple))) 53 | 54 | @override 55 | def map_sum(self, expr: p.Sum) -> m.Sum: 56 | return m.Sum(*[self.rec(child) 57 | for child in expr.children]) 58 | 59 | @override 60 | def map_product(self, expr: p.Product) -> m.Product: 61 | return m.Product(*[self.rec(child) 62 | for child in expr.children]) 63 | 64 | @override 65 | def map_quotient(self, expr: p.Quotient) -> m.TrueDiv: 66 | return m.TrueDiv(self.rec(expr.numerator), self.rec(expr.denominator)) 67 | 68 | @override 69 | def map_floor_div(self, expr: p.FloorDiv) -> m.FloorDiv: 70 | return m.FloorDiv(self.rec(expr.numerator), self.rec(expr.denominator)) 71 | 72 | @override 73 | def map_remainder(self, expr: p.Remainder) -> m.Modulo: 74 | return m.Modulo(self.rec(expr.numerator), self.rec(expr.denominator)) 75 | 76 | @override 77 | def map_power(self, expr: p.Power) -> m.Power: 78 | return m.Power(self.rec(expr.base), self.rec(expr.exponent)) 79 | 80 | @override 81 | def map_left_shift(self, expr: p.LeftShift) -> m.LeftShift: 82 | return m.LeftShift(self.rec(expr.shiftee), self.rec(expr.shift)) 83 | 84 | @override 85 | def map_right_shift(self, expr: p.RightShift) -> m.RightShift: 86 | return m.RightShift(self.rec(expr.shiftee), self.rec(expr.shift)) 87 | 88 | @override 89 | def map_bitwise_not(self, expr: p.BitwiseNot) -> m.BitwiseNot: 90 | return m.BitwiseNot(self.rec(expr.child)) 91 | 92 | @override 93 | def map_bitwise_or(self, expr: p.BitwiseOr) -> m.BitwiseOr: 94 | return m.BitwiseOr(*[self.rec(child) 95 | for child in expr.children]) 96 | 97 | @override 98 | def map_bitwise_and(self, expr: p.BitwiseAnd) -> m.BitwiseAnd: 99 | return m.BitwiseAnd(*[self.rec(child) 100 | for child in expr.children]) 101 | 102 | @override 103 | def map_bitwise_xor(self, expr: p.BitwiseXor) -> m.BitwiseXor: 104 | return m.BitwiseXor(*[self.rec(child) 105 | for child in expr.children]) 106 | 107 | @override 108 | def map_logical_not(self, expr: p.LogicalNot) -> m.LogicalNot: 109 | return m.LogicalNot(self.rec(expr.child)) 110 | 111 | @override 112 | def map_logical_or(self, expr: p.LogicalOr) -> m.LogicalOr: 113 | return m.LogicalOr(*[self.rec(child) 114 | for child in expr.children]) 115 | 116 | @override 117 | def map_logical_and(self, expr: p.LogicalAnd) -> m.LogicalAnd: 118 | return m.LogicalAnd(*[self.rec(child) 119 | for child in expr.children]) 120 | 121 | @override 122 | def map_comparison(self, expr: p.Comparison) -> m.Comparison: 123 | return m.Comparison(self.rec(expr.left), 124 | m.ComparisonOp(expr.operator), 125 | self.rec(expr.right), 126 | ) 127 | 128 | @override 129 | def map_if(self, expr: p.If) -> m.If: 130 | return m.If(self.rec(expr.condition), 131 | self.rec(expr.then), 132 | self.rec(expr.else_)) 133 | 134 | @override 135 | def map_dot_wildcard(self, expr: p.DotWildcard) -> m.Wildcard: 136 | return m.Wildcard.dot(expr.name) 137 | 138 | @override 139 | def map_star_wildcard(self, expr: p.StarWildcard) -> m.Wildcard: 140 | return m.Wildcard.star(expr.name) 141 | 142 | # }}} 143 | 144 | 145 | # {{{ from matchpy 146 | 147 | class FromMatchpyExpressionMapper(BaseMatchPyMapper): 148 | def map_scalar(self, expr: m.Scalar) -> PbScalar: 149 | return expr.value 150 | 151 | def map_variable(self, expr: m.Variable) -> p.Variable: 152 | return p.Variable(expr.id.value) 153 | 154 | def map_call(self, expr: m.Call) -> p.Call: 155 | return p.Call(self.rec(expr.function), 156 | tuple(self.rec(arg) 157 | for arg in expr.args._operands)) 158 | 159 | def map_subscript(self, expr: m.Subscript) -> p.Subscript: 160 | return p.Subscript(self.rec(expr.aggregate), 161 | tuple(self.rec(idx) 162 | for idx in expr.indices)) 163 | 164 | def map_true_div(self, expr: m.TrueDiv) -> p.Quotient: 165 | return p.Quotient(self.rec(expr.x1), self.rec(expr.x2)) 166 | 167 | def map_floor_div(self, expr: m.FloorDiv) -> p.FloorDiv: 168 | return p.FloorDiv(self.rec(expr.x1), self.rec(expr.x2)) 169 | 170 | def map_modulo(self, expr: m.Modulo) -> p.Remainder: 171 | return p.Remainder(self.rec(expr.x1), self.rec(expr.x2)) 172 | 173 | def map_power(self, expr: m.Power) -> p.Power: 174 | return p.Power(self.rec(expr.x1), self.rec(expr.x2)) 175 | 176 | def map_left_shift(self, expr: m.LeftShift) -> p.LeftShift: 177 | return p.LeftShift(self.rec(expr.x1), self.rec(expr.x2)) 178 | 179 | def map_right_shift(self, expr: m.RightShift) -> p.RightShift: 180 | return p.RightShift(self.rec(expr.x1), self.rec(expr.x2)) 181 | 182 | def map_sum(self, expr: m.Sum) -> p.Sum: 183 | return p.Sum(tuple(self.rec(child) 184 | for child in expr.children)) 185 | 186 | def map_product(self, expr: m.Product) -> p.Product: 187 | return p.Product(tuple(self.rec(operand) 188 | for operand in expr.operands)) 189 | 190 | def map_logical_or(self, expr: m.LogicalOr) -> p.LogicalOr: 191 | return p.LogicalOr(tuple(self.rec(operand) 192 | for operand in expr.operands)) 193 | 194 | def map_logical_and(self, expr: m.LogicalAnd) -> p.LogicalAnd: 195 | return p.LogicalAnd(tuple(self.rec(operand) 196 | for operand in expr.operands)) 197 | 198 | def map_bitwise_or(self, expr: m.BitwiseOr) -> p.BitwiseOr: 199 | return p.BitwiseOr(tuple(self.rec(operand) 200 | for operand in expr.operands)) 201 | 202 | def map_bitwise_and(self, expr: m.BitwiseAnd) -> p.BitwiseAnd: 203 | return p.BitwiseAnd(tuple(self.rec(operand) 204 | for operand in expr.operands)) 205 | 206 | def map_bitwise_xor(self, expr: m.BitwiseXor) -> p.BitwiseXor: 207 | return p.BitwiseXor(tuple(self.rec(operand) 208 | for operand in expr.operands)) 209 | 210 | def map_logical_not(self, expr: m.LogicalNot) -> p.LogicalNot: 211 | return p.LogicalNot(self.rec(expr.x)) 212 | 213 | def map_bitwise_not(self, expr: m.BitwiseNot) -> p.BitwiseNot: 214 | return p.BitwiseNot(self.rec(expr.x)) 215 | 216 | def map_comparison(self, expr: m.Comparison) -> p.Comparison: 217 | return p.Comparison(self.rec(expr.left), 218 | expr.operator.value, 219 | self.rec(expr.right)) 220 | 221 | def map_if(self, expr: m.If) -> p.If: 222 | return p.If(self.rec(expr.condition), 223 | self.rec(expr.then), 224 | self.rec(expr.else_)) 225 | 226 | # }}} 227 | 228 | 229 | @dataclass(frozen=True, eq=True) 230 | class ToFromReplacement: 231 | f: Callable[..., PbExpression] 232 | to_matchpy_expr: m.ToMatchpyT 233 | from_matchpy_expr: m.FromMatchpyT 234 | 235 | def __call__(self, **kwargs): 236 | kwargs_to_f = {} 237 | 238 | for kw, arg in kwargs.items(): 239 | if isinstance(arg, MatchpyExpression): 240 | arg = self.from_matchpy_expr(arg) 241 | elif isinstance(arg, multiset.Multiset): 242 | arg = multiset.Multiset({self.from_matchpy_expr(expr): count 243 | for expr, count in arg.items()}) 244 | elif isinstance(arg, tuple): 245 | arg = tuple(self.from_matchpy_expr(el) for el in arg) 246 | else: 247 | raise NotImplementedError(f"Cannot convert back {type(arg)}") 248 | 249 | kwargs_to_f[kw] = arg 250 | 251 | return self.to_matchpy_expr(self.f(**kwargs_to_f)) 252 | 253 | # vim: foldmethod=marker 254 | -------------------------------------------------------------------------------- /pymbolic/mapper/c_code.py: -------------------------------------------------------------------------------- 1 | """ 2 | .. autoclass:: CCodeMapper 3 | :show-inheritance: 4 | """ 5 | from __future__ import annotations 6 | 7 | 8 | __copyright__ = "Copyright (C) 2009-2013 Andreas Kloeckner" 9 | 10 | __license__ = """ 11 | Permission is hereby granted, free of charge, to any person obtaining a copy 12 | of this software and associated documentation files (the "Software"), to deal 13 | in the Software without restriction, including without limitation the rights 14 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 15 | copies of the Software, and to permit persons to whom the Software is 16 | furnished to do so, subject to the following conditions: 17 | 18 | The above copyright notice and this permission notice shall be included in 19 | all copies or substantial portions of the Software. 20 | 21 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 22 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 23 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 24 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 25 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 26 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 27 | THE SOFTWARE. 28 | """ 29 | 30 | from typing import TYPE_CHECKING 31 | 32 | from typing_extensions import override 33 | 34 | import pymbolic.primitives as p 35 | from pymbolic.mapper.stringifier import ( 36 | PREC_CALL, 37 | PREC_LOGICAL_AND, 38 | PREC_LOGICAL_OR, 39 | PREC_NONE, 40 | PREC_UNARY, 41 | SimplifyingSortingStringifyMapper, 42 | ) 43 | 44 | 45 | if TYPE_CHECKING: 46 | from collections.abc import Iterator, Sequence 47 | 48 | from pymbolic.typing import Expression 49 | 50 | 51 | class CCodeMapper(SimplifyingSortingStringifyMapper[[]]): 52 | """Generate C code for expressions, while extracting 53 | :class:`pymbolic.primitives.CommonSubexpression` instances. 54 | 55 | As an example, define a fairly simple expression *expr*: 56 | 57 | .. doctest:: 58 | 59 | >>> import pymbolic.primitives as p 60 | >>> CSE = p.make_common_subexpression 61 | >>> x = p.Variable("x") 62 | >>> u = CSE(3*x**2-5, "u") 63 | >>> expr = u/(u+3)*(u+5) 64 | >>> print(expr) 65 | (CSE(3*x**2 + -5) / (CSE(3*x**2 + -5) + 3))*(CSE(3*x**2 + -5) + 5) 66 | 67 | Notice that if we were to directly generate this code without the added 68 | `CSE`, the subexpression *u* would be evaluated multiple times. Wrapping the 69 | expression as above avoids this unnecessary cost. 70 | 71 | .. doctest:: 72 | 73 | >>> from pymbolic.mapper.c_code import CCodeMapper as CCM 74 | >>> ccm = CCM() 75 | >>> result = ccm(expr) 76 | 77 | >>> for name, value in ccm.cse_name_list: 78 | ... print("%s = %s;" % (name, value)) 79 | ... 80 | _cse_u = 3 * x * x + -5; 81 | >>> print(result) 82 | _cse_u / (_cse_u + 3) * (_cse_u + 5) 83 | 84 | See :class:`pymbolic.mapper.stringifier.CSESplittingStringifyMapperMixin` 85 | for the ``cse_*`` attributes. 86 | """ 87 | 88 | cse_prefix: str 89 | cse_to_name: dict[Expression, str] 90 | cse_names: set[str] 91 | cse_name_list: list[tuple[str, str]] 92 | complex_constant_base_type: str 93 | 94 | def __init__(self, 95 | reverse: bool = True, 96 | cse_prefix: str = "_cse", 97 | complex_constant_base_type: str = "double", 98 | cse_name_list: Sequence[tuple[str, str]] | None = None) -> None: 99 | if cse_name_list is None: 100 | cse_name_list = [] 101 | 102 | super().__init__(reverse) 103 | self.cse_prefix = cse_prefix 104 | 105 | self.cse_to_name = {cse: name for name, cse in cse_name_list} 106 | self.cse_names = {cse for _name, cse in cse_name_list} 107 | self.cse_name_list = list(cse_name_list) 108 | 109 | self.complex_constant_base_type = complex_constant_base_type 110 | 111 | def copy(self, 112 | cse_name_list: Sequence[tuple[str, str]] | None = None, 113 | ) -> CCodeMapper: 114 | if cse_name_list is None: 115 | cse_name_list = self.cse_name_list 116 | 117 | return CCodeMapper(self.reverse, 118 | self.cse_prefix, self.complex_constant_base_type, 119 | cse_name_list) 120 | 121 | def copy_with_mapped_cses( 122 | self, cses_and_values: Sequence[tuple[str, str]]) -> CCodeMapper: 123 | return self.copy((*self.cse_name_list, *cses_and_values)) 124 | 125 | # {{{ mappings 126 | 127 | @override 128 | def map_product(self, expr: p.Product, /, enclosing_prec: int) -> str: 129 | from pymbolic.mapper.stringifier import PREC_PRODUCT 130 | return self.parenthesize_if_needed( 131 | # Spaces prevent '**z' (times dereference z), which 132 | # is hard to read. 133 | 134 | self.join_rec(" * ", expr.children, PREC_PRODUCT), 135 | enclosing_prec, PREC_PRODUCT) 136 | 137 | @override 138 | def map_constant(self, x: object, /, enclosing_prec: int) -> str: 139 | if isinstance(x, complex): 140 | base = self.complex_constant_base_type 141 | real = self.map_constant(x.real, PREC_NONE) 142 | imag = self.map_constant(x.imag, PREC_NONE) 143 | 144 | # NOTE: this will need the include to work 145 | # FIXME: MSVC does not support , so this will not work. 146 | # (AFAIK, it uses a struct instead and does not support arithmetic) 147 | return f"({base} complex)({real} + {imag} * _Imaginary_I)" 148 | else: 149 | return super().map_constant(x, enclosing_prec) 150 | 151 | @override 152 | def map_call(self, expr: p.Call, /, enclosing_prec: int) -> str: 153 | if isinstance(expr.function, p.Variable): 154 | func = expr.function.name 155 | else: 156 | func = self.rec(expr.function, PREC_CALL) 157 | 158 | return self.format("%s(%s)", 159 | func, 160 | self.join_rec(", ", expr.parameters, PREC_NONE)) 161 | 162 | @override 163 | def map_power(self, expr: p.Power, /, enclosing_prec: int) -> str: 164 | from pymbolic.mapper.stringifier import PREC_NONE 165 | from pymbolic.primitives import is_constant, is_zero 166 | if is_constant(expr.exponent): 167 | if is_zero(expr.exponent): 168 | return "1" 169 | elif is_zero(expr.exponent - 1): 170 | return self.rec(expr.base, enclosing_prec) 171 | elif is_zero(expr.exponent - 2): 172 | return self.rec(expr.base*expr.base, enclosing_prec) 173 | 174 | return self.format("pow(%s, %s)", 175 | self.rec(expr.base, PREC_NONE), 176 | self.rec(expr.exponent, PREC_NONE)) 177 | 178 | @override 179 | def map_floor_div(self, expr: p.FloorDiv, /, enclosing_prec: int) -> str: 180 | # Let's see how bad of an idea this is--sane people would only 181 | # apply this to integers, right? 182 | 183 | from pymbolic.mapper.stringifier import PREC_POWER, PREC_PRODUCT 184 | return self.format("(%s/%s)", 185 | self.rec(expr.numerator, PREC_PRODUCT), 186 | self.rec(expr.denominator, PREC_POWER)) # analogous to ^{-1} 187 | 188 | @override 189 | def map_logical_not(self, expr: p.LogicalNot, /, enclosing_prec: int) -> str: 190 | child = self.rec(expr.child, PREC_UNARY) 191 | return self.parenthesize_if_needed(f"!{child}", enclosing_prec, PREC_UNARY) 192 | 193 | @override 194 | def map_logical_and(self, expr: p.LogicalAnd, /, enclosing_prec: int) -> str: 195 | return self.parenthesize_if_needed( 196 | self.join_rec(" && ", expr.children, PREC_LOGICAL_AND), 197 | enclosing_prec, PREC_LOGICAL_AND) 198 | 199 | @override 200 | def map_logical_or(self, expr: p.LogicalOr, /, enclosing_prec: int) -> str: 201 | return self.parenthesize_if_needed( 202 | self.join_rec(" || ", expr.children, PREC_LOGICAL_OR), 203 | enclosing_prec, PREC_LOGICAL_OR) 204 | 205 | @override 206 | def map_common_subexpression( 207 | self, expr: p.CommonSubexpression, /, enclosing_prec: int) -> str: 208 | try: 209 | cse_name = self.cse_to_name[expr.child] 210 | except KeyError: 211 | from pymbolic.mapper.stringifier import PREC_NONE 212 | cse_str = self.rec(expr.child, PREC_NONE) 213 | 214 | if expr.prefix is not None: 215 | def generate_cse_names() -> Iterator[str]: 216 | yield f"{self.cse_prefix}_{expr.prefix}" 217 | i = 2 218 | while True: 219 | yield f"{self.cse_prefix}_{expr.prefix}_{i}" 220 | i += 1 221 | else: 222 | def generate_cse_names() -> Iterator[str]: 223 | i = 0 224 | while True: 225 | yield f"{self.cse_prefix}{i}" 226 | i += 1 227 | 228 | cse_name = None 229 | for cse_name in generate_cse_names(): 230 | if cse_name not in self.cse_names: 231 | break 232 | assert cse_name is not None 233 | 234 | self.cse_name_list.append((cse_name, cse_str)) 235 | self.cse_to_name[expr.child] = cse_name 236 | self.cse_names.add(cse_name) 237 | 238 | assert len(self.cse_names) == len(self.cse_to_name) 239 | 240 | return cse_name 241 | 242 | @override 243 | def map_if(self, expr: p.If, /, enclosing_prec: int) -> str: 244 | from pymbolic.mapper.stringifier import PREC_NONE 245 | return self.format("(%s ? %s : %s)", 246 | self.rec(expr.condition, PREC_NONE), 247 | self.rec(expr.then, PREC_NONE), 248 | self.rec(expr.else_, PREC_NONE), 249 | ) 250 | 251 | # }}} 252 | 253 | # vim: foldmethod=marker 254 | --------------------------------------------------------------------------------