add Irreps.is_scalar

e3nn · Sep 27, 2022 · b631e3e · b631e3e
1 parent 84bbe98
commit b631e3e
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/ChangeLog.md b/ChangeLog.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `e3nn.IrrepsArray.filtered(keep: List[Irrep])`
 - `e3nn.reduced_tensor_product_basis(formula_or_irreps_list: Union[str, List[e3nn.Irreps]], ...)`
 - `IrrepsArray.at[i].set(v)` and `IrrepsArray.at[i].add(v)`
+- add `Irreps.is_scalar`
 
 ## [0.9.0] - 2022-09-04
 ### Added

diff --git a/e3nn_jax/_src/irreps.py b/e3nn_jax/_src/irreps.py
@@ -454,6 +454,20 @@ def count(self, ir: IntoIrrep) -> int:
         ir = Irrep(ir)
         return sum(mul for mul, irrep in self if ir == irrep)
 
+    def is_scalar(self) -> bool:
+        r"""Check if the representation is scalar.
+
+        Returns:
+            `bool`: ``True`` if the representation is scalar
+
+        Examples:
+            >>> Irreps("2x0e + 3x1o").is_scalar()
+            False
+            >>> Irreps("2x0e + 2x0e").is_scalar()
+            True
+        """
+        return {ir for _, ir in self} == {Irrep("0e")}
+
     def index(self, _object):  # noqa: D102
         raise NotImplementedError