re-design cue/cuex API + DCE rule #87

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

mariogeiger merged 110 commits into main from polynomial

Mar 16, 2025

.github/workflows/coverage.yml

This file was deleted.

.github/workflows/tests.yml

-Original file line number
+Diff line change
@@ Expand Up / @@ -27,18 +27,22 @@ jobs: @@
             python -m pip install --upgrade pip
             python -m pip install --upgrade uv
             python -m uv pip install pytest
-            python -m uv pip install ./cuequivariance
+            python -m uv pip install -e ./cuequivariance
         - name: Test with pytest
           run: |
             pytest --doctest-modules -x cuequivariance
       cuequivariance-jax:
-        runs-on: ubuntu-latest
         strategy:
           fail-fast: false
           matrix:
-            python-version: ["3.10", "3.12"]
+            include:
+              - runner: "ubuntu-latest"
+                python-version: "3.10"
+              - runner: "self-hosted"
+                python-version: "3.12"
+        runs-on: ${{ matrix.runner }}
         steps:
         - uses: actions/checkout@v4
@@ Expand All / @@ -50,9 +54,12 @@ jobs: @@
           run: |
             python -m pip install --upgrade pip
             python -m pip install --upgrade uv
-            python -m uv pip install pytest
-            python -m uv pip install ./cuequivariance
-            python -m uv pip install ./cuequivariance_jax
+            python -m uv pip install pytest "jax[cuda12]"
+            python -m uv pip uninstall cuequivariance cuequivariance_jax cuequivariance_torch
+            python -m uv pip install -e ./cuequivariance
+            python -m uv pip install --no-deps -e ./cuequivariance_jax
+            python -c "import cuequivariance; print('cue', cuequivariance.__version__)"
+            python -c "import cuequivariance_jax; print('cuex', cuequivariance_jax.__version__)"
         - name: Test with pytest
           run: |
             pytest --doctest-modules -x cuequivariance_jax
@@ Expand Down Expand Up / @@ -80,8 +87,11 @@ jobs: @@
             python -m pip install --upgrade pip
             python -m pip install --upgrade uv
             python -m uv pip install pytest torch e3nn
-            python -m uv pip install ./cuequivariance --force-reinstall
-            python -m uv pip install ./cuequivariance_torch --force-reinstall
+            python -m uv pip uninstall cuequivariance cuequivariance_jax cuequivariance_torch
+            python -m uv pip install -e ./cuequivariance
+            python -m uv pip install --no-deps -e ./cuequivariance_torch
+            python -c "import cuequivariance; print('cue', cuequivariance.__version__)"
+            python -c "import cuequivariance_torch; print('cuet', cuequivariance_torch.__version__)"
         - name: Test with pytest
           run: |
             pytest --doctest-modules -x cuequivariance_torch

.pre-commit-config.yaml

-Original file line number
+Diff line change
@@ -1,6 +1,6 @@
     repos:
       - repo: https://github.com/astral-sh/ruff-pre-commit
-        rev: v0.9.3
+        rev: v0.9.9
         hooks:
           - id: ruff
             args: ["--extend-select", "I", "--fix"]
@@ Expand Down @@

CHANGELOG.md

-Original file line number
+Diff line change
@@ -1,5 +1,18 @@
     ## Latest Changes
+    ### Breaking Changes
+    - Rename `SegmentedTensorProduct.flop_cost` in `flops`
+    - Rename `SegmentedTensorProduct.memory_cost` in `memory`
+    - Removed `IrrepsArray` in favor of `RepArray`
+    - Change folder structure of cuequivariance and cuequivariance-jax. Now the main subfolders are `segmented_polynomials` and `group_theory`
+    - Deprecate `cue.EquivariantTensorProduct` in favor of `cue.EquivariantPolynomial`
+    - The descriptors return `cue.EquivariantPolynomial` instead of `cue.EquivariantTensorProduct`
+    # Added
+    - Class `cue.SegmentedOperand`, `cue.SegmentedPolynomial`
+    - Class `cue.EquivariantPolynomial` that contains a `cue.SegmentedPolynomial` and the `cue.Rep` of its inputs and outputs
     ## 0.3.0 (2025-03-05)
     The main changes are:
@@ Expand Down @@

cuequivariance/cuequivariance/__init__.py

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -18,7 +18,14 @@
  
        importlib.resources.files(__package__).joinpath("VERSION").read_text().strip()

    )

    from cuequivariance.representation import (

    from .segmented_polynomials import (

        Operation,

        SegmentedOperand,

        SegmentedTensorProduct,

        SegmentedPolynomial,

    )

    from .group_theory import (

        Rep,

        Irrep,

        clebsch_gordan,

    @@ -27,9 +34,6 @@
  
        SU2,

        SO3,

        O3,

    )

    from cuequivariance.irreps_array import (

        get_irrep_scope,

        MulIrrep,

        Irreps,

    @@ -45,18 +49,20 @@
  
        reduced_tensor_product_basis,

        reduced_symmetric_tensor_product_basis,

        reduced_antisymmetric_tensor_product_basis,

        EquivariantPolynomial,

        EquivariantTensorProduct,  # deprecated

    )

    from cuequivariance.segmented_tensor_product import SegmentedTensorProduct

    from cuequivariance.equivariant_tensor_product import EquivariantTensorProduct

    from cuequivariance.operation import Operation

    from cuequivariance import (

        segmented_tensor_product,

        descriptors,

    )

    from cuequivariance import segmented_polynomials as segmented_polynomials

    from cuequivariance import group_theory as group_theory

    from cuequivariance.group_theory import descriptors as descriptors

    __all__ = [

        "__version__",

        "Operation",

        "SegmentedOperand",

        "SegmentedTensorProduct",

        "SegmentedPolynomial",

        "Rep",

        "Irrep",

        "clebsch_gordan",

    @@ -80,9 +86,9 @@
  
        "reduced_tensor_product_basis",

        "reduced_symmetric_tensor_product_basis",

        "reduced_antisymmetric_tensor_product_basis",

        "SegmentedTensorProduct",

        "EquivariantPolynomial",

        "EquivariantTensorProduct",

        "Operation",

        "segmented_tensor_product",

        "segmented_polynomials",

        "group_theory",

        "descriptors",

    ]

0 cuequivariance/cuequivariance/misc/linalg.py → cuequivariance/cuequivariance/etc/linalg.py

File renamed without changes.

cuequivariance/cuequivariance/etc/permutations.py

-Original file line number
+Diff line change
@@ -0,0 +1,44 @@
+    # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+    # SPDX-License-Identifier: Apache-2.0
+    #
+    # Licensed under the Apache License, Version 2.0 (the "License");
+    # you may not use this file except in compliance with the License.
+    # You may obtain a copy of the License at
+    #
+    # http://www.apache.org/licenses/LICENSE-2.0
+    #
+    # Unless required by applicable law or agreed to in writing, software
+    # distributed under the License is distributed on an "AS IS" BASIS,
+    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    # See the License for the specific language governing permissions and
+    # limitations under the License.
+    from typing import Sequence
+    def compose_permutations(p1: tuple[int, ...], p2: tuple[int, ...]) -> tuple[int, ...]:
+        """Compose two permutations"""
+        return tuple(p1[p2[i]] for i in range(len(p1)))
+    def inverse_permutation(p: tuple[int, ...]) -> tuple[int, ...]:
+        """Inverse a permutation"""
+        return tuple(p.index(i) for i in range(len(p)))
+    def generate_permutations_from(
+        generators: Sequence[tuple[int, ...]],
+    ) -> set[tuple[int, ...]]:
+        """Generate all permutations from a list of generators"""
+        result = set(generators)
+        while True:
+            n = len(result)
+            new_result = result.copy()
+            for g in result:
+                for h in result:
+                    new_result.add(compose_permutations(g, h))
+            if len(new_result) == n:
+                break
+            result = new_result
+        return result

0 ...riance/cuequivariance/misc/sympy_utils.py → ...ariance/cuequivariance/etc/sympy_utils.py

File renamed without changes.

cuequivariance/cuequivariance/group_theory/__init__.py

-Original file line number
+Diff line change
@@ -0,0 +1,75 @@
+    # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+    # SPDX-License-Identifier: Apache-2.0
+    #
+    # Licensed under the Apache License, Version 2.0 (the "License");
+    # you may not use this file except in compliance with the License.
+    # You may obtain a copy of the License at
+    #
+    # http://www.apache.org/licenses/LICENSE-2.0
+    #
+    # Unless required by applicable law or agreed to in writing, software
+    # distributed under the License is distributed on an "AS IS" BASIS,
+    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    # See the License for the specific language governing permissions and
+    # limitations under the License.
+    from .representations import (
+        Rep,
+        Irrep,
+        clebsch_gordan,
+        selection_rule_product,
+        selection_rule_power,
+        SU2,
+        SO3,
+        O3,
+    )
+    from .irreps_array import (
+        get_irrep_scope,
+        MulIrrep,
+        Irreps,
+        IrrepsLayout,
+        mul_ir,
+        ir_mul,
+        IrrepsAndLayout,
+        get_layout_scope,
+        assume,
+        NumpyIrrepsArray,
+        from_segments,
+        concatenate,
+        reduced_tensor_product_basis,
+        reduced_symmetric_tensor_product_basis,
+        reduced_antisymmetric_tensor_product_basis,
+    )
+    from .equivariant_polynomial import EquivariantPolynomial
+    from .equivariant_tensor_product import EquivariantTensorProduct
+    __all__ = [
+        "Rep",
+        "Irrep",
+        "clebsch_gordan",
+        "selection_rule_product",
+        "selection_rule_power",
+        "SU2",
+        "SO3",
+        "O3",
+        "get_irrep_scope",
+        "MulIrrep",
+        "Irreps",
+        "IrrepsLayout",
+        "mul_ir",
+        "ir_mul",
+        "IrrepsAndLayout",
+        "get_layout_scope",
+        "assume",
+        "NumpyIrrepsArray",
+        "from_segments",
+        "concatenate",
+        "reduced_tensor_product_basis",
+        "reduced_symmetric_tensor_product_basis",
+        "reduced_antisymmetric_tensor_product_basis",
+        "EquivariantPolynomial",
+        "EquivariantTensorProduct",
+    ]

0 ...ce/cuequivariance/descriptors/__init__.py → ...ance/group_theory/descriptors/__init__.py

File renamed without changes.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

re-design cue/cuex API + DCE rule #87

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!