From 851be80a9156cbec9098440e9c25cf32db8e2f36 Mon Sep 17 00:00:00 2001 From: Mit Kotak Date: Tue, 19 Nov 2024 02:38:15 -0500 Subject: [PATCH 1/3] full_tp descriptor --- .../cuequivariance/descriptors/__init__.py | 2 + .../cuequivariance/descriptors/irreps_tp.py | 79 +++++++++++++++++++ .../tests/equivariant_tensor_products_test.py | 6 ++ 3 files changed, 87 insertions(+) diff --git a/cuequivariance/cuequivariance/descriptors/__init__.py b/cuequivariance/cuequivariance/descriptors/__init__.py index c9592cb..af6e6e3 100644 --- a/cuequivariance/cuequivariance/descriptors/__init__.py +++ b/cuequivariance/cuequivariance/descriptors/__init__.py @@ -14,6 +14,7 @@ # limitations under the License. from .transposition import transpose from .irreps_tp import ( + full_tensor_product, fully_connected_tensor_product, channelwise_tensor_product, elementwise_tensor_product, @@ -35,6 +36,7 @@ __all__ = [ "transpose", + "full_tensor_product", "fully_connected_tensor_product", "channelwise_tensor_product", "elementwise_tensor_product", diff --git a/cuequivariance/cuequivariance/descriptors/irreps_tp.py b/cuequivariance/cuequivariance/descriptors/irreps_tp.py index 1efdf2a..0207379 100644 --- a/cuequivariance/cuequivariance/descriptors/irreps_tp.py +++ b/cuequivariance/cuequivariance/descriptors/irreps_tp.py @@ -20,6 +20,29 @@ from cuequivariance.irreps_array.irrep_utils import into_list_of_irrep +def get_cuequivariance_descriptor(irreps_in1, irreps_in2): + d = stp.SegmentedTensorProduct.from_subscripts("ui,vj,kuv+ijk") + irreps_in1 = cue.Irreps("O3", irreps_in1.__str__()) + irreps_in2 = cue.Irreps("O3", irreps_in2.__str__()) + G = irreps_in1.irrep_class + + for mul, ir in irreps_in1: + d.add_segment(0, (mul, ir.dim)) + for mul, ir in irreps_in2: + d.add_segment(1, (mul, ir.dim)) + + for (i1, (mul1, ir1)), (i2, (mul2, ir2)) in itertools.product( + enumerate(irreps_in1), enumerate(irreps_in2) + ): + for ir3 in ir1 * ir2: + # for loop over the different solutions of the Clebsch-Gordan decomposition + for cg in cue.clebsch_gordan(ir1, ir2, ir3): + d.add_path(i1, i2, None, c=cg) + + d = d.normalize_paths_for_operand(-1) + return d + + def fully_connected_tensor_product( irreps1: cue.Irreps, irreps2: cue.Irreps, irreps3: cue.Irreps ) -> cue.EquivariantTensorProduct: @@ -80,6 +103,62 @@ def fully_connected_tensor_product( ) +def full_tensor_product( + irreps1: cue.Irreps, + irreps2: cue.Irreps, + irreps3_filter: Optional[Sequence[cue.Irrep]] = None, +) -> cue.EquivariantTensorProduct: + """ + subscripts: ``lhs[ui],rhs[vj],output[uvk]`` + + Construct a weightless channelwise tensor product descriptor. + + .. currentmodule:: cuequivariance + + Args: + irreps1 (Irreps): Irreps of the first operand. + irreps2 (Irreps): Irreps of the second operand. + irreps3_filter (sequence of Irrep, optional): Irreps of the output to consider. + + Returns: + EquivariantTensorProduct: Descriptor of the full tensor product. + """ + G = irreps1.irrep_class + + if irreps3_filter is not None: + irreps3_filter = into_list_of_irrep(G, irreps3_filter) + + d = stp.SegmentedTensorProduct.from_subscripts("iu,jv,kuv+ijk") + + for mul, ir in irreps1: + d.add_segment(0, (ir.dim, mul)) + for mul, ir in irreps2: + d.add_segment(1, (ir.dim, mul)) + + irreps3 = [] + + for (i1, (mul1, ir1)), (i2, (mul2, ir2)) in itertools.product( + enumerate(irreps1), enumerate(irreps2) + ): + for ir3 in ir1 * ir2: + # for loop over the different solutions of the Clebsch-Gordan decomposition + for cg in cue.clebsch_gordan(ir1, ir2, ir3): + d.add_path(i1, i2, None, c=cg) + + irreps3.append((mul1 * mul2, ir3)) + + irreps3 = cue.Irreps(G, irreps3) + irreps3, perm, inv = irreps3.sort() + d = d.permute_segments(2, inv) + + d = d.normalize_paths_for_operand(-1) + return cue.EquivariantTensorProduct( + d, + [irreps1, irreps2, irreps3], + layout=cue.ir_mul, + ) + + def channelwise_tensor_product( irreps1: cue.Irreps, irreps2: cue.Irreps, diff --git a/cuequivariance/tests/equivariant_tensor_products_test.py b/cuequivariance/tests/equivariant_tensor_products_test.py index a6f00bd..86126bb 100644 --- a/cuequivariance/tests/equivariant_tensor_products_test.py +++ b/cuequivariance/tests/equivariant_tensor_products_test.py @@ -31,6 +31,12 @@ def test_commutativity_squeeze_flatten(): == d.flatten_coefficient_modes().squeeze_modes() ) + d = descriptors.full_tensor_product(irreps1, irreps2, irreps3).d + assert ( + d.squeeze_modes().flatten_coefficient_modes() + == d.flatten_coefficient_modes().squeeze_modes() + ) + d = descriptors.channelwise_tensor_product(irreps1, irreps2, irreps3).d assert ( d.squeeze_modes().flatten_coefficient_modes() From b73a07b6ec0467034595a5ef7bf191b6847be0b0 Mon Sep 17 00:00:00 2001 From: Mit Kotak Date: Tue, 19 Nov 2024 02:49:00 -0500 Subject: [PATCH 2/3] redundant --- .../cuequivariance/descriptors/irreps_tp.py | 23 ------------------- 1 file changed, 23 deletions(-) diff --git a/cuequivariance/cuequivariance/descriptors/irreps_tp.py b/cuequivariance/cuequivariance/descriptors/irreps_tp.py index 0207379..8a76dd7 100644 --- a/cuequivariance/cuequivariance/descriptors/irreps_tp.py +++ b/cuequivariance/cuequivariance/descriptors/irreps_tp.py @@ -20,29 +20,6 @@ from cuequivariance.irreps_array.irrep_utils import into_list_of_irrep -def get_cuequivariance_descriptor(irreps_in1, irreps_in2): - d = stp.SegmentedTensorProduct.from_subscripts("ui,vj,kuv+ijk") - irreps_in1 = cue.Irreps("O3", irreps_in1.__str__()) - irreps_in2 = cue.Irreps("O3", irreps_in2.__str__()) - G = irreps_in1.irrep_class - - for mul, ir in irreps_in1: - d.add_segment(0, (mul, ir.dim)) - for mul, ir in irreps_in2: - d.add_segment(1, (mul, ir.dim)) - - for (i1, (mul1, ir1)), (i2, (mul2, ir2)) in itertools.product( - enumerate(irreps_in1), enumerate(irreps_in2) - ): - for ir3 in ir1 * ir2: - # for loop over the different solutions of the Clebsch-Gordan decomposition - for cg in cue.clebsch_gordan(ir1, ir2, ir3): - d.add_path(i1, i2, None, c=cg) - - d = d.normalize_paths_for_operand(-1) - return d - - def fully_connected_tensor_product( irreps1: cue.Irreps, irreps2: cue.Irreps, irreps3: cue.Irreps ) -> cue.EquivariantTensorProduct: From 59a397cae86723b1f256f23217d012f003a48b23 Mon Sep 17 00:00:00 2001 From: Mit Kotak Date: Tue, 19 Nov 2024 09:53:14 -0500 Subject: [PATCH 3/3] doc fix --- cuequivariance/cuequivariance/descriptors/irreps_tp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cuequivariance/cuequivariance/descriptors/irreps_tp.py b/cuequivariance/cuequivariance/descriptors/irreps_tp.py index 8a76dd7..631e4c0 100644 --- a/cuequivariance/cuequivariance/descriptors/irreps_tp.py +++ b/cuequivariance/cuequivariance/descriptors/irreps_tp.py @@ -86,7 +86,7 @@ def full_tensor_product( irreps3_filter: Optional[Sequence[cue.Irrep]] = None, ) -> cue.EquivariantTensorProduct: """ - subscripts: ``lhs[ui],rhs[vj],output[uvk]`` + subscripts: ``lhs[iu],rhs[jv],output[kuv]`` Construct a weightless channelwise tensor product descriptor.